Convolution arithmetic and a CNN in PyTorch
By the end: You can compute output shapes and parameter counts by hand, and build and train a small CNN with batch norm and residual connections.
- Output size of a convolution
- Parameters in a layer
- Receptive field
- A small CNN for CIFAR-10
- Batch norm and residual connections
New to convolutional networks? A convolutional network is a neural network built for images. Rather than relating every pixel to every other pixel, it slides a small window across the picture hunting for one local pattern — an edge, then a corner, then a wheel — and hunts for that same pattern everywhere in the frame. Stack those layers and edges become shapes and shapes become objects. Why it matters
A convolution slides a small window of weights across an image, looking for the same pattern everywhere in it, which is why networks built from them work on pictures. This section covers the practical half of that: what size comes out of a layer, how many parameters it holds, and how much of the image a single output unit can see. Then it builds and trains a small network with the two additions that make deep ones trainable at all, batch normalisation and residual connections.
The arithmetic#
output size = floor((input - kernel + 2 x padding) / stride) + 1
32x32 input, 3x3 kernel, padding 1, stride 1 -> 32x32 (padding 1 keeps the size)
32x32 input, 3x3 kernel, padding 1, stride 2 -> 16x16
parameters = kernel x kernel x in_channels x out_channels + out_channels (bias)
3x3 conv, 64 -> 128 channels: 3 x 3 x 64 x 128 + 128 = 73,856Three quantities go in and two come out, and it is worth being deliberate about each of them, because shape confusion is where most people stall. The kernel size is how wide the window is. The stride is how far the window jumps between positions, so stride 2 halves the output. The padding is how many rows and columns of zeros are added round the edge before sliding starts, and the reason it exists is that without it every layer shrinks the image a little, which is why padding of half the kernel size — 1 for a 3x3, 2 for a 5x5, 3 for a 7x7 — is the near-universal choice: it keeps the size exactly the same.
Channels are the part people get wrong, because they behave nothing like the spatial dimensions. A convolution does not slide over channels. It reads all of them at once: a 3x3 kernel applied to a 64-channel input is really a 3 x 3 x 64 block of weights, and one such block produces one output channel. The number of output channels is simply how many different blocks you asked for, and it is a free choice — it has no relationship to the input channel count.
import torch
import torch.nn as nn
x = torch.zeros(1, 3, 32, 32) # batch, channels, height, width
for k, s, p in [(3, 1, 1), (3, 1, 0), (3, 2, 1), (5, 1, 2), (7, 2, 3), (1, 1, 0)]:
c = nn.Conv2d(3, 64, k, stride=s, padding=p)
n = sum(q.numel() for q in c.parameters())
print(f"k={k} s={s} p={p} -> {tuple(c(x).shape)} params={n}")
# k=3 s=1 p=1 -> (1, 64, 32, 32) params=1792
# k=3 s=1 p=0 -> (1, 64, 30, 30) params=1792
# k=3 s=2 p=1 -> (1, 64, 16, 16) params=1792
# k=5 s=1 p=2 -> (1, 64, 32, 32) params=4864
# k=7 s=2 p=3 -> (1, 64, 16, 16) params=9472
# k=1 s=1 p=0 -> (1, 64, 32, 32) params=256When the shapes are wrong PyTorch usually tells you, but the message names the tensors rather than the mistake, so it is worth recognising the two you will actually hit.
try: # channels-last data, as image libraries produce
nn.Conv2d(3, 8, 3)(torch.zeros(1, 32, 32, 3))
except RuntimeError as e:
print("A:", e)
try: # forgot the spatial size after flattening
bad = nn.Sequential(nn.Conv2d(3, 8, 3), nn.Flatten(), nn.Linear(8 * 30 * 30, 10))
bad(torch.zeros(1, 3, 64, 64))
except RuntimeError as e:
print("B:", e)
# A: Given groups=1, weight of size [8, 3, 3, 3], expected input[1, 32, 32, 3]
# to have 3 channels, but got 32 channels instead
# B: mat1 and mat2 shapes cannot be multiplied (1x30752 and 7200x10)A missing batch dimension does not always error
Conv2d accepts an unbatched three-dimensional input and treats it as channels, height, width. So passing a single 3x32x32 image without a batch dimension runs happily and returns an unbatched result, and the mistake only surfaces later when a loss function disagrees about shapes. Check tuple(x.shape) rather than assuming an error would have told you.
Receptive field#
The receptive field is how much of the original image one output unit can see. It matters because a network whose final units see a 30-pixel window cannot recognise anything larger than 30 pixels, however many layers it has. You do not have to trust the formula: put a gradient on one output unit and see which input pixels receive any.
def receptive_field(net, size=41):
inp = torch.zeros(1, 1, size, size, requires_grad=True)
out = net(inp)
mid = out.shape[-1] // 2
out[0, 0, mid, mid].backward()
touched = (inp.grad[0, 0].abs() > 0).nonzero()
return int(touched[:, 0].max() - touched[:, 0].min() + 1)
def ones_conv():
c = nn.Conv2d(1, 1, 3, padding=1, bias=False)
nn.init.constant_(c.weight, 1.0) # so nothing cancels to zero by accident
return c
for depth in (1, 2, 3, 5, 10):
print(depth, receptive_field(nn.Sequential(*[ones_conv() for _ in range(depth)])))
print("with a pool:", receptive_field(nn.Sequential(ones_conv(), nn.MaxPool2d(2), ones_conv())))
# 1 3
# 2 5
# 3 7
# 5 11
# 10 21
# with a pool: 7Two practical consequences. If your objects are large relative to the image, you need downsampling and not merely depth. If your objects are small, aggressive downsampling destroys them before the network can see them, which is why detectors for small objects use higher input resolution or predict from earlier, higher-resolution feature maps.
A small network#
import torch.nn as nn
def block(c_in, c_out):
return nn.Sequential(
nn.Conv2d(c_in, c_out, 3, padding=1, bias=False),
nn.BatchNorm2d(c_out),
nn.ReLU(),
)
class SmallCNN(nn.Module):
def __init__(self, classes=10):
super().__init__()
self.net = nn.Sequential(
block(3, 64), block(64, 64), nn.MaxPool2d(2), # 32 -> 16
block(64, 128), block(128, 128), nn.MaxPool2d(2), # 16 -> 8
block(128, 256), nn.AdaptiveAvgPool2d(1), # 8 -> 1
nn.Flatten(), nn.Linear(256, classes),
)
def forward(self, x): # [B, 3, 32, 32]
return self.net(x)Count its parameters by hand before you run it, because this is the arithmetic an interviewer asks for out loud. Each conv is 3 x 3 x in x out with no bias, because the batch norm that follows has its own shift; each batch norm is two numbers per channel.
conv 3 -> 64 3 x 3 x 3 x 64 = 1,728 + 128 batchnorm
conv 64 -> 64 3 x 3 x 64 x 64 = 36,864 + 128
conv 64 -> 128 3 x 3 x 64 x 128 = 73,728 + 256
conv 128 -> 128 3 x 3 x 128 x 128 = 147,456 + 256
conv 128 -> 256 3 x 3 x 128 x 256 = 294,912 + 512
linear 256 -> 10 2,560 + 10
558,538 total
sum(p.numel() for p in SmallCNN().parameters()) -> 558538Batch norm#
Batch normalisation rescales each channel's activations to a steady mean and spread, which makes deeper networks train faster and much less sensitively to the learning rate. During training it uses the mean and variance of the current batch; during evaluation it uses a running average accumulated over training. That difference is not a detail, and model.eval() matters here more than anywhere else in PyTorch, because in training mode a prediction for one image depends on which other images happened to be in the batch with it.
torch.manual_seed(0)
net = nn.Sequential(nn.Conv2d(3, 8, 3, padding=1), nn.BatchNorm2d(8), nn.ReLU(),
nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(8, 4))
x, y = torch.randn(64, 3, 16, 16), torch.randint(0, 4, (64,))
opt = torch.optim.SGD(net.parameters(), lr=0.1)
for _ in range(50): # enough for the running stats to settle
opt.zero_grad()
nn.functional.cross_entropy(net(x), y).backward()
opt.step()
sample = torch.randn(1, 3, 16, 16) # the image we care about
crowd_a = torch.randn(7, 3, 16, 16) # seven ordinary neighbours
crowd_b = torch.randn(7, 3, 16, 16) * 4 # seven with a very different spread
for mode in ("train", "eval"):
getattr(net, mode)()
with torch.no_grad():
a = net(torch.cat([sample, crowd_a]))[0]
b = net(torch.cat([sample, crowd_b]))[0]
diff = round(float((a - b).abs().max()), 4)
print(mode, "max difference for the same image:", diff)
# train max difference for the same image: 0.1517
# eval max difference for the same image: 0.0Residual connections#
class Residual(nn.Module):
def __init__(self, c):
super().__init__()
self.f = nn.Sequential(block(c, c), nn.Conv2d(c, c, 3, padding=1, bias=False), nn.BatchNorm2d(c))
def forward(self, x):
return nn.functional.relu(x + self.f(x)) # add the input backAdding the input back means each block only has to learn a correction, and gradients have a direct path to early layers. This is what made networks with a hundred or more layers trainable, and the same idea runs through every transformer. There is one detail the code above hides: the addition requires the two tensors to have the same shape, so whenever a block changes the channel count or the stride, the shortcut needs its own 1x1 convolution to match. That is exactly what the downsample branch is in torchvision's ResNet implementation, and it is the first thing to look at when a residual network throws a shape error.
What to skip. You do not need to memorise the list of normalisation variants, or the difference between the original and pre-activation residual orderings, or the schedule that any particular paper used. What you do need, cold, is the output-size formula, the parameter formula, the fact that batch norm behaves differently in the two modes, and why an addition keeps gradients alive.
Where people go wrong
Forgetting model.eval() when evaluating a network with batch norm. Validation uses the statistics of each evaluation batch and scores become noisy and wrong.
If this comes up in an interview
What is the output shape of a 3x3 convolution with stride 2 and padding 1 on a 64x64x32 input with 64 filters, and how many parameters does it have?
Spatially floor((64 - 3 + 2) / 2) + 1 = 32, so 32x32x64. Parameters: 3 x 3 x 32 x 64 = 18,432 weights plus 64 biases, 18,496.
Then they push further
Why do residual connections help very deep networks train?
Each block learns a correction added to its input, so an unhelpful block can stay close to the identity, and gradients flow directly through the additions to early layers without shrinking at every layer.
What gets marked up. Marked up for getting the arithmetic right out loud. Interviews for vision roles often include exactly this.
Build this: Train the small CNN on CIFAR-10 with and without augmentation, then add residual blocks and compare. Compute the parameter count by hand and check it against sum(p.numel() for p in model.parameters()).
Watch it explained
This is one section of Convolutional networks, 21 hours across 5 sections. Read it in full to get the whole topic in order.