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?