Autoencoder in PyTorch for the Fashion-MNIST dataset
By: Javier Fernandez
Tags: Convolutional Network, Machine Learning, Pytorch, Python, Deep Learning

Autoencoder in PyTorch for the Fashion-MNIST dataset

Photo by Negative Space on Pexels

This article is written for people who want to build a basic autoencoder using PyTorch. The dataset in which this article is based on is the Fashion-MNIST dataset.

The article is divided into the following sections:

Introduction to autoencoders. Loading the dataset. Building the neural network architecture. Training the neural network. Visualization in Tensorboard.

1. Introducing autoencoders

As defined in Wikipedia:

An autoencoder is a type of neural network used to learn efficient data codings in an unsupervised manner.

In other words, the aim of an autoencoder is to learn a lower representation of a set of data, which is useful for feature extraction, dimensionality reduction, and image denoising tasks, among others. Its architecture consists of an encoder, which compresses the input into a latent space representation also known as bottleneck, and a decoder, which reconstructs the input from the latent space representation.

Autoencoder architecture

2. Loading the dataset

To showcase how to build an autoencoder in PyTorch, I have decided the well-known Fashion-MNIST dataset.

Fashion-MNIST is a dataset of Zalando’s article images — consisting of a training set of 60,000 examples and a test set of 10,000 examples. Each example is a 28x28 grayscale image, associated with a label from 10 classes. We intend Fashion-MNIST to serve as a direct drop-in replacement for the original MNIST dataset for benchmarking machine learning algorithms. It shares the same image size and structure of training and testing splits.

Firstly, in order to run the rest of the code, it is required to install and import the following libraries.

Libraries

In order to run the code displayed below, it is necessary to download the following files. Once downloaded the data, locate them under a folder named data and run the following code.

Load the data

Once loaded, please check that the shape of train_x is (1000, 28, 28) and test_x is(100, 28, 28). Subsequently, normalize the data in the range[0, 1].

Normalization

Lastly, reshape the arrays with the shape that PyTorch requires for the convolutional layers (train_x: (1000, 1, 28, 28), and test_x: (100, 1, 28, 28)), divide the train_x array into train and validation sets and load the three sets into PyTorch Dataloaders.

3. Building the neural network architecture

As mentioned above, an autoencoder consists of an encoder and a decoder. Despite the possibility of using fully connected layers after flattening the input image, I have decided to use convolutional layers for both the encoder and decoder.

4. Training the neural network

For training the neural network, you can run the following code. Apart from training and evaluating the validation set, it will also save a log file that can be afterwards loaded into Tensorboard.

5. Visualization in Tensorboard

The log file can be opened in Tensorboard running the following command in the terminal: tensorboard --logdir logs/ --port 8080

The training loss and the validation loss are the ones displayed in the following figures.

The x-axis indicates the epoch number, and the y-axis shows the training loss

The x-axis indicates the epoch number, and the y-axis shows the validation loss

For the learning rate, I set an exponential decay which follows the following mathematical formula: lr = lr0 * e^(−kt) , where lr , k are hyperparameters (k=0.999) and t is the iteration number (epoch number).

The x-axis indicates the epoch number and the y-axis shows the learning rate

Here are four images of the validation set:

Ground truth images

The next images display the output of the autoencoder after 1, 125, and 250 epochs. We can observe how after 125 epochs the autoencoder can reconstruct original images accurately.

Reconstruction after the first epoch

Reconstruction after 125 epochs

Reconstruction after 250epochs

Lastly, here are eight reconstructed images of the test set.