Unit 06.02: Convolution, pooling, flattening, and classifier heads

A CNN is convolutions to extract features, pooling to reduce size, then a flatten and a linear layer to classify. The join between the two halves is where the bugs live.

Tracing the shape through every layer

Convolutions preserve spatial layout; pooling shrinks it; flatten collapses whatever remains into a vector for the classifier.

The number the linear layer needs — 16 × 7 × 7 here — depends on every preceding layer. Get it wrong and you get a matrix-size error at the very end.

Run the input through one layer at a time and read the shapes:

import torch
from torch import nn

torch.manual_seed(0)
model = nn.Sequential(
    nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),   # 28 -> 14
    nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),  # 14 -> 7
    nn.Flatten(),
    nn.Linear(16 * 7 * 7, 10),
)

x = torch.randn(4, 1, 28, 28)
# Trace the shape through every stage -- this is how you debug a size error.
h = x
for layer in model:
    h = layer(h)
    print(f"{layer.__class__.__name__:12} -> {tuple(h.shape)}")

print("\nflattened features:", 16 * 7 * 7)
print("total parameters  :", sum(p.numel() for p in model.parameters()))

# Getting 16*7*7 wrong is the most common CNN error. Compute it, do not guess:
with torch.no_grad():
    features = nn.Sequential(*list(model)[:-2])(x)
print("verified flatten size:", features.flatten(1).shape[1])

Two poolings take 28 → 14 → 7, and the final convolution leaves 16 channels, giving 16 × 7 × 7 = 784 flattened features.

The last block shows how to *verify* rather than calculate: run a batch through everything except the head and read the flattened width directly.

The mistake this prevents

Computing the flatten size by hand and getting it wrong, then changing the number until the error goes away. The model runs and the architecture is not what you intended.

Takeaway

Let the code tell you the flatten size. Do not derive it on paper when you can measure it.

Mark complete