PyTorch Tutorial: A Complete Use-case Example

Introduction

This tutorial shows a full use-case of PyTorch in order to explain several concepts by example. The application will be hand-written number detection using MNIST. MNIST is a popular (perhaps the most popular) educational computer vision dataset. It is composed of 70K images of hand-written digits (0-9) split into 60K-10K training and test sets respectively. The images are tiny (28x28), which makes them easy to work with.

Contents:

1. Data loading
   • Loading for tables
   • Loading for text (NLP)
   • Loading for images (CV)
2. Neural Network building
   • Skeleton
   • Layers
   • Activation functions
3. ML components
   • Loss functions
   • Optimizer
4. Training loop
5. Testing
6. Saving/loading models

PyTorch Data Loading

When using PyTorch, there are many ways to load your data. It depends mainly on the type of data (tables, images, text, audio, etc.) and the size. Many text datasets are small enough to load into memory in full. Some image datasets (such as MNIST can also be loaded to memory in full due to the small image size. However, in most real-life applications, the datasets will be too large to load into memory in full.

The way PyTorch handles this problem is simple: datasets, data loaders and batch iterators.

1. A Dataset in PyTorch contains all the data. When we initialize a dataset in PyTorch, we can also specify certain transformations to apply.

2. Data Loaders receive dataset objects as input and create a blueprint of batches.

3. Batch Iterators: Batch iterators loop over the data in batches (of 16, 32, for example) provided by the data loader. Then, a full training loop is performed on this subset. Once finished, the current batch is discarded and a new batch is loaded for training.

By using these above concepts, PyTorch is able to perform preprocessing, transformations, and training on small batches of data without running out of memory.

Let us start by importing the required libraries and tools:

Not that torch.utils.data.Dataset is the dataset class we can extend, whereas torchvision.datasets are just a group of ready to use datasets (such as MNIST) in the PyTorch library.

Since MNIST is already provided as a ready dataset, we just need to download the training and test sets as follows:

The dataset which is stored locally, you can create it as follows:

This is a pseudocode example. You should modify it according to the structure of your dataset. But the key ideas are: image paths and labels are stored, and a __getitem__() method returns an image and its label. The __len__() method is optional but useful.

Let's test training_ds and test_ds to make sure they operate as we expect:

As we can see, len(training_ds) returns the number of paths (or images) in the dataset. And, training_ds[0] returns the first image and its label. So far, so good.

Now, let's visualize a random image.

The image dimensions are 1x28x28. To visualize the image, we must have it in 2D space, or simply 28x28. To remove a dimension from a tensor, use squeeze(). Similarly, to add a dimension, use unsqueeze(). Read the docs for more examples: https://pytorch.org/docs/stable/generated/torch.squeeze.html

Great. Next are the dataloaders. Creating dataloaders in PyTorch is easy:

PyTorch DataLoaders

shuffle = True means that the dataset will be shuffled before being split into batches. This randomizes the batches which is good for generalization.

PyTorch NN Building

Using torch.nn, one can create any kind of model. In this tutorial, we explore the skeleton and guidelines to follow when creating a NN and in the process create a simple feed-forward NN (FFNN).

A NN in PyTorch is a class extending from nn.Module with __init__() and forward() methods. Of course we can add more methods, but these are the key components.

In __init__(), we create the architecture (the layers). A FFNN is composed of several fully-connected layers. Fully-connected layers are created using nn.Linear().

nn.Linear() takes in 2 arguments: number of inputs and number of outputs. When connecting FCs, you must make sure of 3 things:

• The number of inputs in the first layer must match the size of the data.
• The number of outputs of every layer must match the number of inputs in the next layer.
• The number of outputs in the final layer must match the number of classes you are working with.

Since FFNNs expect input as a vector (not a 2D tensor such as images), we cannot simply feed in the 28x28 vectors of MNIST images. We must flatten them into a 28*28 = 784 vector.

In advanced CV projects, the images will be larger than 28x28, and this approach will be unviable. For advanced CV applications, the CNN is a common architecture to use.

forward() takes in a batch and returns predictions for each class for each instance. In the forward() function, we manually pass the data from each layer to the next until the final layer.

F.relu() is a popular activation function used after FC layers. Other alternatives are `F.tanh()` and `F.sigmoid()`, but ReLu has been shown to perform better.

To initialize and use the model, let us first set the device:

Great. The data is ready and the model is ready. All that is left is the training.

Loss Functions and Optimizers

So far, we have prepared the data and the model. But, to train our model, we must perform some sort of gradient descent optimization in order to improve the model. To do so, we have to define a measure of quality. This measure is called the loss function, and it depends on the task. If the task is regression, loss functions such as MSE or RMSE can be used. For multi-label classification such as in MNIST, a common loss function is the cross-entropy loss. In short, it measures the quality of a prediction. Using this measure, we can optimize the loss of the model (reduce the error) iteratively using an optimizer. There are many optimizers proposed by the literature. The classic approach is to use Stochastic Gradient Descent (SGD), but a more popular optimizer is Adam. A lit of all optimizers in PyTorch can be found at https://pytorch.org/docs/stable/optim.html.

Let's create our loss function and optimizer:

As shown, optimizers typically take 2 arguments: the model weights to optimize, and the learning rate. Here, we select a learning rate of 0.003, but other values are also acceptable. However, a really large learning rate may cause the model to never converge, and a very small learning rate my take too long. LRs in the range of 0.001 to 0.0003 are acceptable in most cases. There are more advanced solutions to adaptively change the LR during training.

Training

Training is done in epochs. An epoch is simply 1 iteration over all the training data, typically followed by an iteration over the test data. In each epoch, we iterate over the data in batches. The general skeleton of a training epoch is as follows:

def run_epoch(loader, model, optimizer, criterion):
    for batch in loader:
        imgs, lbls = batch
        optimizer.zero_grad()
        with torch.set_grad_enabled(True): #Or False if testing
            logits = model(imgs)
            loss = criterion(logits, lbls)
            loss.backward()
            optimizer.step()

There are several things to explain in this block:

• optimizer.zero_grad()
• with torch.set_grad_enabled()
• loss.backward()

• optimizer.step()

• optimizer.zero_grad() resets the optimizer gradients to zero. This is necessary before each batch so that backpropagation is optimizing only for current batch.

• with torch.set_grad_enabled() determines whether or not to calculate gradients (i.e. calculate performance). This method takes 1 argument: True or False, depending on wheter or not you are doing training/testing. In the test loop, this must be set to False to avoid training on the test set.

• loss.backward() and optimizer.step() perform backpropagation on the current gradients and update the weights of the model to improve it.

Notice that we never call model.forward() explicitly. The forward function is simply called explicitly by model().

Finally, since this is the loop for one epoch, we can train on more epoch simply by doing this:

def main(epochs):
    for epoch in range(epochs):
        run_epoch()

Using these two functions, we can perform training and testing easily:

As shown by the accuracy scores, the model quickly learns to classify the images. At the end of training, the test accuarcy is ~98%, which is great.

We can visualize the progress of training by plotting the collected accuacies:

In both plots, the accuracy is very good near the end of training.

In classification, accuracy is only 1 metric. In real life applications, we must make sure that the data is balanced and to report recall, precision, and f1-score. These metrics can be found on the sklearn metrics module and they are used the same way we used accuracy_score()

Now, let us test the model to make sure it is actually functioning.

As shown, in almost all random test cases, the model is capable of predicting the correct class.

Saving/Loading Models

Now that we have a trained model, we should save it to disk. That way, we can quickly load it whenever we need predictions without having to train the model again. Saving and loading models is very simple in PyTorch:

Now, let's initialize a new model without loading the weights:

Since this model is untrained, we expect it to perform poorly when predicting:

As expected, it does not perform well.

Now, let's load the trained weights from disk:

Finally, let's make sure that the new model performs properly:

Great! Now we can train models and save them for later use quickly.