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

filter_none

edit
close

play_arrow

link
brightness_4
code

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

chevron_right


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

filter_none

edit
close

play_arrow

link
brightness_4
code

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

chevron_right


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.



filter_none

edit
close

play_arrow

link
brightness_4
code

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)

chevron_right


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.

filter_none

edit
close

play_arrow

link
brightness_4
code

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)

chevron_right


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.

filter_none

edit
close

play_arrow

link
brightness_4
code

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)

chevron_right


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 teh use of tf.data and also defines the hyperparameters for training the AutoEncoder model.

filter_none

edit
close

play_arrow

link
brightness_4
code

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

chevron_right


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!

filter_none

edit
close

play_arrow

link
brightness_4
code

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)

chevron_right


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

Attention reader! Don’t stop learning now. Get hold of all the important CS Theory concepts for SDE interviews with the CS Theory Course at a student-friendly price and become industry ready.




My Personal Notes arrow_drop_up

Check out this Author's contributed articles.

If you like GeeksforGeeks and would like to contribute, you can also write an article using contribute.geeksforgeeks.org or mail your article to contribute@geeksforgeeks.org. See your article appearing on the GeeksforGeeks main page and help other Geeks.

Please Improve this article if you find anything incorrect by clicking on the "Improve Article" button below.