Implementation of mean and log_variance in variational autoencoder keras

user_15156875 0
Viewed 228

As the implementation in Keras for the VAE https://keras.io/examples/generative/vae/, we have to pass a mean and log_variance to calculate the distribution in the latent space.

class Sampling(layers.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

# flatten layer
x = layers.Flatten()(x)
x = layers.Dense(16, activation="relu")(x)
z_mean = layers.Dense(latent_dim, name="z_mean")(x)
z_log_var = layers.Dense(latent_dim, name="z_log_var")(x)
z = Sampling()([z_mean, z_log_var])
encoder = keras.Model(encoder_inputs, [z_mean, z_log_var, z], name="encoder")
encoder.summary()

I don't understand how two dense layers can represent the mean and log variance without doing any special calculation? Because from the code above is just simply create a dense layer and receive result from the previous flatten layer.

autoencoder keras
1 Answers

The dense output layers are trained to output mean and log variance for the input using the Kullback–Leibler divergence loss function.

In the Keras example VAE model it is calculated in the custom train_step using the output of the dense layers:

kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
user_52100750
Related