How to design a neural network for image recognition using PyTorch
By: Harsh R
Tags: Data Science, Towards Data Science, Neural Networks, Deep Learning, Image Recognition

In this post, I am going to show you how to create a neural network from ground up and use it to recognise traffic signs. This project is done by using PyTorch, so you might want to checkout my previous post on:

What is image recognition?

Image recognition/classification is a technique of computer vision, which is used to identify objects, people, writing and, in our case, signs from an image. It is the process of recognising the patterns in the image and classifying to a particular class.

This process includes two parts input and output, where:

Input : An image or set of images

Output: A class or a label to which the image belongs

So, in this project we are going to create a neural net model that will predict the label/class of an image.

Now that it’s clear what we are working on, let’s talk about data!

About Data

The German traffic dataset contains three pickle files i.e. train, test and validation. There are 12,630 images in train and test sets, and validation set has 4,410 images. The images are 32x32 pixels colored image. The class of each image is encoded as an integer ranging from 0 to 42.

German Traffic Signs

Load Data

To get started, we need to import packages.

Our data is in pickle files, so to load them we need to import pickle package.

Data Exploration

Now that we have loaded the data, lets explore the dataset using pandas and NumPy.

Preprocessing and creating loaders

Since we are using PyTorch, so we need to make our data compatible to torch tensors. To do that we need to normalize the data. We do that by copying the data into a NumPy array with integer data type and then use Normalize method to cast the array to float type.

After preprocessing the data, we create a tensor dataset which will later be used to create train, validation and test Data Loaders.

Now, we arrive at the main part.

Model creation

But first let's understand what is feed forward neural network?

A feed forward neural network is a collection of neurons organized in layers. Every node in the input layer is connected to all the nodes in the hidden layers, and some or all nodes in the hidden layer are connected to all nodes in the output layer. The number of nodes in the input layer are often equal to the number of features in a dataset, so each node represents one feature. Analogously, the number of nodes in the output layer are equal to the number of classes in the data, so when input sample of class c is presented, the c-th node in the output layer is active and all others are silent.

A neural network can be represented by just one hidden layer and sufficient number of units, it is also possible to have multiple hidden layers.

Here’s how the model looks like:

In the class definition, you can see the inheritance of the base class ImageClassificationBase. The base class has some essential steps to calculate predictions, loss and accuracy for training and validation steps.

To implement the neural network, we need to follow the official style of PyTorch. We are provided with two methods or functions that are __init__ ()and forward().

__init__()

In the __init__(), we just define the layers of a neural network. As shown, we create a simple 6 layers fully connected neural network to classify and recognize to which class does a traffic sign belong to. The input layer takes 32x32 pixel data of images. The input is then further passed to the 4 layers with 1324, 500, 256 , 128 , 64 with an output layer of size 43 because there are 43 distinct classes.

forward()

Notice that there are some functions used in forward method and not __init__ method. This is because __init__ method doesn’t have the training ability.

In forward method the first thing we do is flatten the image into a 1 dimension vector. We do this to convert a multidimensional image into a single dimension. Later we use F.relu function which is an ReLU activation function to the nodes.

So far so good? We have successfully defined a fully connected feed forward neural network!

Training Time

I recommend trying training a model on different number of epochs, different learning rates, selecting right activation function.

One of the many advantages of using PyTorch is it gives us the ability to use GPU which yields faster computations. So we move the model and weights to the cuda device by using to_device method.

After moving data to cuda device, we start the training process with the fit method. In the first training phase, I tried for 100 epochs with a learning rate of 0.01. This resulted in validation loss of 2.2009 and validation accuracy of 0.4561 .

In the second attempt, I trained the model for 50 epochs with 0.03 learning rate. As a result we got a validation loss of 3.5628 and validation accuracy of 0.5873 .

Now for the last attempt, I used same learning rate and reduced the number of epochs to 10. We see a very slight increase in validation loss to 3.6163, and the validation accuracy is 0.6362. Not bad!

Plotting losses and accuracies

Predictions

This part contains visualization on how the trained model will perform on test set.

We can go on to try with different learning rates and epochs but the accuracy gained will not differ much. This is because its a simple artificial neural network. If we want to improve our model then we need to use Convolutional Neural Networks. When we are dealing with data with higher dimensionality, Convolutional Neural Networks are majorly used to capture the relevant features of images, videos and audios. This means that Convolutional Neural Networks are a good feature extractors.