TensorFlow 2.0 implementation of a variational autoencoder model.

vae.py

class VariationalAutoencoder(tf.keras.Model):
  def __init__(self, latent_dim, original_dim):
    super(VariationalAutoencoder, self).__init__()
    self.encoder = Encoder(latent_dim=latent_dim)
    self.decoder = Decoder(original_dim=original_dim)

  def call(self, input_features):
    z_mean, z_log_var, latent_code = self.encoder(input_features)
    reconstructed = self.decoder(latent_code)
    kl_divergence = -5e-2 * tf.reduce_sum(tf.exp(z_log_var) + tf.square(z_mean) - 1 - z_log_var)
    self.add_loss(kl_divergence)
    return reconstructed