Building a CNN in PyTorch — Deep Learning with PyTorch

You know what a convolution computes; now let's assemble the real thing. In this lesson you'll build the full image-classification pipeline: load a dataset with torchvision, stack conv-relu-pool blocks into a CNN, train it with the recipe from the last module, measure whether data augmentation helps, and — most instructive of all — look at the images it gets wrong.

This lesson needs a GPU

Run the notebook in Google Colab and enable the free GPU first: Runtime → Change runtime type → T4 GPU. Training takes a couple of minutes on GPU versus ~20x longer on CPU. Verify with torch.cuda.is_available() — it should print True.

Loading images with torchvision

torchvision bundles standard datasets and image transforms. We'll use FashionMNIST: 70,000 grayscale 28×28 images of clothing in 10 classes — big enough to be interesting, small enough to train in minutes:

import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
 
device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)
 
transform = transforms.Compose([
    transforms.ToTensor(),                      # PIL image -> [1, 28, 28] float in [0, 1]
    transforms.Normalize((0.286,), (0.353,)),   # (x - mean) / std, per channel
])
 
train_set = datasets.FashionMNIST("data", train=True, download=True, transform=transform)
test_set = datasets.FashionMNIST("data", train=False, download=True, transform=transform)
 
trainloader = DataLoader(train_set, batch_size=128, shuffle=True, num_workers=2)
testloader = DataLoader(test_set, batch_size=256, shuffle=False)
 
classes = train_set.classes
print(classes)          # ['T-shirt/top', 'Trouser', 'Pullover', ...]

Two transforms, two jobs: ToTensor converts the image to a channels-first float tensor scaled to 0–1, and Normalize standardizes it with the dataset's mean and standard deviation — the same "scale your features" habit from tabular data, applied per pixel. Always look at your data before training:

images, labels = next(iter(trainloader))
print(images.shape)     # [128, 1, 28, 28] -> [batch, channels, height, width]
 
fig, axes = plt.subplots(3, 6, figsize=(12, 6))
for img, label, ax in zip(images, labels, axes.flatten()):
    ax.imshow(img.squeeze(), cmap="gray")
    ax.set_title(classes[label])
    ax.axis("off")
plt.show()

The model: conv blocks plus a linear head

A classic small CNN is two parts. The convolutional body extracts features while shrinking the spatial size; the linear head flattens whatever's left and classifies it. The comments track the tensor shape at every stage — do this in every CNN you write, using the size formula from last lesson:

class CNN(nn.Module):
    def __init__(self, n_classes=10):
        super().__init__()
        self.conv = nn.Sequential(
            # in: [1, 28, 28]
            nn.Conv2d(1, 32, kernel_size=3, padding=1),   # -> [32, 28, 28]
            nn.ReLU(),
            nn.MaxPool2d(2),                              # -> [32, 14, 14]
 
            nn.Conv2d(32, 64, kernel_size=3, padding=1),  # -> [64, 14, 14]
            nn.ReLU(),
            nn.MaxPool2d(2),                              # -> [64, 7, 7]
 
            nn.Flatten(),                                 # -> [64 * 7 * 7] = [3136]
        )
        self.fc = nn.Sequential(
            nn.Linear(64 * 7 * 7, 128),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(128, n_classes),                    # raw logits
        )
 
    def forward(self, x):
        return self.fc(self.conv(x))

The bookkeeping matters because nn.Flatten feeds nn.Linear(64 * 7 * 7, ...) — get the arithmetic wrong and you'll meet PyTorch's most famous error, a matrix-shape mismatch on the first forward pass. The pattern to remember: channels grow (1 → 32 → 64) while spatial size shrinks (28 → 14 → 7). The network trades where for what.

Training: same recipe, new data

Nothing about the training loop changes — that's the payoff of the reusable pattern from the DataLoader lesson. Here it is, compact:

def run_epoch(model, loader, criterion, optimizer=None):
    training = optimizer is not None
    model.train() if training else model.eval()
    total_loss, correct, n = 0.0, 0, 0
    with torch.enable_grad() if training else torch.inference_mode():
        for xb, yb in loader:
            xb, yb = xb.to(device), yb.to(device)
            output = model(xb)
            loss = criterion(output, yb)
            if training:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
            total_loss += loss.item() * len(xb)
            correct += (output.argmax(1) == yb).sum().item()
            n += len(xb)
    return total_loss / n, correct / n
 
model = CNN().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
 
for epoch in range(1, 6):
    train_loss, train_acc = run_epoch(model, trainloader, criterion, optimizer)
    test_loss, test_acc = run_epoch(model, testloader, criterion)
    print(f"epoch {epoch} | train acc {train_acc:.3f} | test acc {test_acc:.3f}")
 
torch.save(model.state_dict(), "cnn_baseline.pth")

Five epochs should land around 91–92% test accuracy — already far beyond what an MLP on flattened pixels manages with similar effort. Notice the gap between train and test accuracy creeping open by the last epoch: mild overfitting, our cue for the next section.

Data augmentation: free training data

Data augmentation applies random, label-preserving distortions to each training image, every time it's loaded — flips, small rotations, crops. The model never sees the exact same pixels twice, so memorizing becomes much harder; effectively you've multiplied your dataset. The transforms go in the training transform only — never distort the test set, it's the measuring stick:

train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),       # a mirrored sneaker is still a sneaker
    transforms.RandomRotation(10),           # up to +/-10 degrees
    transforms.ToTensor(),
    transforms.Normalize((0.286,), (0.353,)),
])
 
aug_train_set = datasets.FashionMNIST("data", train=True, transform=train_transform)
aug_trainloader = DataLoader(aug_train_set, batch_size=128, shuffle=True, num_workers=2)
 
torch.manual_seed(0)
model_aug = CNN().to(device)
optimizer = optim.Adam(model_aug.parameters(), lr=0.001)
 
for epoch in range(1, 6):
    train_loss, train_acc = run_epoch(model_aug, aug_trainloader, criterion, optimizer)
    test_loss, test_acc = run_epoch(model_aug, testloader, criterion)
    print(f"epoch {epoch} | train acc {train_acc:.3f} | test acc {test_acc:.3f}")

Compare the two runs. The augmented model's training accuracy is lower (the task got harder — every image is warped), but the train/test gap nearly closes, and with more epochs the augmented model overtakes the baseline on test accuracy. That's the signature of regularization working. Augmentation pays off most when data is scarce or training is long; pick transforms that match reality — horizontal flips make sense for clothing, but would be a terrible idea for digit recognition, where a flipped "3" is no longer a 3.

Read the mistakes

A single accuracy number hides what the model gets wrong. Plotting misclassified images is the fastest way to build intuition:

model_aug.eval()
wrong_imgs, wrong_true, wrong_pred = [], [], []
 
with torch.inference_mode():
    for xb, yb in testloader:
        xb, yb = xb.to(device), yb.to(device)
        preds = model_aug(xb).argmax(1)
        mask = preds != yb
        wrong_imgs.append(xb[mask].cpu())
        wrong_true.append(yb[mask].cpu())
        wrong_pred.append(preds[mask].cpu())
 
wrong_imgs = torch.cat(wrong_imgs)
wrong_true = torch.cat(wrong_true)
wrong_pred = torch.cat(wrong_pred)
print(f"{len(wrong_imgs)} mistakes out of {len(test_set)}")
 
fig, axes = plt.subplots(3, 6, figsize=(14, 7))
for img, t, p, ax in zip(wrong_imgs, wrong_true, wrong_pred, axes.flatten()):
    ax.imshow(img.squeeze(), cmap="gray")
    ax.set_title(f"true: {classes[t]}\npred: {classes[p]}", color="red", fontsize=9)
    ax.axis("off")
plt.tight_layout()
plt.show()

You'll find the errors are systematic, not random: shirts confused with coats, pullovers with shirts — pairs even humans squint at in 28×28 grayscale. That tells you the remaining errors need better inputs or bigger models, not more epochs. And "bigger models" has a shortcut: instead of training a deeper CNN from scratch, borrow one already trained on millions of images.

Check your understanding

5 questions · free

Q1.In the CNN above, why is the first argument of the head's nn.Linear exactly 64 * 7 * 7?
Q2.What is the general shape pattern of a CNN body as data flows through it?
Q3.Why must data augmentation be applied only to the training transform?
Q4.With augmentation on, training accuracy dropped but the train/test gap narrowed. What does this indicate?
Q5.Which augmentation would be a poor choice for handwritten digit recognition?

Exercise: Go deeper: a three-block CNN

Add a third conv block (64 → 128 channels) to the CNN and retrain with the augmented loader for 5 epochs. Before running, compute by hand what spatial size reaches the flatten layer and fix the head's input size to match. Then compare against the two-block model: test accuracy and total parameter count. Which model is bigger — and is the answer what you expected?

Next up: the Transfer Learning module — where we stop training CNNs from scratch and fine-tune networks pretrained on millions of images instead.