How do I run the tensorflow sample code? (throws an error)

Question

I am very new to TF2 and tried to customize the example code on the tensorflow guide documentation:

https://www.tensorflow.org/guide/keras/custom_layers_and_models#putting_it_all_together_an_end-to-end_example

The code given in the guide does not run if the latent dimension is set 1, it runs fine for every latent dimension >1!

For training I tried to use the code given in the example but set the latent dim to 1:

vae = VariationalAutoEncoder(784, 64, 1)

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)

vae.compile(optimizer, loss=tf.keras.losses.MeanSquaredError())
vae.fit(x_train, x_train, epochs=3, batch_size=64)

The error when trying to train is:

ValueError: The last dimension of the inputs to Dense should be defined. Found None and is thrown upon return from the Sample function where I think

epsilon = tf.keras.backend.random_normal(shape=(batch, dim))

can not handle shape=(?,1).

Can someone help I am trying to use the code as a template but I need latent dimension to be 1!

Thanks

Answer 1

OK. The answer came from MarkDoust in Github here

The problem here is the interaction of broadcasting with the z_mean + tf.exp(0.5 * z_log_var) * epsilon line.

Normally the last dimension of z_mean, and z_log_var are known but the last dimension of epsilon is not.

Since you mul and add epsilon it assumes the two have the same shape.

When they have a last dimension of 1, It thinks you might be boradcasting z_mean and z_log_var with epsilon, it can't tell.

So the fix is to tell it that you know the shape of epsilon, and are not broadcastng. Add the following before the z_mean + tf.exp(0.5 * z_log_var) * epsilon line: epsilon.set_shape(z_mean.shape) .