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:

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 28×28 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.

An sample data from Mnist

Installing PyTorch

To install PyTorch you will need to have Python and pip (the package manager for Python) installed on your computer.

You 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.

Step 1: Import the necessary Libraries

We will be importing Pytorch and matplotlib.

• torch: The main 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

Step 2: Load the MNIST Datasets

First, we need to import the necessary libraries and load the datase which can be easily loaded using the torchvision library.

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)

• 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.

Step 3: Build the model

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

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()

• 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 28×28 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.
• ReLU: We use ReLU activation after the first layer to introduce non-linearity.

Step 4: Define 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.

criterion = nn.CrossEntropyLoss()

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

• 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.

Step 5: Train the Model

Now we will train the model. In each epoch we will:

• Perform a forward pass to compute predictions.
• Calculate the loss.
• Perform backpropagation to compute gradients.
• Update the model weights using the optimizer.

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 % 100 == 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:

[Epoch 1, Batch 100] Loss: 0.8981
[Epoch 1, Batch 200] Loss: 0.4278
[Epoch 1, Batch 300] Loss: 0.3847
.
.
[Epoch 5, Batch 900] Loss: 0.1005

• 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.
• 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.

Step 6: Plot Training and Validation Curves

It’s 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:

Training Validation Graph

The Training Loss Curve shows how the model’s loss decreases over epochs indicating that the model is learning and improving its performance. In the first epoch there is a sharp decline in loss suggesting quick learning followed by a more gradual decrease indicating the model is stabilizing and nearing optimal performance. The loss eventually plateaus which suggests the model is converging and no longer improving rapidly a sign that training may be nearing completion without overfitting.

Step 7: Evaluation

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.

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: 96.87%

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

In the above code:

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

In summary deep learning with PyTorch is a powerful tool that can be used to build and train a wide range of models. The MNIST dataset is a good starting point to learn the basics but it’s important to keep in mind that there are many other aspects to consider when working with real-world datasets and problems. With the right knowledge and resources deep learning can be used to solve complex problems and make predictions with high accuracy.