Pooling and the Big Picture

Imagine trying to recognise a friend by looking through a tiny keyhole. You'd need to step back, to see more. A convolutional network does just that. Layer by layer, it shrinks the image, so that each small filter sees a bigger and bigger part of the original.

• Looking through a keyhole
• Step back to see more
• Shrink, then look again

Max pooling

Pooling is simple. Split the map into small squares, and from each one, keep only the largest value. It's saying, somewhere in this patch, this pattern was found, but the exact pixel doesn't matter. The map gets smaller, the next filters see more of the image, and the network also gains a little tolerance to small shifts.

A hierarchy of features

Now stack these up. Early layers have small views, so they find edges. The next layers combine edges into textures and shapes. Further on, shapes become parts, like eyes and wheels. By the end, each neuron effectively sees the whole image, and responds to whole objects. Nobody programmed that hierarchy. It emerges by itself, from training. When researchers first saw this, it was a shock.

Following the shapes

Let's trace an image through three blocks, with each block being a convolution, an activation, and a pool. Run it. Watch the shape. Height and width halve at each block, from thirty two to sixteen, to eight, to four. Meanwhile, the channels grow, from three to sixty four. That's the classic pattern. Smaller images, with more and more meaning per pixel.

import torch
from torch import nn

x = torch.rand(1, 3, 32, 32)

layers = nn.Sequential(
    nn.Conv2d(3, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
    nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
    nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
)

print(f"Start:          {list(x.shape)}")
for i, layer in enumerate(layers):
    x = layer(x)
    if isinstance(layer, nn.MaxPool2d):
        print(f"After block {i // 3 + 1}:  {list(x.shape)}")

Start:          [1, 3, 32, 32]
After block 1:  [1, 16, 16, 16]
After block 2:  [1, 32, 8, 8]
After block 3:  [1, 64, 4, 4]

Our first CNN

Now let's train one. Two convolution blocks, then flatten, and one linear layer. Notice that it has only around sixteen thousand weights, far fewer than our earlier flat network. The training loop is unchanged, not a single word. Run it. In three epochs, it's already well above the flat network on the very same fashion images. Fewer weights, better accuracy. That's the power of building the right idea into your model.

import torch
from torch import nn
from torchvision import datasets, transforms

root = "data"
tf = transforms.ToTensor()
train = datasets.FashionMNIST(root, train=True, transform=tf)
test = datasets.FashionMNIST(root, train=False, transform=tf)
X_te = test.data[:4000].float().unsqueeze(1) / 255
y_te = test.targets[:4000]

model = nn.Sequential(
    nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
    nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),
    nn.Flatten(), nn.Linear(32 * 7 * 7, 10),
)
print("Weights:", sum(p.numel() for p in model.parameters()))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
dl = torch.utils.data.DataLoader(train, batch_size=128, shuffle=True)

for epoch in range(1, 4):
    model.train()
    for x, y in dl:
        loss = nn.functional.cross_entropy(model(x), y)
        opt.zero_grad(); loss.backward(); opt.step()
    model.eval()
    with torch.no_grad():
        acc = (model(X_te).argmax(1) == y_te).float().mean()
    print(f"epoch {epoch}  accuracy {acc:.3f}")

Weights: 20490
epoch 1  accuracy 0.846
epoch 2  accuracy 0.882
epoch 3  accuracy 0.881

Summary