Variational Autoencoders as Generative Models with Keras | by Kartik Chaudhary | Towards Data Science

This article focuses on giving the readers some basic understanding of the Variational Autoencoders and explaining how they are different from the ordinary autoencoders in Machine Learning and Artificial Intelligence. Unlike vanilla autoencoders(like-sparse autoencoders, de-noising autoencoders .etc), Variational Autoencoders (VAEs) are generative models like GANs (Generative Adversarial Networks). This article is primarily focused on the Variational Autoencoders and I will be writing soon about the Generative Adversarial Networks in my upcoming posts.
In this tutorial, we will be discussing how to train a variational autoencoder(VAE) with Keras(TensorFlow, Python) from scratch. We will be concluding our study with the demonstration of the generative capabilities of a simple VAE.
The rest of the content in this tutorial can be classified as the following-
Background: Variational AutoEncoders (VAEs)Building VAE in KerasTraining VAE on the MNIST datasetResultsImage Generation CapabilitiesSummaryFurther reading and resources
Background: Variational AutoEncoders (VAEs)
Building VAE in Keras
Training VAE on the MNIST dataset
Results
Image Generation Capabilities
Summary
Further reading and resources
An autoencoder is basically a neural network that takes a high dimensional data point as input, converts it into a lower-dimensional feature vector(ie., latent vector), and later reconstructs the original input sample just utilizing the latent vector representation without losing valuable information. Any given autoencoder is consists of the following two parts-an Encoder and a Decoder. The Encoder part of the model takes an input data sample and compresses it into a latent vector. While the decoder part is responsible for recreating the original input sample from the learned(learned by the encoder during training) latent representation. To learn more about the basics, do check out my article on Autoencoders in Keras and Deep Learning.
Let’s continue considering that we all are on the same page until now.
One issue with the ordinary autoencoders is that they encode each input sample independently. This means that the samples belonging to the same class (or the samples belonging to the same distribution) might learn very different(distant encodings in the latent space) latent embeddings. Ideally, the latent features of the same class should be somewhat similar (or closer in latent space). This happens because we are not explicitly forcing the neural network to learn the distributions of the input dataset. Due to this issue, our network might not very good at reconstructing related unseen data samples (or less generalizable).
In the past tutorial on Autoencoders in Keras and Deep Learning, we trained a vanilla autoencoder and learned the latent features for the MNIST handwritten digit images. When we plotted these embeddings in the latent space with the corresponding labels, we found the learned embeddings of the same classes coming out quite random sometimes and there were no clearly visible boundaries between the embedding clusters of the different classes. The following figure shows the distribution-
Variational Autoencoder is slightly different in nature. Instead of directly learning the latent features from the input samples, it actually learns the distribution of latent features. The latent features of the input data are assumed to be following a standard normal distribution. This means that the learned latent vectors are supposed to be zero centric and they can be represented with two statistics-mean and variance (as standard normal distribution can be attributed with only these two statistics).
Thus the Variational AutoEncoders(VAEs) calculate the mean and variance of the latent vectors(instead of directly learning latent features) for each sample and forces them to follow a standard normal distribution. Thus the bottleneck part of the network is used to learn mean and variance for each sample, we will define two different fully connected(FC) layers to calculate both. VAEs ensure that the points that are very close to each other in the latent space, are representing very similar data samples(similar classes of data). We are going to prove this fact in this tutorial.
Before jumping into the implementation details let’s first get a little understanding of the KL-divergence which is going to be used as one of the two optimization measures in our model.
In the last section, we were talking about enforcing a standard normal distribution on the latent features of the input dataset. This can be accomplished using KL-divergence statistics. KL-divergence is a statistical measure of the difference between two probabilistic distributions. Thus, we will utilize KL-divergence value as an objective function(along with the reconstruction loss) in order to ensure that the learned distribution is very similar to the true distribution, which we have already assumed to be a standard normal distribution.
In this case, the final objective can be written as-
Objective = Reconstruction Loss + KL-Loss
Here, the reconstruction loss term would encourage the model to learn the important latent features, needed to correctly reconstruct the original image (if not exactly the same, an image of the same class). While the KL-divergence-loss term would ensure that the learned distribution is similar to the true distribution(a standard normal distribution). This further means that the distribution is centered at zero and is well-spread in the space. We will prove this one also in the latter part of the tutorial.
The last section has explained the basic idea behind the Variational Autoencoders(VAEs) in machine learning(ML) and artificial intelligence(AI). In this section, we will build a convolutional variational autoencoder with Keras in Python. This network will be trained on the MNIST handwritten digits dataset that is available in Keras datasets.
This section can be broken into the following parts for step-wise understanding and simplicity-
Data PreparationBuilding EncoderLatent Distribution and SamplingBuilding DecoderBuilding VAELoss
Data Preparation
Building Encoder
Latent Distribution and Sampling
Building Decoder
Building VAE
Loss
In this section, we are going to download and load the MNIST handwritten digits dataset into our Python notebook to get started with the data preparation.
Here are the dependencies, loaded in advance-
import numpy as npimport matplotlib.pyplot as pltimport pandas as pdimport seaborn as snsimport warningswarnings.filterwarnings('ignore')%matplotlib inline
The following python code can be used to download the MNIST handwritten digits dataset. Few sample images are also displayed below-
from tensorflow.keras.datasets import mnist(trainX, trainy), (testX, testy) = mnist.load_data()print('Training data shapes: X=%s, y=%s' % (trainX.shape, trainy.shape))print('Testing data shapes: X=%s, y=%s' % (testX.shape, testy.shape))for j in range(5):    i = np.random.randint(0, 10000)    plt.subplot(550 + 1 + j)    plt.imshow(trainX[i], cmap='gray')    plt.title(trainy[i])plt.show()
Dataset is already divided into the training and test set. The training dataset has 60K handwritten digit images with a resolution of 28*28. While the Test dataset consists of 10K handwritten digit images with similar dimensions-
Each image in the dataset is a 2D matrix representing pixel intensities ranging from 0 to 255. We will first normalize the pixel values(To bring them between 0 and 1) and then add an extra dimension for image channels (as supported by Conv2D layers from Keras). Here is the preprocessing code in python-
train_data = trainX.astype('float32')/255test_data = testX.astype('float32')/255train_data = np.reshape(train_data, (60000, 28, 28, 1))test_data = np.reshape(test_data, (10000, 28, 28, 1))print (train_data.shape, test_data.shape)Out[1]: (60000, 28, 28, 1) (10000, 28, 28, 1)
In this section, we will define the encoder part of our VAE model. The encoder part of the autoencoder usually consists of multiple repeating convolutional layers followed by pooling layers when the input data type is images. The encoder part of a variational autoencoder is also quite similar, it’s just the bottleneck part that is slightly different as discussed above.
Here is the python implementation of the encoder part with Keras-
import tensorflowinput_data = tensorflow.keras.layers.Input(shape=(28, 28, 1))encoder = tensorflow.keras.layers.Conv2D(64, (5,5), activation='relu')(input_data)encoder = tensorflow.keras.layers.MaxPooling2D((2,2))(encoder)encoder = tensorflow.keras.layers.Conv2D(64, (3,3), activation='relu')(encoder)encoder = tensorflow.keras.layers.MaxPooling2D((2,2))(encoder)encoder = tensorflow.keras.layers.Conv2D(32, (3,3), activation='relu')(encoder)encoder = tensorflow.keras.layers.MaxPooling2D((2,2))(encoder)encoder = tensorflow.keras.layers.Flatten()(encoder)encoder = tensorflow.keras.layers.Dense(16)(encoder)
The above snippet compresses the image input and brings down it to a 16 valued feature vector, but these are not the final latent features. The next section will complete the encoder part by adding the latent features computational logic into it.
This section is responsible for taking the convoluted features from the last section and calculating the mean and log-variance of the latent features (As we have assumed that the latent features follow a standard normal distribution, and the distribution can be represented with mean and variance statistical values). Two separate fully connected(FC layers) layers are used for calculating the mean and log-variance for the input samples of a given dataset.
These attributes(mean and log-variance) of the standard normal distribution(SND) are then used to estimate the latent encodings for the corresponding input data points. The function sample_latent_features defined below takes these two statistical values and returns back a latent encoding vector. This latent encoding is passed to the decoder as input for the image reconstruction purpose.
def sample_latent_features(distribution):    distribution_mean, distribution_variance = distribution    batch_size = tensorflow.shape(distribution_variance)[0]    random = tensorflow.keras.backend.random_normal(shape=(batch_size, tensorflow.shape(distribution_variance)[1]))    return distribution_mean + tensorflow.exp(0.5 * distribution_variance) * random distribution_mean = tensorflow.keras.layers.Dense(2, name='mean')(encoder)distribution_variance = tensorflow.keras.layers.Dense(2, name='log_variance')(encoder)latent_encoding = tensorflow.keras.layers.Lambda(sample_latent_features)([distribution_mean, distribution_variance])
These latent features(calculated from the learned distribution) actually complete the Encoder part of the model. Now the Encoder model can be defined as follow-
encoder_model = tensorflow.keras.Model(input_data, latent_encoding)encoder_model.summary()
The encoder is quite simple with just around 57K trainable parameters.
The Encoder part of the model takes an image as input and gives the latent encoding vector for it as output which is sampled from the learned distribution of the input dataset. The job of the decoder is to take this embedding vector as input and recreate the original image(or an image belonging to a similar class as the original image). As the latent vector is a quite compressed representation of the features, the decoder part is made up of multiple pairs of the Deconvolutional layers and upsampling layers. A deconvolutional layer basically reverses what a convolutional layer does. The upsampling layers are used to bring the original resolution of the image back. In this way, it reconstructs the image with original dimensions.
Here is the python implementation of the decoder part with Keras API from TensorFlow-
decoder_input = tensorflow.keras.layers.Input(shape=(2))decoder = tensorflow.keras.layers.Dense(64)(decoder_input)decoder = tensorflow.keras.layers.Reshape((1, 1, 64))(decoder)decoder = tensorflow.keras.layers.Conv2DTranspose(64, (3,3), activation='relu')(decoder) decoder = tensorflow.keras.layers.Conv2DTranspose(64, (3,3), activation='relu')(decoder)decoder = tensorflow.keras.layers.UpSampling2D((2,2))(decoder) decoder = tensorflow.keras.layers.Conv2DTranspose(64, (3,3), activation='relu')(decoder)decoder = tensorflow.keras.layers.UpSampling2D((2,2))(decoder) decoder_output = tensorflow.keras.layers.Conv2DTranspose(1, (5,5), activation='relu')(decoder)
The decoder model object can be defined as below-
decoder_model = tensorflow.keras.Model(decoder_input, decoder_output)decoder_model.summary()
The decoder is again simple with 112K trainable parameters.
Finally, the Variational Autoencoder(VAE) can be defined by combining the encoder and the decoder parts. Here is how you can create the VAE model object by sticking decoder after the encoder.
encoded = encoder_model(input_data)decoded = decoder_model(encoded)autoencoder = tensorflow.keras.models.Model(input_data, decoded)autoencoder.summary()
The overall setup is quite simple with just 170K trainable model parameters. Time to write the objective(or optimization function) function.
As discussed earlier, the final objective(or loss) function of a variational autoencoder(VAE) is a combination of the data reconstruction loss and KL-loss. In this section, we will define our custom loss by combining these two statistics.
The following implementation of the get_loss function returns a total_loss function that is a combination of reconstruction loss and KL-loss as defined below-
Finally, let’s compile the model to make it ready for the training-
autoencoder.compile(loss=get_loss(distribution_mean, distribution_variance), optimizer='adam')
Just like the ordinary autoencoders, we will train it by giving exactly the same images for input as well as the output. The model is trained for 20 epochs with a batch size of 64.
Here is the training summary-
autoencoder.fit(train_data, train_data, epochs=20, batch_size=64, validation_data=(test_data, test_data))
I hope it can be trained a little more, but this is where the validation loss was not changing much and I went ahead with it.
In this section, we will see the reconstruction capabilities of our model on the test images. The following python script will pick 9 images from the test dataset and we will be plotting the corresponding reconstructed images for them.
offset=400print ("Real Test Images")# Real Imagesfor i in range(9):    plt.subplot(330 + 1 + i)    plt.imshow(test_data[i+offset,:,:, -1], cmap='gray')plt.show()# Reconstructed Imagesprint ("Reconstructed Images with Variational Autoencoder")for i in range(9):    plt.subplot(330 + 1 + i)    output = autoencoder.predict(np.array([test_data[i+offset]]))    op_image = np.reshape(output[0]*255, (28, 28))    plt.imshow(op_image, cmap='gray')plt.show()
Here is the output-
The above results confirm that the model is able to reconstruct the digit images with decent efficiency. However, one important thing to notice here is that some of the reconstructed images are very different in appearance from the original images while the class(or digit) is always the same. This happens because, the reconstruction is not just dependent upon the input image, it is the distribution that has been learned. And this learned distribution is the reason for the introduced variations in the model output. This is interesting, isn’t it!
The second thing to notice here is that the output images are a little blurry. This is a common case with variational autoencoders, they often produce noisy(or poor quality) outputs as the latent vectors(bottleneck) is very small and there is a separate process of learning the latent features as discussed before. Variational Autoencoders(VAEs) are not actually designed to reconstruct the images, the real purpose is learning the distribution (and it gives them the superpower to generate fake data, we will see it later in the post).
As we have quoted earlier, the variational autoencoders(VAEs) learn the underlying distribution of the latent features, it basically means that the latent encodings of the samples belonging to the same class should not be very far from each other in the latent space. Secondly, the overall distribution should be standard normal, which is supposed to be centered at zero.
Let’s generate the latent embeddings for all of our test images and plot them(the same color represents the digits belonging to the same class, taken from the ground truth labels). Here is the python code-
x = []y = []z = []for i in range(10000):    z.append(testy[i])    op = encoder_model.predict(np.array([test_data[i]]))    x.append(op[0][0])    y.append(op[0][1])df = pd.DataFrame()df['x'] = xdf['y'] = ydf['z'] = ["digit-"+str(k) for k in z]plt.figure(figsize=(8, 6))sns.scatterplot(x='x', y='y', hue='z', data=df)plt.show()
The above plot shows that the distribution is centered at zero. Embeddings of the same class digits are closer in the latent space. Digit separation boundaries can also be drawn easily. This is pretty much we wanted to achieve from the variational autoencoder. Let’s jump to the final part where we test the generative capabilities of our model.
Variational Autoencoders can be used as generative models. The previous section shows that latent encodings of the input data are following a standard normal distribution and there are clear boundaries visible for different classes of the digits.
Just think for a second-If we already know, which part of the space is dedicated to what class, we don’t even need input images to reconstruct the image. This means that we can actually generate digit images having similar characteristics as the training dataset by just passing the random points from the space (latent distribution space). In this fashion, the variational autoencoders can be used as generative models in order to generate fake data.
As we can see, the spread of latent encodings is in between [-3 to 3 on the x-axis, and also -3 to 3 on the y-axis]. Let’s generate a bunch of digits with random latent encodings belonging to this range only.
generator_model = decoder_modelx_values = np.linspace(-3, 3, 30)y_values = np.linspace(-3, 3, 30)figure = np.zeros((28 * 30, 28 * 30))for ix, x in enumerate(x_values):    for iy, y in enumerate(y_values):        latent_point = np.array([[x, y]])        generated_image = generator_model.predict(latent_point)[0]        figure[ix*28:(ix+1)*28, iy*28:(iy+1)*28,] = generated_image[:,:,-1]plt.figure(figsize=(15, 15))plt.imshow(figure, cmap='gray', extent=[3,-3,3,-3])plt.show()
You can find all the digits(from 0 to 9) in the above image matrix as we have tried to generate images from all the portions of the latent space. The capability of generating handwriting with variations isn’t it awesome!
This tutorial explains the variational autoencoders in Deep Learning and AI. With a basic introduction, it shows how to implement a VAE with Keras and TensorFlow in python. It further trains the model on MNIST handwritten digit dataset and shows the reconstructed results.
We have seen that the latent encodings are following a standard normal distribution (all thanks to KL-divergence) and how the trained decoder part of the model can be utilized as a generative model. We have proved the claims by generating fake digits using only the decoder part of the model.
In case you are interested in reading my article on the Denoising Autoencoders
Convolutional Denoising Autoencoders for image noise reduction
Github code Link: https://github.com/kartikgill/Autoencoders
Originally published on Drops of AI.
Thanks for reading! Hope this was helpful. Kindly let me know your feedback by commenting below. See you in the next article.
Autoencoders in Keras and Deep Learning (Introduction)Optimizers explained for training Neural NetworksOptimizing TensorFlow models with Quantization TechniquesDeep Learning with PyTorch: IntroductionDeep Learning with PyTorch: First Neural Network
Autoencoders in Keras and Deep Learning (Introduction)
Optimizers explained for training Neural Networks
Optimizing TensorFlow models with Quantization Techniques
Deep Learning with PyTorch: Introduction
Deep Learning with PyTorch: First Neural Network
How to Build a Variational Autoencoder in KerasVariational autoencodershttps://keras.io/examples/generative/vae/
How to Build a Variational Autoencoder in Keras
Variational autoencoders
https://keras.io/examples/generative/vae/
Grammar Variational AutoencoderJunction Tree Variational Autoencoder for Molecular Graph GenerationVariational Autoencoder for Deep Learning of Images, Labels, and CaptionsVariational Autoencoder based Anomaly Detection using Reconstruction ProbabilityA Hybrid Convolutional Variational Autoencoder for Text Generation
Grammar Variational Autoencoder
Junction Tree Variational Autoencoder for Molecular Graph Generation
Variational Autoencoder for Deep Learning of Images, Labels, and Captions
Variational Autoencoder based Anomaly Detection using Reconstruction Probability
A Hybrid Convolutional Variational Autoencoder for Text Generation