Machine Learning course menu
Image Classification with CNNs in PyTorch
Classify images with convolutional neural networks in PyTorch: how convolutions and pooling work, building a small CNN, training on handwritten digits, comparing with a plain network, data augmentation, and transfer learning with pretrained models.
On this page
An image is a grid of pixel values. A plain network flattens that grid into one long list and loses the fact that neighbouring pixels belong together. Convolutional neural networks (CNNs) keep the grid and learn small filters that detect local patterns: edges first, then shapes built from edges, then whole objects.
Convolution#
A convolution slides a small filter (for example 3×3) across the image, computing a weighted sum at each position. This filter responds to vertical edges, where pixel values change from left to right:
import torch
import torch.nn.functional as F
image = torch.tensor([[0, 0, 0, 9, 9, 9]] * 5, dtype=torch.float32).view(1, 1, 5, 6)
vertical_edge = torch.tensor([[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3)
print(F.conv2d(image, vertical_edge)[0, 0])tensor([[ 0., 27., 27., 0.],
[ 0., 27., 27., 0.],
[ 0., 27., 27., 0.]])The output is large exactly where the dark-to-bright edge is, and zero in flat regions. In a CNN, the filter values are learned, and each layer has many filters.
Pooling#
Max pooling keeps the largest value in each small window, shrinking the image while keeping the strongest signals:
import torch
import torch.nn.functional as F
x = torch.tensor([[1, 3, 2, 0], [4, 2, 1, 1], [0, 1, 5, 6], [2, 2, 7, 1]], dtype=torch.float32).view(1, 1, 4, 4)
print(F.max_pool2d(x, 2)[0, 0])tensor([[4., 2.],
[2., 7.]])A small CNN for handwritten digits#
scikit-learn's digits dataset has 1,797 handwritten digits as 8×8 grayscale images, small enough to train in seconds on a CPU:
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
X, y = load_digits(return_X_y=True)
X = X.reshape(-1, 1, 8, 8) / 16.0 # (batch, channel, height, width), pixels scaled to 0-1
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0, stratify=y)
X_train, X_test = torch.tensor(X_train, dtype=torch.float32), torch.tensor(X_test, dtype=torch.float32)
y_train, y_test = torch.tensor(y_train), torch.tensor(y_test)
torch.manual_seed(0)
cnn = nn.Sequential(
nn.Conv2d(1, 16, kernel_size=3, padding=1), nn.ReLU(), # 16 filters -> 16 x 8 x 8
nn.MaxPool2d(2), # -> 16 x 4 x 4
nn.Conv2d(16, 32, kernel_size=3, padding=1), nn.ReLU(), # -> 32 x 4 x 4
nn.MaxPool2d(2), # -> 32 x 2 x 2
nn.Flatten(), # -> 128 values
nn.Linear(128, 10),
)
print("output shape:", cnn(X_train[:4]).shape)
optimizer = torch.optim.Adam(cnn.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
for epoch in range(1, 16):
cnn.train()
for xb, yb in loader:
optimizer.zero_grad()
loss_fn(cnn(xb), yb).backward()
optimizer.step()
if epoch in (1, 5, 15):
cnn.eval()
with torch.no_grad():
acc = (cnn(X_test).argmax(1) == y_test).float().mean().item()
print(f"epoch {epoch:>2}: test accuracy {acc:.3f}")output shape: torch.Size([4, 10])
epoch 1: test accuracy 0.767
epoch 5: test accuracy 0.929
epoch 15: test accuracy 0.971padding=1 keeps the image size the same after each 3×3 convolution; each MaxPool2d(2) halves it. nn.Flatten turns the final feature maps into one vector for the last linear layer.
Looking at mistakes#
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import load_digits
from sklearn.metrics import confusion_matrix
from sklearn.model_selection import train_test_split
X, y = load_digits(return_X_y=True)
X = X.reshape(-1, 1, 8, 8) / 16.0
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=0, stratify=y)
X_train, X_test = torch.tensor(X_train, dtype=torch.float32), torch.tensor(X_test, dtype=torch.float32)
y_train, y_test = torch.tensor(y_train), torch.tensor(y_test)
torch.manual_seed(0)
cnn = 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),
nn.Flatten(), nn.Linear(128, 10))
optimizer = torch.optim.Adam(cnn.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
for epoch in range(15):
cnn.train()
for xb, yb in loader:
optimizer.zero_grad()
loss_fn(cnn(xb), yb).backward()
optimizer.step()
cnn.eval()
with torch.no_grad():
pred = cnn(X_test).argmax(1)
wrong = (pred != y_test).nonzero().ravel()
print("wrong predictions:", len(wrong), "of", len(y_test))
for i in wrong[:5].tolist():
print(f" true {y_test[i].item()}, predicted {pred[i].item()}")
cm = confusion_matrix(y_test, pred)
print("correct per digit:", cm.diagonal().tolist())wrong predictions: 13 of 450
true 4, predicted 9
true 5, predicted 3
true 4, predicted 9
true 2, predicted 1
true 8, predicted 7
correct per digit: [45, 46, 43, 45, 43, 45, 43, 45, 38, 44]13 of 450 test digits are wrong. Digit 8 is the hardest here (38 of its test images correct), and a 4 mistaken for a 9 appears twice; at 8×8 pixels those shapes really are close. Looking at the actual mistakes, as here, shows which digits the model confuses and is often the fastest way to decide what to improve.
Data augmentation#
Real image datasets are larger and messier. Data augmentation creates varied copies of training images (small rotations, flips, crops, colour changes) so the model learns to ignore those differences. With the torchvision library:
# requires: pip install torchvision
from torchvision import transforms
train_transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])Augment the training data only; validation and test images get just the resize and normalise steps. Do not flip images when orientation matters (a flipped "6" is not a valid digit).
Transfer learning#
Training a large CNN from scratch needs a very large dataset. Transfer learning starts from a model already trained on millions of images (such as ResNet on ImageNet), keeps its learned filters, and retrains only the last layer for your classes:
# requires torchvision; downloads the pretrained weights on first use
from torch import nn
from torchvision import models
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
for param in model.parameters():
param.requires_grad = False # freeze the pretrained layers
model.fc = nn.Linear(model.fc.in_features, 5) # new final layer for 5 classesThen train as usual, passing only the trainable parameters to the optimiser. With a few hundred images per class, transfer learning commonly reaches accuracy that training from scratch would need far more data for. These last two snippets are shown for reference and were not run here, since they need extra downloads.
Common mistakes#
Practice#
Practice exercise
Work out the output shape of each layer for a batch of 10 RGB images of 32×32 pixels passing through Conv2d(3, 8, 3, padding=1), MaxPool2d(2), Conv2d(8, 16, 3) (no padding), MaxPool2d(2) and Flatten(). Then check by running a random tensor through the layers.
Show hint
Padding 1 with a 3×3 filter keeps the size; no padding shrinks each side by 2; pooling halves (rounding down).
Show solution
import torch
from torch import nn
x = torch.randn(10, 3, 32, 32)
for layer in [nn.Conv2d(3, 8, 3, padding=1), nn.MaxPool2d(2), nn.Conv2d(8, 16, 3), nn.MaxPool2d(2), nn.Flatten()]:
x = layer(x)
print(type(layer).__name__, tuple(x.shape))Conv2d (10, 8, 32, 32)
MaxPool2d (10, 8, 16, 16)
Conv2d (10, 16, 14, 14)
MaxPool2d (10, 16, 7, 7)
Flatten (10, 784)32×32 stays 32×32, pools to 16×16, shrinks to 14×14 without padding, pools to 7×7, and 16 × 7 × 7 = 784 values per image.
Key takeaway#
Next, working with text: Text Classification.
Finished this lesson?