ML | AutoEncoder with TensorFlow 2.0

This tutorial demonstrates how to generate images of handwritten digits using graph mode execution in TensorFlow 2.0 by training an Autoencoder. 
An AutoEncoder is a data compression and decompression algorithm implemented with Neural Networks and/or Convolutional Neural Networks. the data is compressed to a bottleneck that is of a lower dimension than the initial input. The decompression uses the intermediate representation to generate the same input image again. Let us code up a good AutoEncoder using TensorFlow 2.0 which is eager by default to understand the mechanism of this algorithm. AutoEncoders are considered a good pre-requisite for more advanced generative models such as GANs and CVAEs. 
Firstly, download the TensorFlow 2.0 depending on the available hardware. If you are using Google Colab follow along with this IPython Notebook or this colab demo. Make sure that the appropriate versions of CUDA and CUDNN are available for GPU installs. Visit the official downloads instructions on the TensorFlow page here. 
Code: Importing libraries

Python3

# Install TensorFlow 2.0 by using the following command
# For CPU installation
# pip install -q tensorflow == 2.0
# For GPU installation (CUDA and CuDNN must be available)
# pip install -q tensorflow-gpu == 2.0

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function


import tensorflow as tf
print(tf.__version__)

After confirming the appropriate TF download, import the other dependencies for data augmentation and define custom functions as shown below. The standard scaler scales the data by transforming the columns. The get_random_block_from_data function is useful when using tf.GradientTape to perform AutoDiff (Automatic Differentiation) to get the gradients.

Python3

import numpy as np
import sklearn.preprocessing as prep
import tensorflow.keras.layers as layers

def standard_scale(X_train, X_test):
    preprocessor = prep.StandardScaler().fit(X_train)
    X_train = preprocessor.transform(X_train)
    X_test = preprocessor.transform(X_test)
    return X_train, X_test

def get_random_block_from_data(data, batch_size):
    start_index = np.random.randint(0, len(data) - batch_size)
    return data[start_index:(start_index + batch_size)]

AutoEncoders may have a lossy intermediate representation also known as a compressed representation. This dimensionality reduction is useful in a multitude of use cases where lossless image data compression exists. Thus we can say that the encoder part of the AutoEncoder encodes a dense representation of the data. Here we will use TensorFlow Subclassing API to define custom layers for the encoder and decoder.

Python3

class Encoder(tf.keras.layers.Layer):
    '''Encodes a digit from the MNIST dataset'''

    def __init__(self,
                n_dims,
                name ='encoder',
                **kwargs):
        super(Encoder, self).__init__(name = name, **kwargs)
        self.n_dims = n_dims
        self.n_layers = 1
        self.encode_layer = layers.Dense(n_dims, activation ='relu')

    @tf.function
    def call(self, inputs):
        return self.encode_layer(inputs)

class Decoder(tf.keras.layers.Layer):
    '''Decodes a digit from the MNIST dataset'''

    def __init__(self,
                n_dims,
                name ='decoder',
                **kwargs):
        super(Decoder, self).__init__(name = name, **kwargs)
        self.n_dims = n_dims
        self.n_layers = len(n_dims)
        self.decode_middle = layers.Dense(n_dims[0], activation ='relu')
        self.recon_layer = layers.Dense(n_dims[1], activation ='sigmoid')

    @tf.function
    def call(self, inputs):
        x = self.decode_middle(inputs)
        return self.recon_layer(x)

We then extend tf.keras.Model to define a custom model that utilizes our previously defined custom layers to form the AutoEncoder model. The call function is overridden which is the forward passwhen the data is made available to the model object. Notice the @tf.function function decorator. It ensures that the function execution occurs in a graph which speeds up our execution.

Python3

class Autoencoder(tf.keras.Model):
    '''Vanilla Autoencoder for MNIST digits'''

    def __init__(self,
                 n_dims =[200, 392, 784],
                 name ='autoencoder',
                 **kwargs):
        super(Autoencoder, self).__init__(name = name, **kwargs)
        self.n_dims = n_dims
        self.encoder = Encoder(n_dims[0])
        self.decoder = Decoder([n_dims[1], n_dims[2]])

    @tf.function
    def call(self, inputs):
        x = self.encoder(inputs)
        return self.decoder(x)

The following code block prepares the dataset and gets the data ready to be fed into the pre-processing pipeline of functions before training the AutoEncoder.

Python3

mnist = tf.keras.datasets.mnist

(X_train, _), (X_test, _) = mnist.load_data()
X_train = tf.cast(np.reshape(
        X_train, (X_train.shape[0],
                  X_train.shape[1] * X_train.shape[2])), tf.float64)
X_test = tf.cast(
        np.reshape(X_test,
                   (X_test.shape[0],
                    X_test.shape[1] * X_test.shape[2])), tf.float64)

X_train, X_test = standard_scale(X_train, X_test)

It is TensorFlow best practice to use tf.data.Dataset to get tensor slices with a shuffled batch quickly from the dataset for training. The following code block demonstrates the use of tf.data and also defines the hyperparameters for training the AutoEncoder model.

Python3

train_data = tf.data.Dataset.from_tensor_slices(
        X_train).batch(128).shuffle(buffer_size = 1024)
test_data = tf.data.Dataset.from_tensor_slices(
        X_test).batch(128).shuffle(buffer_size = 512)

n_samples = int(len(X_train) + len(X_test))
training_epochs = 20
batch_size = 128
display_step = 1

optimizer = tf.optimizers.Adam(learning_rate = 0.01)
mse_loss = tf.keras.losses.MeanSquaredError()
loss_metric = tf.keras.metrics.Mean()

We have completed every pre-requisite to train our AutoEncoder model! All we have left to do is to define an AutoEncoder object and compile the model with the optimizer and loss before calling model.train on it for the hyperparameters defined above. Voila! You can see the loss reducing and the AutoEncoder improving its performance!

Python3

ae = Autoencoder([200, 392, 784])
ae.compile(optimizer = tf.optimizers.Adam(0.01),
           loss ='categorical_crossentropy')
ae.fit(X_train, X_train, batch_size = 64, epochs = 5)

You can checkout the IPython Notebook here and a colab demo I submitted to TensorFlow here. Follow me on GitHub.