Implementing a Convolutional Neural Network in PyTorch on the CIFAR-10 Dataset

Convolutional neural networks (CNNs) have revolutionized the field of computer vision, enabling machines to classify images with impressive accuracy. In this tutorial, we‘ll walk through how to implement a CNN using the popular deep learning library PyTorch and train it on the classic CIFAR-10 dataset. By the end, you‘ll have a solid understanding of the key components involved in building an image classifier from scratch. Let‘s dive in!

Convolutional Neural Networks: A Primer

Before we get to the implementation, let‘s review what makes CNNs so powerful for image-related tasks. A typical CNN architecture consists of three main types of layers:

1. Convolutional layers – These apply a set of learnable filters to the input, scanning across the image to detect relevant features like edges and textures. The filters are small in width and height but extend through the full depth of the input volume.

2. Pooling layers – Pooling progressively reduces the spatial size of the representation to decrease the number of parameters and amount of computation. This also makes the network invariant to small translations of the input. Max pooling is commonly used, which outputs the max value in each window.

3. Fully-connected (dense) layers – After several convolutional and max pooling layers, the high-level reasoning in the network is done via fully connected layers. These take the output of the final pooling or convolutional layer and connect it to the output classes.

The key aspects that make CNNs advantageous are:

• Local connectivity – Each neuron is connected only to a local region of the input, reducing the number of parameters.
• Parameter sharing – The same filter is used across the entire input, enabling detection of a feature regardless of its position.
• Translational invariance – Pooling gives the CNN some resilience to the exact spatial location of features.

With this foundation in mind, it‘s time to see CNNs in action on a real dataset.

The CIFAR-10 Dataset

CIFAR-10 is a widely used dataset for benchmarking image classification models. It consists of 60,000 32×32 color images across 10 mutually exclusive classes:

• Airplane
• Automobile
• Bird
• Cat
• Deer
• Dog
• Frog
• Horse
• Ship
• Truck

The dataset is already split into 50,000 training images and 10,000 test images. The relatively small size and low resolution of the images makes CIFAR-10 quick to train on, while still posing enough of a challenge to be a good test of a model‘s abilities. Many state-of-the-art CNN architectures have been benchmarked on CIFAR-10 over the years.

Implementing the CNN

Now for the exciting part – building the CNN! We‘ll use PyTorch, a popular open-source machine learning library known for its dynamic computational graphs and ease of use. Here‘s a high-level overview of the steps:

1. Import the necessary libraries and load the dataset
2. Define the architecture of the CNN
3. Set up the training loop
4. Evaluate performance on the test set

Let‘s go through these one by one.

Step 1: Imports and Data Preparation

First, we‘ll import PyTorch along with torchvision, which contains utilities for working with common datasets like CIFAR-10.

import torch 
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

Next, we‘ll load the CIFAR-10 dataset and apply some transformations. We normalize the data by subtracting the mean and dividing by the standard deviation of each channel. This helps the model converge faster. We also augment the training data with random horizontal flips, which reduces overfitting.

transform_train = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261)),
])

transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261)),
])

trainset = torchvision.datasets.CIFAR10(root=‘./data‘, train=True, download=True, transform=transform_train)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root=‘./data‘, train=False, download=True, transform=transform_test)
testloader = torch.utils.data.DataLoader(testset, batch_size=128, shuffle=False, num_workers=2)

Step 2: Defining the CNN Architecture

The architecture we‘ll use consists of:

• 2 convolutional layers with 32 and 64 filters, 3×3 filter size, ReLU activation
• Max pooling layer, 2×2 window
• 2 more convolutional layers with 128 filters each
• Another max pooling layer
• 2 fully connected layers with 512 and 10 units (for the 10 classes)

Here‘s how that looks in PyTorch:

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1) 
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.conv4 = nn.Conv2d(128, 128, 3, padding=1)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.fc1 = nn.Linear(128 * 8 * 8, 512)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = self.pool1(torch.relu(self.conv2(torch.relu(self.conv1(x)))))
        x = self.pool2(torch.relu(self.conv4(torch.relu(self.conv3(x)))))
        x = torch.flatten(x, 1)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

net = Net()

The forward function defines the forward pass of the network. After each pair of convolutional layers, we apply the ReLU activation function and then max pooling. After the final pooling layer, we flatten the output and pass it through the fully connected layers. The final output is 10 values, one per class.

Step 3: Setting Up the Training Loop

With the model architecture defined, we‘re ready to train the network. We‘ll use cross-entropy loss and stochastic gradient descent with momentum.

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

for epoch in range(50):

    running_loss = 0.0
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        optimizer.zero_grad()

        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if i % 500 == 499:
            print(‘[%d, %5d] loss: %.3f‘ % (epoch + 1, i + 1, running_loss / 500))
            running_loss = 0.0

print(‘Finished Training‘)  

In each epoch, we iterate over the training set in batches. For each batch, we:

1. Zero out the parameter gradients
2. Forward pass the input through the network
3. Calculate the loss
4. Backpropagate to compute gradients
5. Update the weights using the optimizer

We print out the running loss every 500 batches to monitor progress.

Step 4: Evaluating on the Test Set

Finally, we‘ll see how well the trained model performs on data it hasn‘t seen before – the test set.

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

print(‘Accuracy of the network on the 10000 test images: %d %%‘ % (100 * correct / total))

Here we disable gradient calculation since we‘re not training. We forward pass the test images through the network, get the highest scoring class for each one, and track how many the model got correct. Printing out the final accuracy gives us a high-level assessment of the model‘s performance.

Results and Next Steps

So how did we do? On a typical run, this simple CNN architecture achieves around 75-80% accuracy on the CIFAR-10 test set. Not bad for a relatively shallow network! Of course, there‘s still plenty of room for improvement. Here are some ideas for taking this further:

• Experiment with deeper architectures like VGG, ResNet, or DenseNet
• Add regularization like dropout to combat overfitting
• Utilize transfer learning by starting with a network pretrained on a larger dataset like ImageNet
• Visualize the learned filters to gain insight into what the network is learning
• Deploy the trained model in an application that can take user images and return predicted classes

Hopefully this tutorial has given you a solid foundation for working with CNNs in PyTorch. The skills you‘ve learned – loading data, defining a model architecture, setting up a training loop, and evaluating results – will serve you well as you continue your deep learning journey. Happy coding!