Convolutional neural networks (CNNs)

Your phone camera recognises faces, programs that assist doctors find changes on X-ray images, cars read road signs. Most of these rely on convolutional neural networks (CNNs). In the previous lesson we saw that a fully connected layer for a 224 × 224 colour image needs over 150 million parameters. CNNs exploit two simple observations: meaningful patterns in an image are local (edges, corners, textures), and the same pattern can appear anywhere in the picture. So a small filter is “slid” over the whole image, and its weights are shared everywhere.

Images as tensors

A grayscale 28 × 28 image is a tensor of shape (1, 28, 28): one channel, 28 rows, 28 columns. A colour image has three channels — red, green, blue: (3, H, W). For batches PyTorch uses the (N, C, H, W) layout. Pixels are usually integers 0–255; transforms.ToTensor() scales them to [0, 1], and then each channel is normalised.

The convolution operation

A filter (kernel) is a small matrix of weights, for example 3 × 3. It slides over the image, and at every position it is multiplied element-wise with the pixels beneath it and summed. The result is a feature map: it has large values wherever the pattern the filter “looks for” is present. A convolutional layer has several filters, each producing its own output channel; every filter spans all the input channels.

The code below computes a convolution from scratch with NumPy. The left half of the image is dark (0) and the right half bright (9); the filter responds to brightness increasing from left to right. Run the code and check the result yourself:

import numpy as np

image = np.array([
    [0, 0, 0, 9, 9, 9],
    [0, 0, 0, 9, 9, 9],
    [0, 0, 0, 9, 9, 9],
    [0, 0, 0, 9, 9, 9],
    [0, 0, 0, 9, 9, 9],
])
kernel = np.array([
    [-1, 0, 1],
    [-1, 0, 1],
    [-1, 0, 1],
])
K = kernel.shape[0]
H, W = image.shape
out = np.zeros((H - K + 1, W - K + 1), dtype=int)
for i in range(out.shape[0]):
    for j in range(out.shape[1]):
        out[i, j] = np.sum(image[i:i + K, j:j + K] * kernel)
print(out)

[[ 0 27 27  0]
 [ 0 27 27  0]
 [ 0 27 27  0]]

In trained CNNs, filters in the first layers typically respond to edges and colour transitions, middle layers to textures and parts (eyes, wheels), and the last layers to whole objects. After each convolution and pooling layer, the region a neuron “sees” — its receptive field — grows.

import torch
from torch import nn

torch.manual_seed(42)
x = torch.randn(8, 1, 28, 28)
conv = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, padding=1)
print(conv(x).shape)
print(nn.Conv2d(1, 16, kernel_size=5)(x).shape)
print(nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1)(x).shape)
print(nn.MaxPool2d(kernel_size=2)(conv(x)).shape)
print(conv.weight.shape, sum(p.numel() for p in conv.parameters()))

torch.Size([8, 16, 28, 28])
torch.Size([8, 16, 24, 24])
torch.Size([8, 16, 14, 14])
torch.Size([8, 16, 14, 14])
torch.Size([16, 1, 3, 3]) 160

Pooling and a CNN for 28 × 28 images

Max pooling keeps only the largest value from each window (usually 2 × 2). It halves the size of the map, makes computation cheaper, adds robustness to small shifts and has no parameters at all. Average pooling takes the mean instead.

Now let us build a small CNN for MNIST: two “convolution → ReLU → max pooling” blocks extract features, and at the end Flatten plus a fully connected layer produce logits for the 10 digits. We print the shape after every layer — the most reliable way to debug a CNN.

import torch
from torch import nn

class SmallCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
        )
        self.classifier = nn.Sequential(nn.Flatten(), nn.Linear(32 * 7 * 7, 10))

    def forward(self, x):
        return self.classifier(self.features(x))

torch.manual_seed(42)
model = SmallCNN()
x = torch.randn(64, 1, 28, 28)
for layer in model.features:
    x = layer(x)
    print(f'{type(layer).__name__:9s} {tuple(x.shape)}')
print(model(torch.randn(64, 1, 28, 28)).shape)
print(sum(p.numel() for p in model.parameters()))

Conv2d    (64, 16, 28, 28)
ReLU      (64, 16, 28, 28)
MaxPool2d (64, 16, 14, 14)
Conv2d    (64, 32, 14, 14)
ReLU      (64, 32, 14, 14)
MaxPool2d (64, 32, 7, 7)
torch.Size([64, 10])
20490

MNIST with torchvision

The torchvision package (pip install torchvision) provides ready-made datasets (MNIST, CIFAR-10, the ImageNet format), image transforms (transforms) and pretrained models. MNIST consists of 70,000 images of handwritten digits: 60,000 for training and 10,000 for testing, each 28 × 28 grayscale. download=True downloads the files into the data folder the first time and reads them from disk afterwards.

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

torch.manual_seed(42)
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,)),
])
train_set = datasets.MNIST('data', train=True, download=True, transform=transform)
test_set = datasets.MNIST('data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_set, batch_size=64, shuffle=True)
test_loader = DataLoader(test_set, batch_size=1000)
images, labels = next(iter(train_loader))
print(len(train_set), len(test_set))
print(images.shape, labels.shape)

60000 10000
torch.Size([64, 1, 28, 28]) torch.Size([64])

Transfer learning

Training a large CNN from scratch needs millions of images and powerful GPUs. But the edges, textures and shapes learned by the first layers of a network trained on ImageNet (1000 classes, over a million training images) are useful for almost any image task. Transfer learning works like this: take a pretrained model, freeze its layers, replace the last classification layer with a new one for your own classes and train only that.

import torch
from torch import nn
from torchvision import models

model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
for p in model.parameters():
    p.requires_grad = False
print(model.fc)
model.fc = nn.Linear(model.fc.in_features, 3)
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
print(trainable, total)

Linear(in_features=512, out_features=1000, bias=True)
1539 11178051

The new layer has (512 + 1) · 3 = 1,539 parameters. So few parameters can be trained quickly with a few hundred images, even on a CPU. With more data, in a second stage you can “unfreeze” the last few layers too and fine-tune the whole network with a very small learning rate. Give the optimizer only the trainable parameters: torch.optim.Adam(model.fc.parameters(), lr=1e-3).