Train a Deep Learning Model With Pytorch

Neural Network is a type of machine learning model inspired by the structure and function of human brain. It consists of layers of interconnected nodes called neurons which process and transmit information. Neural networks are particularly well-suited for tasks such as image and speech recognition, natural language processing and making predictions based on large amounts of data.

Below is an image of a simple neural network:

In this article we will train a neural network on the MNIST dataset. It is a dataset of handwritten digits consisting of 60,000 training examples and 10,000 test examples. Each example is a 28x28 grayscale image of a handwritten digit with values ranging from 0 (white) to 255 (black). The label for each example is the digit that the image represents with values ranging from 0 to 9. It is a dataset commonly used for training and evaluating image classification models particularly in the field of computer vision.

Installing PyTorch

To install PyTorch you will need to have Python and pip (the package manager for Python) installed on your computer. We can install PyTorch and necessary libraries by running the following command in your command prompt or terminal:

pip install torch
pip install torchvision
pip install torchsummary

Implementing Deep Neural Network

The model will be a simple feedforward neural network and we’ll walk through the steps to implement it.

1. Importing necessary Libraries

We will be importing Pytorch and matplotlib.

• torch: PyTorch package for tensor operations.
• torchvision: A package for vision related datasets and models.

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt

2. Loading MNIST Datasets

First we need to import the necessary libraries and load the dataset which can be loaded using the torchvision library. We will then convert images to tensors and normalizes their pixel values to have a mean of 0.5 and standard deviation of 0.5.

• transforms.Compose: A sequence of transformations to apply to the images.
• ToTensor(): Converts images into PyTorch tensors.
• Normalize(): Normalizes the pixel values to the range [-1, 1].
• DataLoader: DataLoader is a utility class that handles the batching, shuffling and loading of data for training and evaluation.
• batch_size=64 means that in each iteration 64 samples will be processed at once.

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)

trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False)

3. Building model

Let’s define a simple feedforward neural network with one hidden layer. This model will take 28x28 images, flatten them and pass them through a fully connected layer followed by an output layer with 10 units (for the 10 digits).

• nn.Module: The base class for all neural network modules in PyTorch. By inheriting from this class we create a custom model with layers and a forward pass.
• nn.Linear: This is a basic layer where each input is connected to every output node.
• fc1: The first fully connected layer transforms the 28x28 image (flattened to a 784-length vector) into a 128-dimensional vector.
• fc2: The second fully connected layer outputs a 10-dimensional vector corresponding to the 10 possible digit classes.
• We use ReLU activation after the first layer to introduce non-linearity.

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(28*28, 128) 
        self.fc2 = nn.Linear(128, 10)  

    def forward(self, x):
        x = x.view(-1, 28*28)  
        x = torch.relu(self.fc1(x))
        x = self.fc2(x) 
        return x

model = SimpleNN()

4. Defining the Loss Function and Optimizer

For classification tasks we use CrossEntropyLoss as the loss function which is appropriate for multi-class classification. The Adam optimizer will be used to update the weights during training.

• CrossEntropyLoss: Measures the difference between the predicted and true labels.
• Adam: A popular optimization algorithm that adapts the learning rate during training.
• lr=0.001 sets the learning rate which determines how much we adjust the weights during each update.

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(model.parameters(), lr=0.001)

5. Training the Model

Now we will train the model. In each epoch we will perform a forward pass to compute predictions, calculate the loss then perform backpropagation to compute gradients. Finally we update the model weights using the optimizer. After every 100 mini-batches the average loss for that mini-batch is printed. The total loss for each epoch is also tracked for later visualization.

• loss.backward(): Computes gradients for each model parameter.
• optimizer.step(): Updates the model's weights using the gradients.
• cuda(): If a GPU is available we move the images, labels and model to the GPU for faster computation.
• optimizer.zero_grad(): Before performing a backward pass we need to set all the gradients to zero. Otherwise the gradients will accumulate from previous iterations.
• loss.backward(): This computes the gradients for the model’s weights indicating how much the weights need to be adjusted.
• optimizer.step(): This updates the model’s weights based on the gradients calculated during the backward pass.

train_losses = []

num_epochs = 5 
for epoch in range(num_epochs):
    running_loss = 0.0
    for i, (images, labels) in enumerate(trainloader, 0):
        if torch.cuda.is_available():
            images, labels = images.cuda(), labels.cuda()
            model.cuda()
        else:
            model.cpu()

        optimizer.zero_grad()

        outputs = model(images)
        loss = criterion(outputs, labels)

        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        
        if i % 200 == 99:
            print(f"[Epoch {epoch+1}, Batch {i+1}] Loss: {running_loss / 100:.4f}")
            running_loss = 0.0

    train_losses.append(running_loss / len(trainloader))

Output:

6. Ploting Training and Validation Curves

The Training Loss Curve shows how the model's loss decreases over epochs indicating that the model is learning and improving its performance. It is useful to visualize how the training loss changes over time. This helps us understand if the model is learning effectively.

plt.plot(train_losses, label="Training Loss")
plt.title('Training Loss Curve')
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.legend()
plt.show()

Output:

The training loss decreases steadily across epochs with a sharper drop between epochs 1 and 2. A slight increase at epoch 3 suggests a temporary plateau but the loss resumes its decline by epoch 4 indicating overall improvement in model performance.

7. Evaluating the Model

Once the model is trained we evaluate its performance on the test dataset. We compute the accuracy by comparing the predicted labels with the true labels.

• model.eval(): Sets the model to evaluation mode disabling dropout layers.
• torch.no_grad(): Ensures that no gradients are calculated during evaluation saving memory.

correct = 0
total = 0

model.eval()

with torch.no_grad():
    for images, labels in testloader:
        if torch.cuda.is_available():
            images, labels = images.cuda(), labels.cuda()
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = 100 * correct / total
print(f"Test Accuracy: {accuracy:.2f}%")

Output:

Test Accuracy: 97.48%

We achieved a good accuracy of 97.48% which is great for a deep learning model.