Batch Norm as an nn.Module

Write your own BatchNorm1d module that behaves like torch.nn.BatchNorm1d.

Batch norm has two modes, and the module must handle both:

• Training (model.train()): normalise each feature with the mean and the variance of the current batch. Then update the running statistics.
• Evaluation (model.eval()): normalise with the running statistics. Do not change them.

For an input x of shape (N, C):

mean = x.mean(dim=0)
var  = x.var(dim=0, correction=0)          # biased, used to normalise
y    = weight * (x - mean) / sqrt(var + eps) + bias

running_mean = (1 - momentum) * running_mean + momentum * mean
running_var  = (1 - momentum) * running_var  + momentum * var * N / (N - 1)

The running variance uses the unbiased variance, which is why it has the N / (N - 1) factor. PyTorch does the same.

weight (starts at ones) and bias (starts at zeros) are learnable, so make them nn.Parameters. running_mean (starts at zeros) and running_var (starts at ones) are state, not parameters. Register them with self.register_buffer(...), so that they are in the state_dict and move with .to(device), but the optimiser does not see them.

Function signature

class BatchNorm1d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1): ...
    def forward(self, x): ...

Example

Input

bn = BatchNorm1d(1, momentum=0.1)
bn(torch.tensor([[1.0], [3.0]]))

Output

tensor([[-1.0000],
        [ 1.0000]])
bn.running_mean  ->  tensor([0.2000])
bn.running_var   ->  tensor([1.1000])

Reasoning

The batch mean is 2 and the biased variance is 1, so the outputs are -1 and 1. The running mean moves from 0 to 0.1 * 2 = 0.2. The unbiased variance is 2, so the running variance moves from 1 to 0.9 * 1 + 0.1 * 2 = 1.1.

Constraints and notes

• The input has shape (N, C) with N >= 2.
• The running statistics must not be part of the autograd graph. Update them under torch.no_grad(), or from detached values.
• Gradients must flow to x, weight and bias in training mode. The tests compare them with nn.BatchNorm1d.
• Do not use nn.BatchNorm1d or F.batch_norm inside your module.