Skip to content
Related Articles

Related Articles

Improve Article

ML | AutoEncoder with TensorFlow 2.0

  • Last Updated : 13 Jul, 2021

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 

Attention reader! Don’t stop learning now. Get hold of all the important Machine Learning Concepts with the Machine Learning Foundation Course at a student-friendly price and become industry ready.


# 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

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.


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.


class Encoder(tf.keras.layers.Layer):
    '''Encodes a digit from the MNIST dataset'''
    def __init__(self,
                name ='encoder',
        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')
    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,
                name ='decoder',
        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')
    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 avaialable 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.


class Autoencoder(tf.keras.Model):
    '''Vanilla Autoencoder for MNIST digits'''
    def __init__(self,
                 n_dims =[200, 392, 784],
                 name ='autoencoder',
        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]])
    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.


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(
                    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 to get tensor slices with a shuffled batch quickly from the dataset for training. The following code block demonstrates teh use of and also defines the hyperparameters for training the AutoEncoder model.


train_data =
        X_train).batch(128).shuffle(buffer_size = 1024)
test_data =
        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!


ae = Autoencoder([200, 392, 784])
ae.compile(optimizer = tf.optimizers.Adam(0.01),
           loss ='categorical_crossentropy'), 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.

My Personal Notes arrow_drop_up
Recommended Articles
Page :