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.0from __future__ import absolute_importfrom __future__ import divisionfrom __future__ import print_functionimport tensorflow as tfprint(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 npimport sklearn.preprocessing as prepimport tensorflow.keras.layers as layersdef 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_testdef 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 = 20batch_size = 128display_step = 1optimizer = 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.
