TopolearnCNNs remain the workhorse for vision tasks on modest compute, and they teach how architectures are composed.
A convolutional network stacks blocks of Conv2d -> (BatchNorm2d) -> ReLU -> pooling. Early layers detect edges and textures, later ones larger patterns, while pooling shrinks the spatial size and the channel count grows. A small head (flatten or global pooling, then Linear) produces class logits.
The output size of a conv or pool along one spatial dimension is:
floor((H + 2*padding - kernel) / stride) + 1
so 3x3 kernels with padding=1 keep the size, and MaxPool2d(2) halves it. Instead of computing the flattened size by hand, end with nn.AdaptiveAvgPool2d(1) (global average pooling), which works for any input resolution, or use nn.LazyLinear.
The torchvision package supplies datasets (MNIST, CIFAR10, ImageFolder), image transforms and pretrained models. Its torchvision.transforms.v2 API converts images to tensors, normalises them and applies augmentation such as random crops and flips, which is the cheapest way to reduce overfitting.
Use BatchNorm2d after convolutions in deeper nets (then set bias=False on the conv), and remember that training images need augmentation but validation images only need resizing and normalising.
Example
import torch
from torch import nn
class SmallCNN(nn.Module):
def __init__(self, n_classes=10):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 16, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 16 x 14 x 14
nn.Conv2d(16, 32, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2), # 32 x 7 x 7
)
self.head = nn.Sequential(nn.Flatten(), nn.Linear(32 * 7 * 7, n_classes))
def forward(self, x):
return self.head(self.features(x))
model = SmallCNN()
x = torch.randn(8, 1, 28, 28) # a batch of MNIST-sized images
print(model.features(x).shape) # torch.Size([8, 32, 7, 7])
print(model(x).shape) # torch.Size([8, 10])Remember
- Conv-ReLU-pool blocks shrink space and grow channels
- Size = floor((H + 2p - k) / s) + 1
- Global average pooling or LazyLinear avoid hand-computed sizes
Watch out for
- Applying training augmentations to the validation set, which makes metrics noisy and pessimistic.