Inside Batch Norm

Here's a classic bug. A model gets excellent results in training. You save it, and load it up, and it predicts nonsense. Nothing has changed, except for one line. You forgot to switch it to evaluation mode. To see why that matters, we need to look inside batch normalisation.

• Excellent in training
• Nonsense after loading
• You forgot one line

What batch norm does

Here are three columns on very different scales. Normalising subtracts the mean, and divides by the spread. Run it. Spreads of one, ten, and ninety seven all become exactly one. That's the first job of batch norm. It gives each layer steady, well behaved inputs. But now look at the third line. The stored mean isn't fifty. It's five! That's the second life. Batch norm keeps a slow moving average during training, because at test time you might have a single image, and one image has no batch to average. And it updates slowly, by ten percent each time, so here we've only seen one batch. That's why training mode and evaluation mode give different answers.

import torch

x = torch.randn(1000, 3) * torch.tensor([1.0, 10.0, 100.0]) + 50
print("Before:", x.mean(0).round().tolist(), x.std(0).round().tolist())

mean, var = x.mean(0), x.var(0)
z = (x - mean) / var.sqrt()
print("After: ", z.mean(0).abs().round().tolist(), z.std(0).round().tolist())

bn = torch.nn.BatchNorm1d(3)
bn(x)
print("Stored mean:", bn.running_mean.round().tolist())
bn.eval()
print("One item at test time works:", bn(x[:1]).shape)

Before: [50.0, 50.0, 48.0] [1.0, 10.0, 103.0]
After:  [0.0, 0.0, 0.0] [1.0, 1.0, 1.0]
Stored mean: [5.0, 5.0, 5.0]
One item at test time works: torch.Size([1, 3])

Batch norm rules

Here's how to stay safe. Always call model dot eval before testing, and model dot train before training. Very small batches give noisy statistics, so with a batch size of two or four, use group norm or layer norm instead. Transformers use layer norm for exactly that reason. And remember that when you fine tune a model, its stored statistics matter too.

Summary