🎓 A1Learner
🔍

PyTorch Intermediate Capstone

The capstone is one full training pipeline for image classification. It combines synthetic data, augmentation, a schedule, early stopping, a saved best checkpoint, and a metrics report.

One ring, drawn with distances

A ring is the set of cells whose distance from the centre sits in a band. The band from 4.5 to 6.5 lights 64 cells.

import torch

rows = torch.arange(16).view(16, 1).float()
cols = torch.arange(16).view(1, 16).float()
dist = torch.sqrt((rows - 7.5) ** 2 + (cols - 7.5) ** 2)
ring = ((dist >= 4.5) & (dist <= 6.5)).float()
print(ring.shape, int(ring.sum()))

A one-pixel shift changes the image

Augmentation shifts a training image by one pixel and clears the wrapped edge. The shifted square is a new training example with the same label.

img = torch.zeros(1, 16, 16)
img[0, 3:11, 3:11] = 1.0
shifted = torch.roll(img, shifts=(0, 1), dims=(1, 2))
shifted[:, :, 0] = 0
print(not torch.equal(img, shifted))

The CNN shape flow

The capstone CNN takes batches of 1 by 16 by 16 images. It returns one score per class for each image.

model = ShapesCNN()
out = model(torch.zeros(2, 1, 16, 16))
print(out.shape)

The schedule, epoch by epoch

StepLR holds the rate at 0.01 for five epochs. It then halves the rate to 0.005.

rates = []
for epoch in range(7):
    rates.append(round(optimizer.param_groups[0]["lr"], 4))
    scheduler.step()
print(rates)

Example

import copy

import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(7)


def make_shapes(count, seed):
    g = torch.Generator().manual_seed(seed)
    images = torch.zeros(count, 1, 16, 16)
    labels = torch.randint(0, 3, (count,), generator=g)
    rows = torch.arange(16).view(16, 1).float()
    cols = torch.arange(16).view(1, 16).float()
    for i in range(count):
        cls = labels[i].item()
        if cls == 0:
            top = int(torch.randint(5, 10, (1,), generator=g))
            images[i, 0, top:top + 3, :] = 1.0
        elif cls == 1:
            cy = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
            cx = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
            dist = torch.sqrt((rows - cy) ** 2 + (cols - cx) ** 2)
            images[i, 0] = ((dist >= 4.5) & (dist <= 6.5)).float()
        else:
            top = int(torch.randint(3, 6, (1,), generator=g))
            left = int(torch.randint(3, 6, (1,), generator=g))
            images[i, 0, top:top + 8, left:left + 8] = 1.0
    images += 0.1 * torch.rand(images.shape, generator=g)
    return images.clamp(0, 1), labels


def augment(images, g):
    out = images.clone()
    flip = torch.rand(len(out), generator=g) < 0.5
    out[flip] = torch.flip(out[flip], dims=[3])
    dys = torch.randint(-1, 2, (len(out),), generator=g)
    dxs = torch.randint(-1, 2, (len(out),), generator=g)
    for dy in (-1, 0, 1):
        for dx in (-1, 0, 1):
            if dy == 0 and dx == 0:
                continue
            mask = (dys == dy) & (dxs == dx)
            if mask.any():
                shifted = torch.roll(out[mask], shifts=(dy, dx), dims=(2, 3))
                if dy == 1:
                    shifted[:, :, 0, :] = 0
                if dy == -1:
                    shifted[:, :, -1, :] = 0
                if dx == 1:
                    shifted[:, :, :, 0] = 0
                if dx == -1:
                    shifted[:, :, :, -1] = 0
                out[mask] = shifted
    return out


class ShapesCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 8, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(8, 16, 3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.head = nn.Linear(16, 3)

    def forward(self, x):
        x = self.features(x).flatten(1)
        return self.head(x)


def evaluate(model, loader):
    model.eval()
    correct = 0
    total = 0
    hits = [0, 0, 0]
    totals = [0, 0, 0]
    with torch.no_grad():
        for images, labels in loader:
            preds = model(images).argmax(dim=1)
            correct += (preds == labels).sum().item()
            total += len(labels)
            for c in range(3):
                mask = labels == c
                totals[c] += mask.sum().item()
                hits[c] += (preds[mask] == c).sum().item()
    recalls = [hits[c] / totals[c] for c in range(3)]
    return correct / total, recalls


train_images, train_labels = make_shapes(240, 11)
val_images, val_labels = make_shapes(90, 23)
train_loader = DataLoader(
    TensorDataset(train_images, train_labels),
    batch_size=32, shuffle=True, generator=torch.Generator().manual_seed(3),
)
val_loader = DataLoader(TensorDataset(val_images, val_labels), batch_size=32)

model = ShapesCNN()
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.5)
aug_g = torch.Generator().manual_seed(5)

best_loss = float("inf")
best_state = None
wait = 0
last_epoch = 0
for epoch in range(1, 16):
    last_epoch = epoch
    model.train()
    for images, labels in train_loader:
        images = augment(images, aug_g)
        optimizer.zero_grad()
        loss = loss_fn(model(images), labels)
        loss.backward()
        optimizer.step()
    scheduler.step()
    model.eval()
    val_total = 0.0
    seen = 0
    with torch.no_grad():
        for images, labels in val_loader:
            val_total += loss_fn(model(images), labels).item() * len(labels)
            seen += len(labels)
    val_loss = val_total / seen
    print(f"Epoch {epoch}: val loss {val_loss:.4f}")
    if val_loss < best_loss - 0.01:
        best_loss = val_loss
        best_state = copy.deepcopy(model.state_dict())
        wait = 0
    else:
        wait += 1
        if wait >= 2:
            print(f"Early stopping at epoch {epoch}")
            break

torch.save(best_state, "best_shapes.pt")
model.load_state_dict(torch.load("best_shapes.pt"))
accuracy, recalls = evaluate(model, val_loader)
names = ["bar", "ring", "square"]
print(f"Epochs run: {last_epoch}")
print(f"Best val loss: {best_loss:.4f}")
print(f"Validation accuracy: {accuracy:.3f}")
for c in range(3):
    print(f"Recall for {names[c]}: {recalls[c]:.3f}")

Key points

  • Shapes data: 16 by 16 single-channel images of a bar, a ring, or a square
  • Augmentation: random flips and one-pixel shifts applied to training batches
  • Schedule: StepLR halves the learning rate every 5 epochs
  • Early stopping: training stops when validation loss stops improving by at least 0.01
  • Best checkpoint: the state_dict with the lowest validation loss, saved to a file
  • Recall: the share of one class that the model gets right
  • In short, the pipeline trains with help, stops on evidence, and reports per class

In this example, make_shapes builds 240 training images and 90 validation images with fixed seeds. Bars, rings, and squares are drawn with masks and distances, and a little noise is added. The ShapesCNN stacks two convolution blocks and ends with a linear head. Each training batch passes through augment, which flips some images and shifts others by one pixel. Adam trains the model, and StepLR halves the rate every five epochs. After each epoch, the loop scores the validation slice and keeps a copy of the best state_dict. Gains under 0.01 count as no improvement, so two flat epochs stop the run at epoch 13. The best state is saved to a file and loaded back before scoring. The report prints accuracy and one recall line per class. PyTorch Advanced assumes this whole pipeline as a starting point. It assumes a model that trains with a schedule and stops on validation evidence. The model saves its best state and reports per-class results.

Note

The report scores the validation slice. A separate test slice would be touched only once, at the very end.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why does the capstone load the saved best state before reporting metrics?