🎓 A1Learner
🔍

PyTorch Training a CNN

Training a CNN is the standard training loop run on a convolutional network. The loop keeps the full discipline of batches, a learning rate schedule, and a validation score after every epoch.

One forward pass

The model takes a batch of images. It returns one score row per image, with one score per class.

model = ShapesCNN()
x = torch.zeros(4, 1, 16, 16)
print(model(x).shape)

Shapes through the stack

The first conv keeps the 16 by 16 size because padding is 1. The pool then halves each side to 8.

h = torch.relu(model.conv1(torch.zeros(2, 1, 16, 16)))
print(h.shape)
h = model.pool(h)
print(h.shape)

GAP ends the stack

The second conv makes 16 feature maps of 8 by 8. Global average pooling averages each map to a single number.

h = torch.relu(model.conv2(h))
h = model.gap(h)
print(h.shape)

The schedule halves the rate

StepLR with step_size 4 and gamma 0.5 halves the rate after four scheduler steps. The print shows the rate after those steps.

optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.5)
for _ in range(4):
    scheduler.step()
print(optimizer.param_groups[0]["lr"])

Example

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


def draw_shape(label, g):
    rows = torch.arange(16).view(16, 1).float()
    cols = torch.arange(16).view(1, 16).float()
    if label == 0:
        centre = 8 + int(torch.randint(-2, 3, (1,), generator=g))
        mask = (rows >= centre - 1) & (rows <= centre + 1) & (cols >= 0)
    elif label == 1:
        cr = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
        cc = 7.5 + int(torch.randint(-1, 2, (1,), generator=g))
        dist = torch.sqrt((rows - cr) ** 2 + (cols - cc) ** 2)
        mask = (dist >= 4.5) & (dist <= 6.5)
    else:
        top = int(torch.randint(2, 7, (1,), generator=g))
        left = int(torch.randint(2, 7, (1,), generator=g))
        mask = (rows >= top) & (rows <= top + 7) & (cols >= left) & (cols <= left + 7)
    return mask.float()


def make_shapes(count, seed):
    g = torch.Generator().manual_seed(seed)
    labels = torch.arange(count) % 3
    labels = labels[torch.randperm(count, generator=g)]
    images = [draw_shape(int(label), g) for label in labels]
    return torch.stack(images).unsqueeze(1), labels


class ShapesCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 8, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(8, 16, 3, padding=1)
        self.gap = nn.AdaptiveAvgPool2d(1)
        self.head = nn.Linear(16, 3)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = torch.relu(self.conv2(x))
        x = self.gap(x)
        return self.head(torch.flatten(x, 1))


train_x, train_y = make_shapes(240, seed=42)
val_x, val_y = make_shapes(60, seed=43)
loader = DataLoader(TensorDataset(train_x, train_y), batch_size=32,
                    shuffle=True, generator=torch.Generator().manual_seed(0))

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


def accuracy():
    model.eval()
    with torch.no_grad():
        predictions = model(val_x).argmax(dim=1)
    return float((predictions == val_y).float().mean())


for epoch in range(1, 9):
    model.train()
    total_loss = 0.0
    for xb, yb in loader:
        optimizer.zero_grad()
        loss = loss_fn(model(xb), yb)
        loss.backward()
        optimizer.step()
        total_loss += float(loss) * len(xb)
    scheduler.step()
    print(f"Epoch {epoch}: train loss {total_loss / len(train_x):.4f}, "
          f"validation accuracy {accuracy():.3f}")

print(f"Final validation accuracy: {accuracy():.3f}")

Key points

  • Conv stack: Conv2d, ReLU, MaxPool2d, Conv2d, ReLU, then global average pooling
  • Global average pooling: each feature map becomes one number before the Linear head
  • Use CrossEntropyLoss on the raw scores from the Linear head
  • Use StepLR to halve the learning rate halfway through the run
  • Validation accuracy: the fraction of held-back images the model labels right

In this example, make_shapes builds 240 train images and 60 validation images with fixed seeds. The DataLoader serves batches of 32 and shuffles with its own seeded generator. The model stacks two conv layers around a pool and ends with global average pooling and a Linear head. Adam starts at a learning rate of 0.01. StepLR halves that rate after epoch 4. Each epoch prints the train loss and the validation accuracy side by side. The loss falls from 1.1037 to 0.1008. Validation accuracy reaches 1.000 at epoch 3 and stays there.

Note

The model trains on the train split only, and validation images are scored but never trained on.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does global average pooling hand to the Linear head?