Pytorch Tutorial Part 2: Model, Training and testing

Hi, welcome back! Today, we'll learn about building, training, and testing a model in PyTorch. Building the model was the most challenging part for me. Let's get started:

1. Model
2. Training and Testing

Model

The model can be divided into two parts: Easy part and n

Easy part

Let's start from the beginning. At first glance, the code may seem long, especially when compared to TensorFlow code for CNN. However, the process involves a series of consistent steps.

• Define the Class: Start by defining the class for your CNN, for example, class SimpleCNN(nn.Module).
• Initialize the Class: Initialize the class with the __init__ method. In this method, you define all the layers of your network. By breaking it down into these steps, the process becomes more manageable and structured.

def __init__(self):
    super(SimpleCNN, self).__init__()
    # Define layers here

• Forward Method: Write a forward method to specify the forward pass of your network. In this method, you connect the layers defined in the __init__ method.

def forward(self, x):
    # Connect layers here
    return x

Easy part done !! And it is the most of the work.

Difficult part

The challenging aspect is the nn.Linear(32 * 62 * 62, 128) layer, known as the fully connected layer. In this layer, you need to determine the number 32 * 62 * 62, which represents the input size from the preceding CNN layer.

There are two common methods to determine this number. The first method involves using a formula to calculate the output size of each layer sequentially. The second method, which we've demonstrated below, involves printing the output shape of the last CNN layer. To see this in action, click the 'Play in Colab' button. You can also experiment with different numbers to understand how the input size affects the fully connected layer.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

# Model definition
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        # First convolutional layer: 3 input channels (RGB), 16 output channels, 3x3 kernel
        self.conv1 = nn.Conv2d(3, 16, 3, 1)
        # Second convolutional layer: 16 input channels, 32 output channels, 3x3 kernel
        self.conv2 = nn.Conv2d(16, 32, 3, 1)
        # Fully connected layer: input size 32*62*62, output size 128
        self.fc1 = nn.Linear(32*62*62, 128)
        # Fully connected layer: input size 128, output size 2 (number of classes)
        self.fc2 = nn.Linear(128, 2)

    def forward(self, x):
        # Apply first convolution, followed by ReLU activation
        x = F.relu(self.conv1(x))
        # Apply second convolution, followed by ReLU activation, then max pooling with 2x2 kernel
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        # Flatten the output from the convolutional layers
        print(x.shape)
        x = x.view(-1, 32*62*62)
        # Apply first fully connected layer followed by ReLU activation
        x = F.relu(self.fc1(x))
        # Apply second fully connected layer to get the final output
        x = self.fc2(x)
        return x

# Instantiate the model
model = SimpleCNN()

Training and Testing

The training function for a PyTorch model, defined as train(model, train_loader, criterion, optimizer, epochs=5), sets the model to training mode using model.train(). It iterates over a specified number of epochs, and within each epoch, it processes each batch of images and labels from the train_loader. For each batch, the function resets the gradients with optimizer.zero_grad(), performs a forward pass to compute the outputs, calculates the loss using the specified criterion, performs a backward pass to compute the gradients, and updates the model's parameters using the optimizer's step function. The running loss for each batch is accumulated and averaged at the end of each epoch, and the loss is printed. To test the model, the test(model, test_loader) function sets the model to evaluation mode using model.eval() and disables gradient computation with torch.no_grad(). It then processes each batch of images and labels from the test_loader, performs a forward pass to compute the outputs, extracts the predicted class using torch.max(), and calculates the total and correct predictions. The function finally prints the accuracy as a percentage of correct predictions out of the total samples. To use these functions, you call train(model, train_loader, criterion, optimizer) to train the model and test(model, test_loader) to evaluate its performance.

# Function to train the model
def train(model, train_loader, criterion, optimizer, epochs=5):
    model.train()  # Set the model to training mode
    for epoch in range(epochs):
        running_loss = 0.0
        for images, labels in train_loader:
            optimizer.zero_grad()  # Zero the parameter gradients
            outputs = model(images)  # Forward pass
            loss = criterion(outputs, labels)  # Compute the loss
            loss.backward()  # Backward pass
            optimizer.step()  # Optimize the parameters
            running_loss += loss.item()  # Accumulate the loss
        print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}")

# Train the model
train(model, train_loader, criterion, optimizer)

# Function to test the model
def test(model, test_loader):
    model.eval()  # Set the model to evaluation mode
    correct = 0
    total = 0
    with torch.no_grad():  # Disable gradient computation
        for images, labels in test_loader:
            outputs = model(images)  # Forward pass
            _, predicted = torch.max(outputs.data, 1)  # Get the predicted class
            total += labels.size(0)  # Update total number of labels
            correct += (predicted == labels).sum().item()  # Update number of correct predictions
    print(f'Accuracy: {100 * correct / total}%')

# Test the model
test(model, test_loader)

Thats it 😄, I hope I make it easy for you, please comment below if you want to know more about it.