A convolutional image classifier

Why convolutions suit images

A Linear layer on raw pixels connects every pixel to every unit: a 224 x 224 x 3 image into 1000 units needs about 150 million weights, and it ignores the fact that nearby pixels belong together. A convolution learns small filters, for example 3 x 3, and slides each one across the whole image. The same weights are reused at every position (weight sharing), so a pattern such as an edge is detected wherever it appears, with far fewer parameters.

A common building block is Conv -> BatchNorm -> ReLU -> Pool. As you go deeper, the number of channels usually grows (more kinds of features) while height and width shrink (each feature summarises a larger region). At the end, nn.AdaptiveAvgPool2d(1) averages each channel over all positions, giving one number per channel regardless of the input size, and a Linear layer maps those to class logits. So the same network accepts 32 x 32 or 64 x 64 images unchanged.

Write the expected shape next to each layer as a comment. Most CNN errors are shape errors.

Inside a small CNN

An image batch passes through stacked convolution blocks, which produce feature maps with more channels and smaller spatial size. Global average pooling turns each map into one number, a linear head maps those features to class logits.

SmallCNN

Three blocks take a 32 x 32 image down to 4 x 4 with 128 channels; global average pooling and a Linear layer finish the job. This class is reused in later topics.

import torch
from torch import nn

def conv_block(c_in, c_out):
    return nn.Sequential(
        nn.Conv2d(c_in, c_out, kernel_size=3, padding=1, bias=False),
        nn.BatchNorm2d(c_out),
        nn.ReLU(inplace=True),
        nn.MaxPool2d(2),               # halves height and width
    )

class SmallCNN(nn.Module):
    def __init__(self, num_classes: int = 10):
        super().__init__()
        self.features = nn.Sequential(
            conv_block(3, 32),         # (N,3,32,32)  -> (N,32,16,16)
            conv_block(32, 64),        # -> (N,64,8,8)
            conv_block(64, 128),       # -> (N,128,4,4)
        )
        self.head = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),   # -> (N,128,1,1)
            nn.Flatten(),              # -> (N,128)
            nn.Linear(128, num_classes),
        )

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

model = SmallCNN()
print(model(torch.randn(4, 3, 32, 32)).shape)   # torch.Size([4, 10])

Small details that matter

bias=False on a conv followed by BatchNorm is deliberate: BatchNorm has its own learned shift, so the conv bias would be redundant. And inplace=True on ReLU saves memory, but never use in-place operations on a tensor that autograd still needs; when unsure, leave it off.

Quick check: What is the main benefit of ending a CNN with `nn.AdaptiveAvgPool2d(1)` before the Linear head?

• It removes the need for a loss function
• It makes the head independent of the input image size
• It doubles the number of channels
• It replaces BatchNorm