Variational Autoencoder - ValueError: No gradients provided for any variable (TensorFlow2.6)
I am implementing a toy Variational Autoencoder in TensorFlow 2.6, Python 3.9 for MNIST dataset. The code is:
# Specify latent space-
latent_dim = 3
class Sampling(layers.Layer):
'''
Create a sampling layer.
Uses (z_mean, z_log_var) to sample z - the vector encoding a digit.
'''
def call(self, inputs):
z_mean, z_log_var = inputs
batch = tf.shape(z_mean)[0]
dim = tf.shape(z_mean)[1]
epsilon = tf.keras.backend.random_normal(shape = (batch, dim))
return z_mean + tf.exp(0.5 * z_log_var) * epsilon
class Encoder(Model):
def __init__(self):
super(Encoder, self).__init__()
self.conv1 = Conv2D(
filters = 32, kernel_size = (3, 3),
activation = None , strides = 2,
padding="same")
self.conv2 = Conv2D(
filters = 64, kernel_size = (3, 3),
activation = "relu", strides = 2,
padding = "same")
self.flatten = Flatten()
self.dense = Dense(
units = 16, activation = None
)
def call(self, x):
x = tf.keras.activations.relu(self.conv1(x))
x = tf.keras.activations.relu(self.conv2(x))
x = self.flatten(x)
x = tf.keras.activations.relu(self.dense(x))
return x
class Decoder(Model):
def __init__(self):
super(Decoder, self).__init__()
self.dense = Dense(
units = 7 * 7 * 64, activation = None)
self.conv_tran_1 = Conv2DTranspose(
filters = 64, kernel_size = (3, 3),
activation = None, strides = 2,
padding = "same")
self.conv_tran_2 = Conv2DTranspose(
filters = 32, kernel_size = (3, 3),
activation = None, strides = 2,
padding = "same")
self.decoder_outputs = Conv2DTranspose(
filters = 1, kernel_size = (3, 3),
activation = None, padding = "same")
def call(self, x):
x = tf.keras.activations.relu(self.dense(x))
x = layers.Reshape((7, 7, 64))(x)
x = tf.keras.activations.relu(self.conv_tran_1(x))
x = tf.keras.activations.relu(self.conv_tran_2(x))
x = self.decoder_outputs(x)
return x
class VAE(Model):
def __init__(self, latent_space = 3):
super(VAE, self).__init__()
self.latent_space = latent_space
self.encoder = Encoder()
self.z_mean = Dense(units = self.latent_space, activation = None)
self.z_log_var = Dense(units = self.latent_space, activation = None)
self.decoder = Decoder()
def reparameterize(self, encoded_mean, encoded_log_var):
# NOT USED!
# encoded_mean = self.z_mean(x)
# encoded_log_var = self.z_log_var(x)
batch = tf.shape(encoded_mean)[0]
encoded_dim = tf.shape(encoded_mean)[1]
epsilon = tf.keras.backend.random_normal(shape = (batch, encoded_dim))
return encoded_mean + tf.exp(0.5 * encoded_log_var) * epsilon
def call(self, x):
x = self.encoder(x)
mu = self.z_mean(x)
log_var = self.z_log_var(x)
# z = self. 338