Pooling, BatchNorm and the Conv Block

Context

Max pooling keeps the largest value of every small block of a map and so halves the map with no weights at all. Batch normalization (BatchNorm) rescales each channel of a batch to a mean of 0 and a standard deviation of 1, then by two numbers it learns, and keeps running averages of what it saw for the model to use once training ends. A conv block is a convolution, a BatchNorm, a ReLU and a pooling layer, repeated.

The conv block is the unit the network of Lesson 5.1.3 repeats, and torchvision's pretrained ResNets, a widely used family of vision models, print a BatchNorm2d after each of their convolutions. Every time such a model is scored, served one image at a time or saved with torch.save, what BatchNorm does in training and what it does in evaluation decides the number that comes out.

Lesson 5.1.1 halved its maps with a stride and stacked layers with a ReLU between them, and stopped short of what a working stack needs. Eight seeded layers with nothing else between them shrink a batch's values to a few hundredths. A model left in training mode answers for one image differently depending on the images beside it in the batch, with no error. A bias placed before a BatchNorm is a parameter that trains and changes nothing. And a map of odd size, pooled, loses its last row as quietly as a stride did.

Goal

• Pool a map with nn.MaxPool2d(2) and nn.AvgPool2d, predict the output side, and show from a run how much of a one-pixel move pooling absorbs.
• Compute a BatchNorm2d output by hand from a batch's per-channel mean and variance, and name its two parameters and its three buffers.
• Show the same image answered differently in two batches in training mode and identically in evaluation mode, and say which of Lesson 4.5.2's functions already switches modes.
• Measure how the spread of a batch's values changes through eight convolutions with and without BatchNorm.
• Build conv_block in blocks.py, trace a batch through two blocks, and prove that a bias before BatchNorm is wasted and that ReLU and max pooling may swap places.

Prerequisites

None of this is required, and the lesson stands on its own. Both are here if you want pooling drawn on a small grid first, or BatchNorm's exact contract before the lesson takes it apart.

• YouTube - "C4W1L09 Pooling Layers" by DeepLearningAI, Andrew Ng's course lecture. Max pooling with a 2 by 2 window and a stride of 2 on a small grid, its two settings and the absence of anything to learn, average pooling, and pooling applied to each channel on its own.
• Official docs - PyTorch 2.14, torch.nn.BatchNorm2d. The formula with its eps, the mean and variance taken per channel over the mini-batch, the learned scale and shift, the running estimates kept with momentum=0.1, and which statistics evaluation mode uses.

Table of Contents

• Max Pooling
• Batch Normalization, Training and Evaluation
• The Conv Block

Theory

Every Theory block runs in ~/dev/bootcamp, beside fashion.py, Lesson 4.5.1's, and image_loop.py, Lesson 4.5.2's. The last subtitle saves blocks.py, which Lesson 5.1.3 imports. Nothing here trains: every layer runs on seeded weights, and the statistics BatchNorm keeps come from batches it is shown, never from a fit. A seeded draw prints what it prints on torch 2.14.1 on a processor. A Python block that opens with its imports is a script of its own, and a block that does not continues the script above it.

Max Pooling

What. nn.MaxPool2d(2) cuts each map into blocks of 2 by 2 and writes the largest of the four values, so a 28 by 28 map becomes 14 by 14. Its stride defaults to its kernel size, so the blocks do not overlap, the output side follows Lesson 5.1.1's formula with k = s = 2, and every channel is pooled on its own, so the channel count stays. nn.AvgPool2d(2) writes the mean of each block instead. Neither holds a parameter. Why. Halving a map cuts the work of every later layer by four and grows their receptive fields as a stride does, at no cost in weights. The maximum keeps the strongest response of each block, so a feature that moves inside its block keeps its value, a small amount of tolerance for movement that a stride of 2 does not give. When. After the ReLU of a conv block, in the networks of this section and in most classic image networks.

import torch
from torch import nn
from fashion import train_full

toy = torch.tensor([[1.0, 5.0, 2.0, 0.0], [3.0, 4.0, 8.0, 1.0], [0.0, 2.0, 6.0, 7.0], [9.0, 1.0, 3.0, 2.0]]).view(1, 1, 4, 4)
pool = nn.MaxPool2d(2)
print(pool, list(pool.parameters()))
print(pool(toy))
print(nn.AvgPool2d(2)(toy))

MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False) []
tensor([[[[5., 8.],
          [9., 7.]]]])
tensor([[[[3.2500, 2.7500],
          [3.0000, 4.5000]]]])

The top left block holds 1, 5, 3 and 4, and max pooling keeps 5 where average pooling keeps their mean, 3.25. The printout names every setting, the stride equal to the kernel among them, and ceil_mode=False is the floor of Lesson 5.1.1's formula, kept. The parameter list is empty, so nothing about the layer trains. What it hands back to the layers below it in training, the gradient of the pooled values' sum with respect to the grid, Lesson 4.1.4's backward():

grid = toy.clone().requires_grad_(True)
pool(grid).sum().backward()
print(grid.grad)
spread = toy.clone().requires_grad_(True)
nn.AvgPool2d(2)(spread).sum().backward()
print(spread.grad)

tensor([[[[0., 1., 0., 0.],
          [0., 0., 1., 0.],
          [0., 0., 0., 1.],
          [1., 0., 0., 0.]]]])
tensor([[[[0.2500, 0.2500, 0.2500, 0.2500],
          [0.2500, 0.2500, 0.2500, 0.2500],
          [0.2500, 0.2500, 0.2500, 0.2500],
          [0.2500, 0.2500, 0.2500, 0.2500]]]])

Max pooling sends the whole gradient of each block to the one value that won it, the 5, the 8, the 7 and the 9, and nothing to the other three, because nudging a value that lost moves no output. Average pooling shares each block's gradient equally, a quarter to each of the four. So in training, through max pooling, the filter below learns from the positions where it responded most and from no others, one value in four on every step. Sizes:

for side in (28, 14, 7, 15):
    print(side, pool(torch.zeros(1, 8, side, side)).shape)

28 torch.Size([1, 8, 14, 14])
14 torch.Size([1, 8, 7, 7])
7 torch.Size([1, 8, 3, 3])
15 torch.Size([1, 8, 7, 7])

Eight channels in, eight out, at every size. 28 halves to 14 and 14 to 7, and 7 gives 3, (7 - 2) // 2 + 1: three blocks fit, and the seventh row and column are dropped exactly as Lesson 5.1.1's 2 by 2 stride dropped them. One lit pixel, moved along its row:

dot = torch.zeros(1, 1, 4, 4)
dot[0, 0, 0, 0] = 1.0
for columns in range(4):
    moved = torch.roll(dot, shifts=columns, dims=3)
    print(columns, pool(moved)[0, 0].flatten().tolist())

0 [1.0, 0.0, 0.0, 0.0]
1 [1.0, 0.0, 0.0, 0.0]
2 [0.0, 1.0, 0.0, 0.0]
3 [0.0, 1.0, 0.0, 0.0]

Moved by one column the pixel stays inside its block, and the pooled map is unchanged. Moved by two it crosses into the next block, and the pooled map moves by one position, the move Lesson 4.5.3's shift measured, at half the scale. Whether a move is absorbed depends on where the feature sits in its block, so on real maps it is a share:

image = train_full[0][0].unsqueeze(0)
torch.manual_seed(42)
conv = nn.Conv2d(1, 8, kernel_size=3, padding=1)
with torch.no_grad():
    maps = torch.relu(conv(image))
    moved = torch.relu(conv(torch.roll(image, shifts=1, dims=3)))
print(f"{(maps == moved).float().mean().item():.4f}")
print(f"{(pool(maps) == pool(moved)).float().mean().item():.4f}")

0.5408
0.6084

Eight seeded filters on the ankle boot and on the boot moved one pixel right: 0.5408 of the 6,272 values in the eight 28 by 28 maps are the same, most of them zeros the ReLU writes on both, and after pooling 0.6084 of the 1,568 values are. Pooling absorbs part of a one-pixel move and leaves most of what changed still changed.

DON'T:

• DON'T read max pooling as making a network ignore a move. One pixel left 0.6084 of the pooled values unchanged, and two pixels move the pooled map by a whole position.
• DON'T expect every value of a pooled block to learn on a step. Max pooling's gradient reached one value of each 2 by 2 block, 1.0 there and 0.0 at the other three.
• DON'T pool a map of odd side and expect to keep its edge. A 7 by 7 map pools to 3 by 3, and its last row and column are never read.

Batch Normalization, Training and Evaluation

What. nn.BatchNorm2d(C) works on each of C channels alone. In training mode it takes the mean and the variance of a channel over every image of the batch and every position, subtracts the mean, divides by the square root of the variance plus eps, a small constant that keeps the division safe, then multiplies by a learned weight and adds a learned bias, 2 * C parameters that start at 1 and 0. Each training batch also updates its running mean and running variance, buffers (tensors a module keeps and saves that are not parameters and take no gradient), by momentum=0.1: the new running value is 0.9 of the old one plus 0.1 of the batch's. In evaluation mode it normalizes with the running values instead of the batch's. Why. The values a convolution writes take whatever scale its weights give them, and each layer hands the next a scale that changes as the weights below it train. BatchNorm puts every channel back to a mean of 0 and a standard deviation of 1 on every batch, then lets the network choose its own scale through the two learned numbers. Lesson 4.5.1's Normalize standardized the pixels once, with the training set's statistics. BatchNorm does the same after every layer, with statistics it gathers itself. When. After a convolution and before its ReLU, in every conv block of this section, and every time a model holding one is evaluated, which needs model.eval() first.

import torch
from torch import nn
from fashion import train_loader, train_set, val_set

bn = nn.BatchNorm2d(8)
print(bn)
for name, p in bn.named_parameters():
    print(name, tuple(p.shape), p.detach()[:3])
for name, b in bn.named_buffers():
    print(name, tuple(b.shape), b[:3] if b.dim() else b)

BatchNorm2d(8, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
weight (8,) tensor([1., 1., 1.])
bias (8,) tensor([0., 0., 0.])
running_mean (8,) tensor([0., 0., 0.])
running_var (8,) tensor([1., 1., 1.])
num_batches_tracked () tensor(0)

The weight starts at 1 and the bias at 0, so a fresh BatchNorm normalizes and changes nothing more. The running mean starts at 0 and the running variance at 1, and num_batches_tracked counts the training batches it has seen, none yet. affine=True is the learned weight and bias, track_running_stats=True the buffers. One batch of 64 images through eight seeded filters, before and after:

x, y = next(iter(train_loader))
torch.manual_seed(42)
conv = nn.Conv2d(1, 8, kernel_size=3, padding=1)
with torch.no_grad():
    maps = conv(x)
    out = bn(maps)
mean = maps.mean(dim=(0, 2, 3))
var = maps.var(dim=(0, 2, 3), unbiased=False)
print(x.shape, maps.shape)
print(mean.round(decimals=3))
print(var.sqrt().round(decimals=3))
print(out.mean(dim=(0, 2, 3)).abs().max().item() < 1e-5)
print(out.std(dim=(0, 2, 3)).round(decimals=3))

torch.Size([64, 1, 28, 28]) torch.Size([64, 8, 28, 28])
tensor([ 0.6260, -0.0430, -0.5250, -0.2800,  0.0250,  0.0300,  0.1910,  0.2800])
tensor([0.3470, 0.2740, 0.2270, 0.0950, 0.2890, 0.1380, 0.1020, 0.0590])
True
tensor([1.0000, 1.0000, 1.0000, 0.9990, 1.0000, 1.0000, 1.0000, 0.9990])

dim=(0, 2, 3) takes each statistic over the 64 images and the 784 positions, leaving one number per channel, and unbiased=False divides by the count itself, as BatchNorm does. The eight channels arrive with means from -0.525 to 0.626 and standard deviations from 0.059 to 0.347, and leave with means within 1e-5 of zero and standard deviations of 1. Channels 3 and 7 leave at 0.999, because their variances, about 0.009 and 0.0035, are small enough for the eps of 1e-5 added to them to show. The same output by hand, and the buffers after this one batch:

by_hand = (maps - mean.view(1, 8, 1, 1)) / torch.sqrt(var.view(1, 8, 1, 1) + bn.eps)
print(torch.allclose(out, by_hand, atol=1e-5))
print(torch.allclose(bn.running_mean, 0.1 * mean), bn.num_batches_tracked)
print(torch.allclose(bn.running_var, 0.9 + 0.1 * var))

True
True tensor(1)
True

The output is the formula, with the weight of 1 and the bias of 0 leaving it as it stands. The running mean moved a tenth of the way from 0 to the batch's means, and the running variance a tenth of the way from 1 to the batch's variances. The two learned numbers are what let the network keep a channel's own scale where the loss wants it. Set to the batch's standard deviation and mean, in a second BatchNorm so the first keeps its buffers, they undo the normalization:

undo = nn.BatchNorm2d(8)
with torch.no_grad():
    undo.weight.copy_(torch.sqrt(var + undo.eps))
    undo.bias.copy_(mean)
    back = undo(maps)
print(torch.allclose(back, maps, atol=1e-5))

True

The output is the convolution's own maps again, within 1e-5. A trained network sets these 16 numbers by gradient like any others, so BatchNorm costs it nothing it could not get back: the normalization fixes the scale each batch arrives at, and the weight and the bias choose the scale the next layer is handed. Evaluation mode on the same batch:

bn.eval()
with torch.no_grad():
    judged = bn(maps)
print(judged.mean(dim=(0, 2, 3)).round(decimals=3))
print(judged.std(dim=(0, 2, 3)).round(decimals=3))
bn.train()

tensor([ 0.5900, -0.0400, -0.4960, -0.2650,  0.0230,  0.0290,  0.1810,  0.2650])
tensor([0.3640, 0.2880, 0.2380, 0.1000, 0.3030, 0.1450, 0.1080, 0.0620])

In evaluation mode BatchNorm subtracts a running mean that has moved a tenth of the way and divides by a running variance still near 1, so after one batch it leaves the values close to where they came in. A fresh BatchNorm shown a hundred training batches, then put in evaluation mode and handed 1,000 validation images no batch held:

settled = nn.BatchNorm2d(8)
for i, (batch, labels) in enumerate(train_loader):
    with torch.no_grad():
        settled(conv(batch))
    if i == 99:
        break
settled.eval()
val = torch.stack([val_set[i][0] for i in range(1000)])
with torch.no_grad():
    held = settled(conv(val))
print(settled.num_batches_tracked)
print(held.mean(dim=(0, 2, 3)).round(decimals=3))
print(held.std(dim=(0, 2, 3)).round(decimals=3))

tensor(100)
tensor([-0.0130, -0.0120,  0.0120,  0.0030, -0.0120,  0.0110, -0.0100,  0.0000])
tensor([0.9970, 0.9960, 0.9960, 0.9960, 0.9960, 0.9980, 0.9960, 1.0020])

After 100 updates the starting zero and one carry a weight of 0.9 to the hundredth power, under 0.00003, and the running statistics are the training batches' own. With them the validation images leave every channel with a mean within 0.013 of zero and a standard deviation within 0.004 of 1, statistics gathered on training batches normalizing images none of those batches held. An epoch of train_loader is 782 such updates. What the normalization does to a deep stack, eight convolutions of 16 filters on 128 training images, with nothing between them but the ReLU, and with a BatchNorm before every ReLU:

images = torch.stack([train_set[i][0] for i in range(128)])
torch.manual_seed(42)
plain = [nn.Conv2d(1 if i == 0 else 16, 16, 3, padding=1) for i in range(8)]
torch.manual_seed(42)
normed = [nn.Conv2d(1 if i == 0 else 16, 16, 3, padding=1, bias=False) for i in range(8)]
norms = [nn.BatchNorm2d(16) for _ in range(8)]
with torch.no_grad():
    a = b = images
    for i in range(8):
        a = torch.relu(plain[i](a))
        b = torch.relu(norms[i](normed[i](b)))
        print(i + 1, f"{a.std().item():.4f}", f"{b.std().item():.4f}")

1 0.2115 0.5527
2 0.0737 0.5718
3 0.0271 0.5952
4 0.0284 0.6148
5 0.0219 0.5818
6 0.0258 0.5771
7 0.0340 0.5951
8 0.0322 0.6103

1 if i == 0 else 16 gives the first layer one input channel and every later one sixteen. The two stacks hold the same seeded weights. Without BatchNorm the spread of the values falls from 0.2115 after one layer to 0.0271 after three and stays near that, a seventh of where it started. With it the spread stays between 0.55 and 0.62 at every depth, where the ReLU of values with a mean of 0 and a standard deviation of 1 leaves it. What this run shows is the scale at the start, before any training, and not how either stack would train. Training mode has a cost: the same image, in two batches.

torch.manual_seed(42)
model = nn.Sequential(nn.Conv2d(1, 8, 3, padding=1, bias=False), nn.BatchNorm2d(8), nn.ReLU())
batch_a = images[:64]
batch_b = torch.cat([images[:1], images[64:127]])
with torch.no_grad():
    out_a, out_b = model(batch_a)[0], model(batch_b)[0]
print(torch.equal(batch_a[0], batch_b[0]), round((out_a - out_b).abs().max().item(), 4))
model.eval()
with torch.no_grad():
    out_a, out_b = model(batch_a)[0], model(batch_b)[0]
    alone = model(images[:1])[0]
print(torch.allclose(out_a, out_b), torch.allclose(out_a, alone))

True 0.1028
True True

Both batches open with the same image, and in training mode its eight maps differ by up to 0.1028 between them, because each batch's mean and variance come from its own 64 images. In evaluation mode the running statistics are fixed, and the image gets the same maps in both batches and alone. A classifier left in training mode gives an image a prediction that depends on what else is in its batch. accuracy in image_loop.py, Lesson 4.5.2's, already switches to evaluation mode before it scores and back to training mode after, which the fully connected network never needed and every network of this section does. And one value per channel is not a batch:

try:
    nn.BatchNorm2d(8)(torch.zeros(1, 8, 1, 1))
except ValueError as err:
    print(err)

Expected more than 1 value per channel when training, got input size torch.Size([1, 8, 1, 1])

A batch of one image whose maps are 1 by 1 leaves each channel one value, whose variance is zero, and training mode refuses it. A batch of one 28 by 28 image passes, with each channel's statistics taken over that one image's 784 positions.

DON'T:

• DON'T score a model holding BatchNorm in training mode. The same image's maps differed by up to 0.1028 between two batches, and in evaluation mode they matched.
• DON'T read the normalization as throwing away the scale a convolution wrote. The weight and the bias set to the batch's standard deviation and mean gave back the convolution's own maps within 1e-5, and training sets them wherever the loss wants them.
• DON'T trust evaluation mode on running statistics that have seen one batch. After one batch the running mean had moved a tenth of the way and the outputs stayed close to the unnormalized values, where after a hundred the validation images came out within 0.013 of a zero mean.

The Conv Block

What. A conv block is nn.Conv2d(c_in, c_out, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(c_out), nn.ReLU() and nn.MaxPool2d(2), in one nn.Sequential. It sets the channel count, keeps the side through the convolution and halves it through the pooling, so each block takes a map from (N, c_in, H, W) to (N, c_out, H // 2, W // 2). blocks.py holds it as one function, conv_block(c_in, c_out). Why. A network stacks the same four steps several times, and a function that builds them keeps every block identical and every network built from them one line per block. The convolution carries no bias because BatchNorm removes it, and the ReLU sits before the pooling, though the two may swap, both proved below. When. As the repeated unit of Lesson 5.1.3's network, and wherever a model reads an image through a stack of convolutions.

import torch
from torch import nn
from fashion import train_loader, train_set

images = torch.stack([train_set[i][0] for i in range(128)])
torch.manual_seed(42)
with_bias = nn.Conv2d(1, 8, 3, padding=1)
without = nn.Conv2d(1, 8, 3, padding=1, bias=False)
bn = nn.BatchNorm2d(8)
with torch.no_grad():
    without.weight.copy_(with_bias.weight)
    print(with_bias.bias.round(decimals=4))
    print(torch.allclose(bn(with_bias(images)), bn(without(images)), atol=1e-5))

tensor([ 0.3208, -0.2751, -0.3306, -0.2608, -0.2242,  0.1350,  0.1194,  0.2770])
True

The two layers share their 72 weights, and one adds a bias from 0.3208 down to -0.3306 to each channel. A bias adds the same number to every value of its channel, so it raises the channel's batch mean by exactly that number, and BatchNorm subtracts the mean: after it, the two layers' outputs agree within 1e-5. A bias before a BatchNorm trains, takes its share of every step, and changes nothing, so the block's convolution goes without one, eight parameters saved on eight channels. ReLU and max pooling, in both orders:

pool = nn.MaxPool2d(2)
with torch.no_grad():
    normed = bn(without(images))
print(torch.equal(pool(torch.relu(normed)), torch.relu(pool(normed))))

True

The largest of four values passed through a ReLU is the ReLU of the largest of the four, because the ReLU never changes which of two values is larger, so the two orders give the same tensor to the last bit. This block keeps the ReLU first, the order most models print, and the other order would apply the ReLU to a quarter of the values. Save it as blocks.py in ~/dev/bootcamp:

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(),
        nn.MaxPool2d(2),
    )

A batch from train_loader through two blocks, the first from 1 channel to 16 and the second from 16 to 32:

import torch
from torch import nn
from fashion import train_loader
from blocks import conv_block

torch.manual_seed(42)
body = nn.Sequential(conv_block(1, 16), conv_block(16, 32))
print(body[0])
x, y = next(iter(train_loader))
h = x
for block in body:
    h = block(h)
    print(tuple(h.shape), sum(p.numel() for p in block.parameters()))

Sequential(
  (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
  (1): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
  (2): ReLU()
  (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(64, 16, 14, 14) 176
(64, 32, 7, 7) 4672

The first block turns 64 one-channel images of 28 by 28 into 16 maps of 14 by 14 each, holding 144 weights and BatchNorm's 32. The second turns those into 32 maps of 7 by 7, holding 32 * 16 * 9, 4,608 weights and 64. What the two convolutions compute per image, counted the way Lesson 4.5.3 counted its first layer's 100,352 multiplications, one per weight per position:

counts = [28 * 28 * 16 * 1 * 9, 14 * 14 * 32 * 16 * 9]
print(counts, sum(counts), sum(counts) / 100352)

[112896, 903168] 1016064 10.125

The first convolution writes 16 values at each of 784 positions from 9 weights each, and the second writes 32 at each of 196 positions from 144 weights each, 903,168 multiplications, eight times the first's, because it reads sixteen channels. The two blocks compute 10.125 times what the fully connected network's first layer computed, with 4,848 parameters against its 100,480: a filter is reused at every position, so it costs few parameters and many multiplications. What torch.save would keep, Lesson 4.4.3's state_dict:

state = body.state_dict()
print(list(state)[:6])
print(len(state), sum(p.numel() for p in body.parameters()), sum(b.numel() for b in body.buffers()))
print(body[0][1].num_batches_tracked)

['0.0.weight', '0.1.weight', '0.1.bias', '0.1.running_mean', '0.1.running_var', '0.1.num_batches_tracked']
12 4848 98
tensor(1)

0.1.running_mean is block 0's layer 1, its BatchNorm, and the state_dict holds twelve entries: per block the convolution's weight, BatchNorm's weight and bias, and its three buffers. The 4,848 parameters train, and the 98 buffer values, 16 and 16 and 1 for the first block, 32 and 32 and 1 for the second, are saved beside them, so a model reloaded from its file evaluates with the running statistics it trained with. The one batch traced above was a training-mode batch, and num_batches_tracked counted it. The body saved and loaded into a fresh one, beside a fresh body handed the 4,848 parameters and none of the buffers:

torch.save(state, "body.pt")
reloaded = nn.Sequential(conv_block(1, 16), conv_block(16, 32))
reloaded.load_state_dict(torch.load("body.pt", weights_only=True))
params_only = nn.Sequential(conv_block(1, 16), conv_block(16, 32))
with torch.no_grad():
    for mine, saved in zip(params_only.parameters(), body.parameters()):
        mine.copy_(saved)
for model in (body, reloaded, params_only):
    model.eval()
with torch.no_grad():
    ref = body(x)
    print(torch.equal(reloaded(x), ref), round((params_only(x) - ref).abs().max().item(), 4))
print(reloaded[0][1].num_batches_tracked, params_only[0][1].num_batches_tracked)

True 0.0782
tensor(1) tensor(0)

load_state_dict copies every entry of the file, buffers included, so in evaluation mode the reloaded body writes exactly what the saved one writes. zip pairs the two bodies' parameters in order and copies each, which leaves the copy's buffers at their fresh 0 and 1: the same 4,848 parameters, in evaluation mode, write values up to 0.0782 away, from statistics one batch apart. After a full epoch of training the gap between a body's running statistics and fresh ones is the whole of what the BatchNorms learned about their channels' scale.

DON'T:

• DON'T give a convolution a bias when a BatchNorm follows it. The outputs with and without the bias agreed within 1e-5, and the bias is eight parameters that train for nothing.
• DON'T read 4,848 parameters as a network that is cheap to run. The two blocks multiply 1,016,064 times per image, 10.125 times the fully connected network's first layer, which held more than twenty times their parameters.
• DON'T push a batch through a model in training mode to look at its shapes. The trace updated every BatchNorm's running statistics, and num_batches_tracked counted the batch.
• DON'T copy a BatchNorm model by its parameters alone. The body handed every parameter and no buffer wrote values up to 0.0782 away from the saved body in evaluation mode, where load_state_dict matched it exactly.

Lab Exercises

Lab Exercise 1: A Grid Pooled Two Ways, and Three Running Means

In pool_norm.py, build the 6 by 6 grid torch.arange(36.0).view(1, 1, 6, 6) and print it pooled by nn.MaxPool2d(2) and by nn.AvgPool2d(3), then the shape a 28 by 28 map takes through nn.MaxPool2d(3). Say in a comment why every max-pooled value of this grid sits in a block's bottom right corner.

Seed 42 and build nn.Conv2d(1, 4, kernel_size=3, padding=1) and nn.BatchNorm2d(4). Run the first three batches of train_loader through both under torch.no_grad(), keeping each batch's per-channel mean of the convolution's output, then stop. Print num_batches_tracked and the running mean to four decimals, then compute the running mean by hand from the three batch means with BatchNorm's update, starting from zero, and print whether it matches.

Lab Exercise 2: Three Blocks, and a Block That Never Trained

In block_check.py, seed 42 and build a fresh conv_block(1, 8) from blocks.py. Put it in evaluation mode and print whether its BatchNorm's output on the first validation image is its convolution's output divided by sqrt(1 + eps), and say in a comment what a BatchNorm that has never seen a training batch does in evaluation mode.

Then seed 42 and build a body of three blocks, 1 to 8, 8 to 16 and 16 to 32 channels. Trace one batch from train_loader through it, printing each block's output shape, parameter count and buffer count. Stack the first 127 validation images and, in training mode, print the largest difference between the body's output for validation image 0 in a batch of images 0 to 63 and in a batch of image 0 with images 64 to 126. Switch to evaluation mode and print whether the two now agree, and what the first block's num_batches_tracked reads.

import torch
from torch import nn
from fashion import train_loader

grid = torch.arange(36.0).view(1, 1, 6, 6)
print(nn.MaxPool2d(2)(grid)[0, 0])
print(nn.AvgPool2d(3)(grid)[0, 0])
print(nn.MaxPool2d(3)(torch.zeros(1, 1, 28, 28)).shape)
# The grid grows left to right and top to bottom, so in every
# block the largest value is the one furthest right and lowest.

tensor([[ 7.,  9., 11.],
        [19., 21., 23.],
        [31., 33., 35.]])
tensor([[ 7., 10.],
        [25., 28.]])
torch.Size([1, 1, 9, 9])

torch.manual_seed(42)
conv = nn.Conv2d(1, 4, kernel_size=3, padding=1)
bn = nn.BatchNorm2d(4)
means = []
for i, (x, y) in enumerate(train_loader):
    with torch.no_grad():
        maps = conv(x)
        bn(maps)
    means.append(maps.mean(dim=(0, 2, 3)))
    if i == 2:
        break
running = torch.zeros(4)
for m in means:
    running = 0.9 * running + 0.1 * m
print(bn.num_batches_tracked, bn.running_mean.round(decimals=4))
print(torch.allclose(bn.running_mean, running))

tensor(3) tensor([ 0.0989,  0.1375, -0.0440, -0.0337])
True

import torch
from torch import nn
from fashion import train_loader, val_set
from blocks import conv_block

val = torch.stack([val_set[i][0] for i in range(127)])
torch.manual_seed(42)
fresh = conv_block(1, 8)
fresh.eval()
with torch.no_grad():
    raw = fresh[0](val[:1])
    normed = fresh[1](raw)
print(torch.allclose(normed, raw / (1 + fresh[1].eps) ** 0.5))
# With a running mean of 0, a running variance of 1, a weight of 1
# and a bias of 0, an untrained BatchNorm in evaluation mode divides
# by sqrt(1 + eps) and otherwise passes its input through.

True

torch.manual_seed(42)
body = nn.Sequential(conv_block(1, 8), conv_block(8, 16), conv_block(16, 32))
x, y = next(iter(train_loader))
h = x
for block in body:
    h = block(h)
    params = sum(p.numel() for p in block.parameters())
    print(tuple(h.shape), params, sum(b.numel() for b in block.buffers()))
batch_a = val[:64]
batch_b = torch.cat([val[:1], val[64:127]])
with torch.no_grad():
    gap = (body(batch_a)[0] - body(batch_b)[0]).abs().max().item()
print(round(gap, 4))
body.eval()
with torch.no_grad():
    same = torch.allclose(body(batch_a)[0], body(batch_b)[0])
print(same, body[0][1].num_batches_tracked)

(64, 8, 14, 14) 88 17
(64, 16, 7, 7) 1184 33
(64, 32, 3, 3) 4672 65
0.435
True tensor(3)

Theory Assessment

Q1: Two Pixels or One

A teammate reads that max pooling gives a network tolerance for small moves and concludes that a network with two pooling layers ignores any move of up to four pixels. Using the lit pixel of the first subtitle and the share of pooled values a one-pixel move left unchanged on the ankle boot's maps, say what one nn.MaxPool2d(2) does with a move of one pixel and of two, why the answer depends on where a feature sits, and what the 0.6084 says about the teammate's claim.

Hint: Ask which block the lit pixel is in at columns 0, 1, 2 and 3, and what a pooled map can and cannot do when its input crosses from one block into the next.

Q2: A Score That Depends on the Batch Size

A teammate scores a network built from conv blocks with a loop of their own that skips model.eval(), and gets one validation accuracy with batches of 1,000 and another with batches of 10. Say what each BatchNorm uses to normalize in that loop, why the two accuracies differ, what changes when model.eval() is called first, and which function of image_loop.py already does this.

Hint: Compare where a batch's mean and variance come from with where the running statistics come from, and ask which of the two depends on the other images in the batch.

Q3: Forty-Eight Parameters That Do Nothing

Lesson 5.1.3's network uses two conv blocks, 1 to 16 and 16 to 32 channels. A teammate rewrites conv_block with the convolution's default bias=True. Say how many parameters that adds to the two blocks, why the outputs after each BatchNorm do not change, and what the BatchNorm's own bias does that the convolution's could not.

Hint: Add a constant to every value of one channel of a batch and ask what happens to that channel's mean and to its value minus the mean.

Q4: One Layer at Serving Time

A team serving the network from Lesson 5.1.3 wants to drop every BatchNorm from the served model to save a step per block, by folding it into the convolution before it: a new convolution whose weight is the old weight times bn.weight / sqrt(running_var + eps) per output channel, and whose bias is bn.bias - running_mean * bn.weight / sqrt(running_var + eps). Say why the fold is exact in evaluation mode, why it cannot be exact in training mode, and show it on a seeded conv_block(1, 8) after twenty training batches.

Hint: Write BatchNorm's evaluation output for one channel as a * x + b and ask which of a and b depend on the batch. Then write the convolution's output as a weighted sum and multiply it by a.

import torch
from torch import nn
from fashion import train_loader
from blocks import conv_block

torch.manual_seed(42)
block = conv_block(1, 8)
for i, (x, y) in enumerate(train_loader):
    with torch.no_grad():
        block(x)
    if i == 19:
        break
block.eval()
conv, bn = block[0], block[1]
fused = nn.Conv2d(1, 8, kernel_size=3, padding=1)
with torch.no_grad():
    scale = bn.weight / torch.sqrt(bn.running_var + bn.eps)
    fused.weight.copy_(conv.weight * scale.view(8, 1, 1, 1))
    fused.bias.copy_(bn.bias - bn.running_mean * scale)
    gap = (bn(conv(x)) - fused(x)).abs().max().item()
print(bn.num_batches_tracked, gap < 1e-5)

tensor(20) True

Assignment

Assignment: A Normalization That Ignores the Batch

BatchNorm's training-mode output depends on the batch, and it refuses a batch that leaves one value per channel. Write group_norm.py, which replaces it in one block with a normalization that takes its statistics from each image alone.

One layer this lesson did not use. nn.GroupNorm(num_groups, num_channels) splits each image's channels into num_groups groups and normalizes every group with the mean and variance of that one image's values in that group, over its channels and positions, then applies a learned weight and bias per channel, as BatchNorm does. It keeps no running statistics, so training and evaluation mode compute the same thing, and num_channels must be divisible by num_groups.

Build nn.GroupNorm(4, 8) and print it, the shapes of its parameters, and the keys of its state_dict, beside a BatchNorm2d(8)'s keys. Seed 42 and build nn.Sequential(nn.Conv2d(1, 8, 3, padding=1, bias=False), nn.GroupNorm(4, 8), nn.ReLU()), and in training mode print whether validation image 0's output agrees across a batch of images 0 to 63, a batch of image 0 with images 64 to 126, and image 0 alone. Show that nn.GroupNorm(4, 8) takes the (1, 8, 1, 1) batch BatchNorm refused, and print the one image's group statistics by hand for group 0 of the convolution's output, its mean and its variance over channels 0 and 1 and all 784 positions, to four decimals, then confirm the block's GroupNorm output for those two channels matches the formula.

Hints:

• With 8 channels in 4 groups, group 0 is channels 0 and 1, so maps[0, 0:2] holds its 2 times 784 values, and unbiased=False gives the variance GroupNorm divides by.
• A batch of (1, 8, 1, 1) leaves each group two values, one per channel, enough for a variance, which is why GroupNorm takes it.

import torch
from torch import nn
from fashion import val_set

gn = nn.GroupNorm(4, 8)
print(gn, [tuple(p.shape) for p in gn.parameters()])
print(list(gn.state_dict()), list(nn.BatchNorm2d(8).state_dict()))
val = torch.stack([val_set[i][0] for i in range(127)])
torch.manual_seed(42)
block = nn.Sequential(nn.Conv2d(1, 8, 3, padding=1, bias=False), nn.GroupNorm(4, 8), nn.ReLU())
with torch.no_grad():
    in_a = block(val[:64])[0]
    in_b = block(torch.cat([val[:1], val[64:127]]))[0]
    alone = block(val[:1])[0]
print(block.training, torch.allclose(in_a, in_b, atol=1e-6), torch.allclose(in_a, alone, atol=1e-6))
print(gn(torch.zeros(1, 8, 1, 1)).shape)

GroupNorm(4, 8, eps=1e-05, affine=True, bias=True) [(8,), (8,)]
['weight', 'bias'] ['weight', 'bias', 'running_mean', 'running_var', 'num_batches_tracked']
True True True
torch.Size([1, 8, 1, 1])

with torch.no_grad():
    maps = block[0](val[:1])
    group = maps[0, 0:2]
    mean, var = group.mean(), group.var(unbiased=False)
    by_hand = (group - mean) / torch.sqrt(var + block[1].eps)
    out = block[1](maps)[0, 0:2]
print(f"{mean.item():.4f}", f"{var.item():.4f}")
print(torch.allclose(out, by_hand, atol=1e-5))

0.4091 0.1031
True

Additional Resources

Optional, and nothing later depends on either. Worth opening for the argument over why BatchNorm helps a network train, which this lesson measured the start of and never settled.

• YouTube - "Why Does Batch Norm Work? (C2W3L06)" by DeepLearningAI, Andrew Ng's course lecture. Covariate shift, a later layer's inputs moving as the earlier layers train, as the course's account of why normalizing each layer helps, and the small regularizing effect of a mini-batch's noisy mean and variance.
• Book chapter - Dive into Deep Learning, "Batch Normalization" by Zhang, Lipton, Li and Smola, under CC BY-SA 4.0. BatchNorm written from scratch for fully connected and convolution layers, its prediction-time behavior, layer normalization beside it, and the discussion of why it works.

Summary

Pooling and BatchNorm are small layers with large consequences. Pooling holds no weight and decides what every later layer sees and how much a feature may move before it shows: some, a share that depends on where each feature sits, and never the full tolerance the word promises. A move larger than a block still arrives as a move, half its size.

BatchNorm is the first layer in this bootcamp whose output depends on more than its own input. In training it takes its numbers from the batch, which is what keeps eight convolutions deep at a usable scale, and in evaluation it takes them from a memory of every batch it saw, which is what lets one image get one answer. The switch between the two was always in accuracy, and it is now load-bearing: a network that forgets it is wrong by an amount that depends on what else was in the batch.

And the conv block shows how much of a design is arithmetic rather than taste. The bias goes because a subtraction undoes it, the ReLU and the pooling may trade places because one never reorders what the other compares, and the fold of a BatchNorm into its convolution is exact only once the statistics stop moving. The same arithmetic prices the block: a filter reused at every position is cheap to store and expensive to run, so two blocks with under a twentieth of the fully connected network's first-layer parameters do ten times its multiplications, and that is the bill an epoch will present. Lesson 5.1.3 stacks two of these blocks, trains them, and finds out whether they beat the network that read pixels one weight at a time.

Cheat Sheet

Part A - Reference Table

Part B - Compact Code Block

def conv_block(c_in, c_out):
    return nn.Sequential(
        nn.Conv2d(c_in, c_out, 3, padding=1, bias=False),
        nn.BatchNorm2d(c_out),  # mean 0, std 1 per channel
        nn.ReLU(),
        nn.MaxPool2d(2),        # 28 to 14
    )

body = nn.Sequential(conv_block(1, 16), conv_block(16, 32))
body.eval()         # running statistics from here
with torch.no_grad():
    maps = body(x)  # (N, 32, 7, 7)
body.train()        # batch statistics again