title variational autoencoder fchollethttps twitter comfchollet date created 20200503 last modified 20231122 description convolutional variational autoencoder vae trained on mnist digits accelerator gpu setup create a sampling layer uses zmean zlogvar to sample z the vector encoding a digit def callself inputs zmean zlogvar inputs batch tf shapezmean0 dim tf shapezmean1 epsilon tf random normalshapebatch dim return zmean tf exp0 5 zlogvar epsilon build the encoder latentdim 2 encoderinputs keras inputshape28 28 1 x layers conv2d32 3 activationrelu strides2 paddingsame encoderinputs x layers conv2d64 3 activationrelu strides2 paddingsamex x layers flattenx x layers dense16 activationrelux zmean layers denselatentdim namezmeanx zlogvar layers denselatentdim namezlogvarx z samplingzmean zlogvar encoder keras modelencoderinputs zmean zlogvar z nameencoder encoder summary build the decoder latentinputs keras inputshapelatentdim x layers dense7 7 64 activationrelulatentinputs x layers reshape7 7 64x x layers conv2dtranspose64 3 activationrelu strides2 paddingsame x x layers conv2dtranspose32 3 activationrelu strides2 paddingsame x decoderoutputs layers conv2dtranspose 1 3 activationsigmoid paddingsame x decoder keras modellatentinputs decoderoutputs namedecoder decoder summary define the vae as a model with a custom trainstep class vaekeras model def initself encoder decoder kwargs super initkwargs self encoder encoder self decoder decoder self totallosstracker keras metrics meannametotalloss self reconstructionlosstracker keras metrics mean namereconstructionloss self kllosstracker keras metrics meannameklloss property def metricsself return self totallosstracker self reconstructionlosstracker self kllosstracker def trainstepself data with tf gradienttape as tape zmean zlogvar z self encoderdata reconstruction self decoderz reconstructionloss tf reducemean tf reducesum keras losses binarycrossentropydata reconstruction axis1 2 klloss 0 5 1 zlogvar tf squarezmean tf expzlogvar klloss tf reducemeantf reducesumklloss axis1 totalloss reconstructionloss klloss grads tape gradienttotalloss self trainableweights self optimizer applygradientszipgrads self trainableweights self totallosstracker updatestatetotalloss self reconstructionlosstracker updatestatereconstructionloss self kllosstracker updatestateklloss return loss self totallosstracker result reconstructionloss self reconstructionlosstracker result klloss self kllosstracker result train the vae xtrain xtest keras datasets mnist loaddata mnistdigits np concatenatextrain xtest axis0 mnistdigits np expanddimsmnistdigits 1 astypefloat32 255 vae vaeencoder decoder vae compileoptimizerkeras optimizers adam vae fitmnistdigits epochs30 batchsize128 display a grid of sampled digits import matplotlib pyplot as plt def plotlatentspacevae n30 figsize15 display a nn 2d manifold of digits digitsize 28 scale 1 0 figure np zerosdigitsize n digitsize n linearly spaced coordinates corresponding to the 2d plot of digit classes in the latent space gridx np linspacescale scale n gridy np linspacescale scale n 1 for i yi in enumerategridy for j xi in enumerategridx zsample np arrayxi yi xdecoded vae decoder predictzsample verbose0 digit xdecoded0 reshapedigitsize digitsize figure i digitsize i 1 digitsize j digitsize j 1 digitsize digit plt figurefigsizefigsize figsize startrange digitsize 2 endrange n digitsize startrange pixelrange np arangestartrange endrange digitsize samplerangex np roundgridx 1 samplerangey np roundgridy 1 plt xtickspixelrange samplerangex plt ytickspixelrange samplerangey plt xlabelz0 plt ylabelz1 plt imshowfigure cmapgreysr plt show plotlatentspacevae display how the latent space clusters different digit classes def plotlabelclustersvae data labels display a 2d plot of the digit classes in the latent space zmean vae encoder predictdata verbose0 plt figurefigsize12 10 plt scatterzmean 0 zmean 1 clabels plt colorbar plt xlabelz0 plt ylabelz1 plt show xtrain ytrain keras datasets mnist loaddata xtrain np expanddimsxtrain 1 astypefloat32 255 plotlabelclustersvae xtrain ytrain setup create a sampling layer uses z_mean z_log_var to sample z the vector encoding a digit build the encoder build the decoder define the vae as a model with a custom train_step train the vae display a grid of sampled digits display a n n 2d manifold of digits linearly spaced coordinates corresponding to the 2d plot of digit classes in the latent space display how the latent space clusters different digit classes display a 2d plot of the digit classes in the latent space
import os
os.environ["KERAS_BACKEND"] = "tensorflow"

import numpy as np
import tensorflow as tf
import keras
from keras import layers




class Sampling(layers.Layer):
    

    def call(self, inputs):
        z_mean, z_log_var = inputs
        batch = tf.shape(z_mean)[0]
        dim = tf.shape(z_mean)[1]
        epsilon = tf.random.normal(shape=(batch, dim))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon




latent_dim = 2

encoder_inputs = keras.Input(shape=(28, 28, 1))
x = layers.Conv2D(32, 3, activation="relu", strides=2, padding="same")(
    encoder_inputs
)
x = layers.Conv2D(64, 3, activation="relu", strides=2, padding="same")(x)
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()



latent_inputs = keras.Input(shape=(latent_dim,))
x = layers.Dense(7 * 7 * 64, activation="relu")(latent_inputs)
x = layers.Reshape((7, 7, 64))(x)
x = layers.Conv2DTranspose(64, 3, activation="relu", strides=2, padding="same")(
    x
)
x = layers.Conv2DTranspose(32, 3, activation="relu", strides=2, padding="same")(
    x
)
decoder_outputs = layers.Conv2DTranspose(
    1, 3, activation="sigmoid", padding="same"
)(x)
decoder = keras.Model(latent_inputs, decoder_outputs, name="decoder")
decoder.summary()




class VAE(keras.Model):
    def __init__(self, encoder, decoder, **kwargs):
        super().__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.total_loss_tracker = keras.metrics.Mean(name="total_loss")
        self.reconstruction_loss_tracker = keras.metrics.Mean(
            name="reconstruction_loss"
        )
        self.kl_loss_tracker = keras.metrics.Mean(name="kl_loss")

    @property
    def metrics(self):
        return [
            self.total_loss_tracker,
            self.reconstruction_loss_tracker,
            self.kl_loss_tracker,
        ]

    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(data)
            reconstruction = self.decoder(z)
            reconstruction_loss = tf.reduce_mean(
                tf.reduce_sum(
                    keras.losses.binary_crossentropy(data, reconstruction),
                    axis=(1, 2),
                )
            )
            kl_loss = -0.5 * (
                1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var)
            )
            kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
            total_loss = reconstruction_loss + kl_loss
        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        return {
            "loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result(),
        }




(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()
mnist_digits = np.concatenate([x_train, x_test], axis=0)
mnist_digits = np.expand_dims(mnist_digits, -1).astype("float32") / 255

vae = VAE(encoder, decoder)
vae.compile(optimizer=keras.optimizers.Adam())
vae.fit(mnist_digits, epochs=30, batch_size=128)



import matplotlib.pyplot as plt


def plot_latent_space(vae, n=30, figsize=15):
    
    digit_size = 28
    scale = 1.0
    figure = np.zeros((digit_size * n, digit_size * n))
    
    
    grid_x = np.linspace(-scale, scale, n)
    grid_y = np.linspace(-scale, scale, n)[::-1]

    for i, yi in enumerate(grid_y):
        for j, xi in enumerate(grid_x):
            z_sample = np.array([[xi, yi]])
            x_decoded = vae.decoder.predict(z_sample, verbose=0)
            digit = x_decoded[0].reshape(digit_size, digit_size)
            figure[
                i * digit_size : (i + 1) * digit_size,
                j * digit_size : (j + 1) * digit_size,
            ] = digit

    plt.figure(figsize=(figsize, figsize))
    start_range = digit_size // 2
    end_range = n * digit_size + start_range
    pixel_range = np.arange(start_range, end_range, digit_size)
    sample_range_x = np.round(grid_x, 1)
    sample_range_y = np.round(grid_y, 1)
    plt.xticks(pixel_range, sample_range_x)
    plt.yticks(pixel_range, sample_range_y)
    plt.xlabel("z[0]")
    plt.ylabel("z[1]")
    plt.imshow(figure, cmap="Greys_r")
    plt.show()


plot_latent_space(vae)




def plot_label_clusters(vae, data, labels):
    
    z_mean, _, _ = vae.encoder.predict(data, verbose=0)
    plt.figure(figsize=(12, 10))
    plt.scatter(z_mean[:, 0], z_mean[:, 1], c=labels)
    plt.colorbar()
    plt.xlabel("z[0]")
    plt.ylabel("z[1]")
    plt.show()


(x_train, y_train), _ = keras.datasets.mnist.load_data()
x_train = np.expand_dims(x_train, -1).astype("float32") / 255

plot_label_clusters(vae, x_train, y_train)