Building and Training a Simple CNN for Image Classification

1. The Five Main Steps

Building and training a model in PyTorch generally follows a consistent, five-step process. We'll walk through each one:

1. Load and Prepare Data: We need to get our images and labels, convert them into tensors, and group them into batches.
2. Define the Model: We'll create a Python class for our CNN architecture, defining the layers in the __init__ method.
3. Define the Forward Pass: Inside our model class, we'll specify how data flows through the layers in the forward method.
4. Set up the Training Components: We'll choose a loss function and an optimizer to guide the training process.
5. Build the Training and Evaluation Loop: We'll write the code to repeatedly feed data to the model, calculate loss, backpropagate, update weights, and check performance.

Let's get started.

2. Step 1: Loading the MNIST Dataset

First, we need our data. PyTorch's torchvision library makes it incredibly easy to download standard datasets like MNIST. We'll also use DataLoader to efficiently manage the data by creating shuffled batches for training.

Here is a summary of the code structure for loading the data:

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# Define a transform to convert images to tensors
transform = transforms.ToTensor()

# Download and load the training data
train_data = datasets.MNIST(root='data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_data, batch_size=100, shuffle=True)

# Download and load the test data
test_data = datasets.MNIST(root='data', train=False, download=True, transform=transform)
test_loader = DataLoader(test_data, batch_size=100, shuffle=False)

The train_loader and test_loader are now iterators that will provide us with batches of images and their corresponding labels.

3. Steps 2 & 3: Defining the CNN Model and Forward Pass

This is where we build the heart of our project. We will create a CNN class that inherits from PyTorch's nn.Module.

• In the __init__ method, we define the layers we'll use: nn.Conv2d for convolutions, nn.MaxPool2d for pooling, and nn.Linear for the final classification layers.
• In the forward method, we define the computational graph—how an input x flows through these layers and activation functions (F.relu) to produce an output.

A critical step is calculating the input size for the first fully connected (nn.Linear) layer. This depends on the output dimensions of the final pooling layer. For example, if the last pooling layer outputs a feature map of size 16 x 4 x 4, we flatten it into a vector of size 16 * 4 * 4 = 256.

The resulting model class will look something like this. Notice how the forward method is a recipe for passing data through the layers defined in __init__.

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

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # Input for MNIST is 1 channel (grayscale), output is 10 feature maps
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=10, kernel_size=5)
        # Second conv layer takes 10 channels (from conv1), outputs 20
        self.conv2 = nn.Conv2d(in_channels=10, out_channels=20, kernel_size=5)
        # Max pooling layer
        self.pool = nn.MaxPool2d(kernel_size=2)
        # Dropout layer for regularization
        self.conv2_drop = nn.Dropout2d()
        # Fully connected layers
        # The input features (320) is calculated from the output of the conv/pool layers
        # (20 channels * 4x4 image size)
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10) # 10 output classes for digits 0-9

    def forward(self, x):
        # First conv layer -> relu -> pool
        x = self.pool(F.relu(self.conv1(x)))
        # Second conv layer -> dropout -> relu -> pool
        x = self.pool(F.relu(self.conv2_drop(self.conv2(x))))
        # Flatten the feature maps to a 1D vector
        x = x.view(-1, 320)
        # Pass through fully connected layers
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        # Use log_softmax for the final output, which works well with NLLLoss
        return F.log_softmax(x, dim=1)

4. Step 4: Setting Up the Training Components

With our model and data ready, we need to define our loss function and optimizer.

• Loss Function: Since this is a multi-class classification problem, nn.CrossEntropyLoss is a standard choice. If your model's last layer is log_softmax (as in the video example), you should use nn.NLLLoss (Negative Log Likelihood Loss).
• Optimizer: The Adam optimizer is a robust and popular choice that often works well with little hyperparameter tuning.

This setup is very similar to what you did for the feedforward network.

5. Step 5: Building the Training and Evaluation Loop

Now we orchestrate the entire process. We will create a train function and a test function.

The training loop will:

1. Set the model to training mode: model.train().
2. Iterate through the train_loader.
3. For each batch, perform the forward pass, calculate the loss, perform backpropagation (loss.backward()), and update the weights (optimizer.step()).

The evaluation loop will:

1. Set the model to evaluation mode: model.eval().
2. Disable gradient calculations with with torch.no_grad():.
3. Iterate through the test_loader.
4. Calculate the total loss and the number of correct predictions to determine the model's accuracy.

Finally, we put it all together by calling these functions inside a loop that runs for a specified number of epochs.

You'll see the test accuracy climb with each epoch, likely reaching over 98%, which is a great result and significantly better than what a simple FNN could achieve with the same amount of effort.

For reference, the articles from DataCamp and Nextjournal provide excellent text-based walkthroughs of this entire process. They are great resources to review alongside the video.

Conclusion

Congratulations! You have just walked through the complete process of building and training a Convolutional Neural Network in PyTorch. This is a milestone in your deep learning journey.

Key Takeaways:

• Model Definition: You now know how to define a custom CNN architecture by creating a class that inherits from nn.Module, specifying layers like nn.Conv2d and nn.MaxPool2d in __init__.
• Forward Pass: You learned to define the data flow through your network in the forward method, including applying activations and flattening feature maps.
• Training and Evaluation: You have seen how to plug your new CNN model into a standard PyTorch training loop to train its parameters and then evaluate its performance on unseen test data.
• Putting it all Together: You understand the end-to-end workflow: loading data, defining a model, setting up the optimizer/loss, and executing the training and evaluation process.

Preview of the Next Lesson:
So far, we've focused on data that is spatial, like images. However, many important problems in AI involve data that is sequential, where order matters, such as text, speech, or time-series data. CNNs are not naturally suited for this.

In our next lesson, we will explore Recurrent Neural Networks (RNNs) and their powerful variant, Long Short-Term Memory (LSTM) networks. You will learn how these architectures are designed to handle sequential data by maintaining an internal "memory" of past information.