How to fix error where a KerasTensor is passed to a TF API? - tensorflow

Question

I've downloaded code for a Wasserstein GAN with Gradient Policy (WGAN-GP) from Keras-GAN (GitHub). Some of the imports appeared to be of outdated syntax, as I was getting errors and they were based on the pre-Tensorflow Keras. After a while of searching and tinkering, I have determined that I have no idea what to do next.
What I do know is that, in the following code, both the,
interpolated_img = RandomWeightedAverage()([real_img, fake_img])
and the,
validity_interpolated = self.critic(interpolated_img)
are both of the type KerasTensor or, more specifically, of type,
<class 'keras.engine.keras_tensor.KerasTensor'>
and that immediately after printing both their types, the program crashes. So, it certainly seems to be caused by these objects.
Here is the code:
from __future__ import print_function, division
import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.layers import Concatenate # _Merge
from tensorflow.keras.layers import Input, Dense, Reshape, Flatten, Dropout
from tensorflow.keras.layers import BatchNormalization, Activation, ZeroPadding2D
from tensorflow.keras.layers import LeakyReLU
from tensorflow.keras.layers import Conv2D, Conv2DTranspose, UpSampling2D
from tensorflow.keras.models import Sequential, Model
from tensorflow.keras.optimizers import RMSprop
from functools import partial
import tensorflow.keras.backend as K
from keras.layers.merge import _Merge
import matplotlib.pyplot as plt
import sys
import numpy as np
class RandomWeightedAverage(tf.keras.layers.Layer):
"""Provides a (random) weighted average between real and generated image samples"""
def __init__(self, batch_size=32):
super().__init__()
self.batch_size = batch_size
def call(self, inputs, **kwargs):
alpha = tf.random.uniform((32, 1, 1, 1))
return (alpha * inputs[0]) + ((1 - alpha) * inputs[1])
def comput_output_shape(self, input_shape):
return input_shape[0]
class WGANGP():
def __init__(self, height=128, width=128, channels=3, noise_dim=100, batch_size=64):
self.img_height = height
self.img_width = width
self.channels = channels
self.img_shape = (self.img_height, self.img_width, self.channels)
self.noise_dim = noise_dim
self.batch_size = batch_size
# Following parameter and optimizer set as recommended in paper
self.n_critic = 5
optimizer = RMSprop(lr=0.00005)
# Build the generator and critic
self.generator = self.build_generator()
self.critic = self.build_critic()
#-------------------------------
# Construct Computational Graph
# for the Critic
#-------------------------------
# Freeze generator's layers while training critic
self.generator.trainable = False
# Image input (real sample)
real_img = Input(shape=self.img_shape)
# Noise input
z_disc = Input(shape=(self.noise_dim,))
# Generate image based of noise (fake sample)
fake_img = self.generator(z_disc)
# Discriminator determines validity of the real and fake images
fake = self.critic(fake_img)
valid = self.critic(real_img)
# Construct weighted average between real and fake images
interpolated_img = RandomWeightedAverage()([real_img, fake_img])
# Determine validity of weighted sample
validity_interpolated = self.critic(interpolated_img)
# Use Python partial to provide loss function with additional
# 'averaged_samples' argument
partial_gp_loss = partial(self.gradient_penalty_loss,
averaged_samples=interpolated_img)
partial_gp_loss.__name__ = 'gradient_penalty' # Keras requires function names
self.critic_model = Model(inputs=[real_img, z_disc],
outputs=[valid, fake, validity_interpolated])
self.critic_model.compile(loss=[self.wasserstein_loss,
self.wasserstein_loss,
partial_gp_loss],
optimizer=optimizer,
loss_weights=[1, 1, 10])
#-------------------------------
# Construct Computational Graph
# for Generator
#-------------------------------
# For the generator we freeze the critic's layers
self.critic.trainable = False
self.generator.trainable = True
# Sampled noise for input to generator
z_gen = Input(shape=(self.noise_dim,))
# Generate images based of noise
img = self.generator(z_gen)
# Discriminator determines validity
valid = self.critic(img)
# Defines generator model
self.generator_model = Model(z_gen, valid)
self.generator_model.compile(loss=self.wasserstein_loss, optimizer=optimizer)
def gradient_penalty_loss(self, y_true, y_pred, averaged_samples):
"""
Computes gradient penalty based on prediction and weighted real / fake samples
"""
gradients = K.gradients(y_pred, averaged_samples)[0]
# compute the euclidean norm by squaring ...
gradients_sqr = K.square(gradients)
# ... summing over the rows ...
gradients_sqr_sum = K.sum(gradients_sqr,
axis=np.arange(1, len(gradients_sqr.shape)))
# ... and sqrt
gradient_l2_norm = K.sqrt(gradients_sqr_sum)
# compute lambda * (1 - ||grad||)^2 still for each single sample
gradient_penalty = K.square(1 - gradient_l2_norm)
# return the mean as loss over all the batch samples
return K.mean(gradient_penalty)
def wasserstein_loss(self, y_true, y_pred):
return K.mean(y_true * y_pred)
def build_generator(self):
model = Sequential()
model.add(Dense(128 * 7 * 7, activation="relu", input_dim=self.noise_dim))
model.add(Reshape((7, 7, 128)))
model.add(UpSampling2D())
model.add(Conv2D(128, kernel_size=4, padding="same"))
model.add(BatchNormalization(momentum=0.8))
model.add(Activation("relu"))
model.add(UpSampling2D())
model.add(Conv2D(64, kernel_size=4, padding="same"))
model.add(BatchNormalization(momentum=0.8))
model.add(Activation("relu"))
model.add(Conv2D(self.channels, kernel_size=4, padding="same"))
model.add(Activation("tanh"))
model.summary()
#
# noise = Input(shape=(self.noise_dim,))
# img = model(noise)
return model # Model(noise, img)
def build_critic(self):
model = Sequential()
model.add(Conv2D(16, kernel_size=3, strides=2, input_shape=self.img_shape, padding="same"))
model.add(LeakyReLU(alpha=0.2))
model.add(Dropout(0.25))
model.add(Conv2D(32, kernel_size=3, strides=2, padding="same"))
model.add(ZeroPadding2D(padding=((0,1),(0,1))))
model.add(BatchNormalization(momentum=0.8))
model.add(LeakyReLU(alpha=0.2))
model.add(Dropout(0.25))
model.add(Conv2D(64, kernel_size=3, strides=2, padding="same"))
model.add(BatchNormalization(momentum=0.8))
model.add(LeakyReLU(alpha=0.2))
model.add(Dropout(0.25))
model.add(Conv2D(128, kernel_size=3, strides=1, padding="same"))
model.add(BatchNormalization(momentum=0.8))
model.add(LeakyReLU(alpha=0.2))
model.add(Dropout(0.25))
model.add(Flatten())
model.add(Dense(1))
model.summary()
# img = Input(shape=self.img_shape)
# validity = model(img)
return model # Model(img, validity)
def train(self, epochs, batch_size, sample_interval=50):
# Load the dataset
(X_train, _), (_, _) = mnist.load_data()
# Rescale -1 to 1
X_train = (X_train.astype(np.float32) - 127.5) / 127.5
X_train = np.expand_dims(X_train, axis=3)
# Adversarial ground truths
valid = -np.ones((batch_size, 1))
fake = np.ones((batch_size, 1))
dummy = np.zeros((batch_size, 1)) # Dummy gt for gradient penalty
for epoch in range(epochs):
for _ in range(self.n_critic):
# ---------------------
# Train Discriminator
# ---------------------
# Select a random batch of images
idx = np.random.randint(0, X_train.shape[0], batch_size)
imgs = X_train[idx]
# Sample generator input
noise = np.random.normal(0, 1, (batch_size, self.noise_dim))
# Train the critic
d_loss = self.critic_model.train_on_batch([imgs, noise],
[valid, fake, dummy])
# ---------------------
# Train Generator
# ---------------------
g_loss = self.generator_model.train_on_batch(noise, valid)
# Plot the progress
print ("%d [D loss: %f] [G loss: %f]" % (epoch, d_loss[0], g_loss))
# If at save interval => save generated image samples
if epoch % sample_interval == 0:
self.sample_images(epoch)
def sample_images(self, epoch):
r, c = 5, 5
noise = np.random.normal(0, 1, (r * c, self.noise_dim))
gen_imgs = self.generator.predict(noise)
# Rescale images 0 - 1
gen_imgs = 0.5 * gen_imgs + 0.5
fig, axs = plt.subplots(r, c)
cnt = 0
for i in range(r):
for j in range(c):
axs[i,j].imshow(gen_imgs[cnt, :,:,0], cmap='gray')
axs[i,j].axis('off')
cnt += 1
fig.savefig("images/mnist_%d.png" % epoch)
plt.close()
if __name__ == '__main__':
img_width = 28
img_height = 28
channels = 1
wgan = WGANGP(height=img_height, width=img_width, channels=channels)
wgan.train(epochs=30000, batch_size=32, sample_interval=100)
I get the following error first:
Traceback (most recent call last):
File "[REDACTED PATH]", line 255, in <module>
wgan.train(epochs=30000, batch_size=32, sample_interval=100)
File "[REDACTED PATH]", line 215, in train
d_loss = self.critic_model.train_on_batch([imgs, noise],
File "J:\Anaconda3\lib\site-packages\keras\engine\training.py", line 2093, in train_on_batch
logs = self.train_function(iterator)
File "J:\Anaconda3\lib\site-packages\tensorflow\python\util\traceback_utils.py", line 153, in error_handler
raise e.with_traceback(filtered_tb) from None
File "J:\Anaconda3\lib\site-packages\tensorflow\python\framework\func_graph.py", line 1147, in autograph_handler
raise e.ag_error_metadata.to_exception(e)
Which is immediately followed by:
TypeError: in user code:
File "J:\Anaconda3\lib\site-packages\keras\engine\training.py", line 1021, in train_function *
return step_function(self, iterator)
File "[REDACTED PATH]", line 117, in gradient_penalty_loss *
gradients = K.gradients(y_pred, averaged_samples)[0]
File "J:\Anaconda3\lib\site-packages\keras\backend.py", line 4352, in gradients **
return tf.compat.v1.gradients(
File "J:\Anaconda3\lib\site-packages\numpy\core\_asarray.py", line 102, in asarray
return array(a, dtype, copy=False, order=order)
File "J:\Anaconda3\lib\site-packages\keras\engine\keras_tensor.py", line 254, in __array__
raise TypeError(
TypeError: You are passing KerasTensor(type_spec=TensorSpec(shape=(32, 28, 28, 1), dtype=tf.float32, name=None), name='random_weighted_average/add:0', description="created by layer 'random_weighted_average'"), an intermediate Keras symbolic input/output, to a TF API that does not allow registering custom dispa
tchers, such as `tf.cond`, `tf.function`, gradient tapes, or `tf.map_fn`. Keras Functional model construction only supports TF API calls that *do* support dispatching, such as `tf.math.add` or `tf.reshape`. Other APIs cannot be called directly on symbolic Kerasinputs/outputs. You can work around this limitation by
putting the operation in a custom Keras layer `call` and calling that layer on this symbolic input/output.

Answer 1

I solved the problem disabling eager execution
from tensorflow.python.framework.ops import disable_eager_execution
disable_eager_execution()
I also read some answers which suggested that this problem might be due to numpy 1.20>= , If the solution above doesn't work try downgrading numpy to like 1.19.5