Variational AutoEncoders and Image Generation with Keras

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-

1. Background: Variational AutoEncoders (VAEs)
2. Building VAE in Keras
3. Training VAE on the MNIST dataset
4. Results
5. Image Generation Capabilities
6. Summary
7. Further reading and resources

1. Variational AutoEncoders (VAEs)

Background

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.

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 AutoEncoders

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.

Kullback–Leibler (KL) divergence

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 later part of the tutorial.

2. Building VAE in Keras

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-

1. Data Preparation
2. Building Encoder
3. Latent Distribution and Sampling
4. Building Decoder
5. Building VAE
6. Loss

Data Preparation

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 np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import warnings
warnings.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 the 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')/255
test_data = testX.astype('float32')/255

train_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)

(60000, 28, 28, 1) (10000, 28, 28, 1)

Building Encoder

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 tensorflow

input_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.

Latent Distribution and Sampling

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()

Model: "functional_1"
__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_1 (InputLayer)            [(None, 28, 28, 1)]  0                                            
__________________________________________________________________________________________________
conv2d (Conv2D)                 (None, 24, 24, 64)   1664        input_1[0][0]                    
__________________________________________________________________________________________________
max_pooling2d (MaxPooling2D)    (None, 12, 12, 64)   0           conv2d[0][0]                     
__________________________________________________________________________________________________
conv2d_1 (Conv2D)               (None, 10, 10, 64)   36928       max_pooling2d[0][0]              
__________________________________________________________________________________________________
max_pooling2d_1 (MaxPooling2D)  (None, 5, 5, 64)     0           conv2d_1[0][0]                   
__________________________________________________________________________________________________
conv2d_2 (Conv2D)               (None, 3, 3, 32)     18464       max_pooling2d_1[0][0]            
__________________________________________________________________________________________________
max_pooling2d_2 (MaxPooling2D)  (None, 1, 1, 32)     0           conv2d_2[0][0]                   
__________________________________________________________________________________________________
flatten (Flatten)               (None, 32)           0           max_pooling2d_2[0][0]            
__________________________________________________________________________________________________
dense (Dense)                   (None, 16)           528         flatten[0][0]                    
__________________________________________________________________________________________________
mean (Dense)                    (None, 2)            34          dense[0][0]                      
__________________________________________________________________________________________________
log_variance (Dense)            (None, 2)            34          dense[0][0]                      
__________________________________________________________________________________________________
lambda (Lambda)                 (None, 2)            0           mean[0][0]                       
                                                                 log_variance[0][0]               
==================================================================================================
Total params: 57,652
Trainable params: 57,652
Non-trainable params: 0

Encoder is quite simple with just around 57K trainable parameters.

Building Decoder

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()

Model: "functional_3"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_2 (InputLayer)         [(None, 2)]               0         
_________________________________________________________________
dense_1 (Dense)              (None, 64)                192       
_________________________________________________________________
reshape (Reshape)            (None, 1, 1, 64)          0         
_________________________________________________________________
conv2d_transpose (Conv2DTran (None, 3, 3, 64)          36928     
_________________________________________________________________
conv2d_transpose_1 (Conv2DTr (None, 5, 5, 64)          36928     
_________________________________________________________________
up_sampling2d (UpSampling2D) (None, 10, 10, 64)        0         
_________________________________________________________________
conv2d_transpose_2 (Conv2DTr (None, 12, 12, 64)        36928     
_________________________________________________________________
up_sampling2d_1 (UpSampling2 (None, 24, 24, 64)        0         
_________________________________________________________________
conv2d_transpose_3 (Conv2DTr (None, 28, 28, 1)         1601      
=================================================================
Total params: 112,577
Trainable params: 112,577
Non-trainable params: 0
_________________________________________________________________

Decoder is again simple with 112K trainable parameters.

Building VAE (Variational Autoencoders)

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()

Model: "functional_5"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 28, 28, 1)]       0         
_________________________________________________________________
functional_1 (Functional)    (None, 2)                 57652     
_________________________________________________________________
functional_3 (Functional)    (None, 28, 28, 1)         112577    
=================================================================
Total params: 170,229
Trainable params: 170,229
Non-trainable params: 0
_________________________________________________________________

The overall setup is quite simple with just 170K trainable model parameters. Time to write the objective(or optimization function) function.

Loss

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-

def get_loss(distribution_mean, distribution_variance):
    
    def get_reconstruction_loss(y_true, y_pred):
        reconstruction_loss = tensorflow.keras.losses.mse(y_true, y_pred)
        reconstruction_loss_batch = tensorflow.reduce_mean(reconstruction_loss)
        return reconstruction_loss_batch*28*28
    
    def get_kl_loss(distribution_mean, distribution_variance):
        kl_loss = 1 + distribution_variance - tensorflow.square(distribution_mean) - tensorflow.exp(distribution_variance)
        kl_loss_batch = tensorflow.reduce_mean(kl_loss)
        return kl_loss_batch*(-0.5)
    
    def total_loss(y_true, y_pred):
        reconstruction_loss_batch = get_reconstruction_loss(y_true, y_pred)
        kl_loss_batch = get_kl_loss(distribution_mean, distribution_variance)
        return reconstruction_loss_batch + kl_loss_batch
    
    return total_loss

Finally let’s compile the model to make it ready for the training-

autoencoder.compile(loss=get_loss(distribution_mean, distribution_variance), optimizer='adam')

3. Training VAE (Variational Autoencoders)

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))

Train on 60000 samples, validate on 10000 samples
Epoch 1/20
60000/60000 [==============================] - 82s 1ms/sample - loss: 43.5876 - val_loss: 39.2346
Epoch 2/20
60000/60000 [==============================] - 83s 1ms/sample - loss: 37.9965 - val_loss: 37.2783
Epoch 3/20
60000/60000 [==============================] - 84s 1ms/sample - loss: 36.7326 - val_loss: 36.1599
Epoch 4/20
60000/60000 [==============================] - 85s 1ms/sample - loss: 36.1045 - val_loss: 35.7573
Epoch 5/20
60000/60000 [==============================] - 86s 1ms/sample - loss: 35.6746 - val_loss: 35.6806
Epoch 6/20
60000/60000 [==============================] - 87s 1ms/sample - loss: 35.3674 - val_loss: 35.6007
Epoch 7/20
60000/60000 [==============================] - 88s 1ms/sample - loss: 35.1614 - val_loss: 34.8265
Epoch 8/20
60000/60000 [==============================] - 91s 2ms/sample - loss: 34.9543 - val_loss: 34.8720
Epoch 9/20
60000/60000 [==============================] - 93s 2ms/sample - loss: 34.7366 - val_loss: 34.9417
Epoch 10/20
60000/60000 [==============================] - 84s 1ms/sample - loss: 34.6108 - val_loss: 34.6359
Epoch 11/20
60000/60000 [==============================] - 85s 1ms/sample - loss: 34.4663 - val_loss: 34.6572
Epoch 12/20
60000/60000 [==============================] - 85s 1ms/sample - loss: 34.4012 - val_loss: 34.4730
Epoch 13/20
60000/60000 [==============================] - 85s 1ms/sample - loss: 34.2638 - val_loss: 34.2877
Epoch 14/20
60000/60000 [==============================] - 86s 1ms/sample - loss: 34.1820 - val_loss: 34.2911
Epoch 15/20
60000/60000 [==============================] - 86s 1ms/sample - loss: 34.0327 - val_loss: 34.0680
Epoch 16/20
60000/60000 [==============================] - 86s 1ms/sample - loss: 33.9865 - val_loss: 34.1897
Epoch 17/20
60000/60000 [==============================] - 85s 1ms/sample - loss: 33.9136 - val_loss: 34.0406
Epoch 18/20
60000/60000 [==============================] - 86s 1ms/sample - loss: 33.8506 - val_loss: 34.2033
Epoch 19/20
60000/60000 [==============================] - 86s 1ms/sample - loss: 33.7726 - val_loss: 33.8551
Epoch 20/20
60000/60000 [==============================] - 86s 1ms/sample - loss: 33.6502 - val_loss: 33.7459

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.

4. Results

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=400
print ("Real Test Images")
# Real Images
for i in range(9):
    plt.subplot(330 + 1 + i)
    plt.imshow(test_data[i+offset,:,:, -1], cmap='gray')
plt.show()

# Reconstructed Images
print ("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 like 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).

Latent Feature Clusters

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(same color represents the digits belonging to 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'] = x
df['y'] = y
df['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.

5. Fake Image Generation

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_model

x_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. A capability of generating handwriting with variations, isn’t it awesome!

Summary

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 Link: https://github.com/kartikgill/Autoencoders

Thanks for reading! Hope this was helpful. Kindly let me know your feedback by commenting below. See you in the next article.

Read Next >>>

1. Autoencoders in Keras and Deep Learning (Introduction)
2. Optimizers explained for training Neural Networks
3. Optimizing TensorFlow models with Quantization Techniques
4. Deep Learning with PyTorch: Introduction
5. Deep Learning with PyTorch: First Neural Network

References

1. How to Build a Variational Autoencoder in Keras
2. Variational autoencoders
3. https://keras.io/examples/generative/vae/

Related Papers

1. Grammar Variational Autoencoder
2. Junction Tree Variational Autoencoder for Molecular Graph Generation
3. Variational Autoencoder for Deep Learning of Images, Labels and Captions
4. Variational Autoencoder based Anomaly Detection using Reconstruction Probability
5. A Hybrid Convolutional Variational Autoencoder for Text Generation