[Python] tried changing the output of vae model to satisfy loss function but something doesnt...

I tried to create a VAE model which generally contains a custom loss for which GradientTape() or a class is used for implementation. I didn't want to use these methods and instead tried a workaround which isn't working as expected and I wanted to know why.

the model code I tried - (encoder and decoder are just basic models where encoder outputs z_mean, z_log_var, z and decoder outputs a batch of images) -

def build_vae(encoder, decoder, input_shape):
input_layer = layers.Input(shape=input_shape)
z_mean, z_log_var, z = encoder(input_layer)
generated_images = decoder(z)

return models.Model(input_layer, [generated_images, z_mean, z_log_var])


custom loss -

def custom_vae_loss(y_true, y_pred):
generated_images = y_pred[0]
z_mean = y_pred[1]
z_log_var = y_pred[2]

r_loss = tf.reduce_mean(tf.keras.losses.binary_crossentropy(y_true, generated_images))
kl_loss = -0.5 * (tf.reduce_mean(1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)))

return r_loss+kl_loss


training code - vae.train_on_batch(batch, batch)

the error is that target_shape = (8, 256, 256, 3) doesn't match output_shape = (256, 256, 3)

what i was thinking is that the list which should have been y_pred, is getting squeezed or something so that z_mean and z_log_var are being removed which means y_pred (from custom_vae_loss) is the batch of generated images and so generated_images is just the 1st generated image of the batch. after adding print statements in the loss function, y_pred is Tensor("vae_1/decoder_1/conv2d_transpose_31_1/Sigmoid:0", shape=(8, 256, 256, 3), dtype=float32) shape of y_true: (8, 256, 256, 3) shape of generated_images: (256, 256, 3) shape of z_mean: (256, 256, 3) shape of z_log_var: (256, 256, 3)

how should i fix this? preferably without GradientTape() or a class. (I wanted to do this without those methods)

Continue reading...