Batch Normalisation Explained: Internals, Gotchas and Production Reality

import torch
import torch.nn as nn

# ─────────────────────────────────────────────
# Manual BatchNorm implementation — fully from scratch.
# This mirrors exactly what nn.BatchNorm1d does internally
# so you can reason about it under a debugger.
# ─────────────────────────────────────────────

class ManualBatchNorm1d(nn.Module):
    def __init__(self, num_features: int, eps: float = 1e-5, momentum: float = 0.1):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum

        # Learnable affine parameters — gamma starts at 1, beta at 0
        # so initially the layer behaves like pure normalisation.
        self.gamma = nn.Parameter(torch.ones(num_features))
        self.beta = nn.Parameter(torch.zeros(num_features))

        # Running stats — NOT parameters, so they won't appear in
        # model.parameters() and won't be touched by the optimiser.
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))

    def forward(self, activations: torch.Tensor) -> torch.Tensor:
        # activations: shape (batch_size, num_features)
        if activations.dim() != 2:
            raise ValueError(f'ManualBatchNorm1d expects 2D input, got {activations.dim()}D')

        if self.training:
            # ── TRAINING PATH ─────────────────────────────────────
            # Compute statistics over the BATCH dimension (dim=0)
            batch_mean = activations.mean(dim=0)          # shape: (num_features,)
            batch_var  = activations.var(dim=0, unbiased=False)  # biased variance, matches PyTorch

            # Normalise using BATCH statistics
            activations_normalised = (activations - batch_mean) / torch.sqrt(batch_var + self.eps)

            # Update running stats with exponential moving average.
            # Note: no_grad not needed here — buffers aren't tracked by autograd.
            self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * batch_mean.detach()
            self.running_var  = (1 - self.momentum) * self.running_var  + self.momentum * batch_var.detach()

        else:
            # ── INFERENCE PATH ────────────────────────────────────
            # Use the accumulated running stats — NOT the current batch.
            # The batch might be size 1 at inference; that's fine here.
            activations_normalised = (activations - self.running_mean) / torch.sqrt(self.running_var + self.eps)

        # Apply learnable scale and shift — this is what lets the network
        # undo the normalisation if that's actually optimal.
        output = self.gamma * activations_normalised + self.beta
        return output


# ─── Smoke test: compare our manual layer to PyTorch's built-in ───
torch.manual_seed(42)
batch_size, features = 32, 16

sample_activations = torch.randn(batch_size, features) * 5 + 3  # deliberately off-centre

manual_bn = ManualBatchNorm1d(num_features=features)
pytorch_bn = nn.BatchNorm1d(num_features=features)

# Force both to have identical weights (defaults are already the same,
# but let's be explicit so this is reproducible)
with torch.no_grad():
    pytorch_bn.weight.copy_(manual_bn.gamma)
    pytorch_bn.bias.copy_(manual_bn.beta)

manual_bn.train()
pytorch_bn.train()

manual_output  = manual_bn(sample_activations)
pytorch_output = pytorch_bn(sample_activations)

max_diff = (manual_output - pytorch_output).abs().max().item()
print(f'Max absolute difference (training mode): {max_diff:.2e}')
print(f'Manual running mean (first 4):  {manual_bn.running_mean[:4].tolist()}')
print(f'PyTorch running mean (first 4): {pytorch_bn.running_mean[:4].tolist()}')

# Switch to eval and verify inference path
manual_bn.eval()
pytorch_bn.eval()

inference_input = torch.randn(1, features)  # batch size 1 — totally fine in eval
manual_eval  = manual_bn(inference_input)
pytorch_eval = pytorch_bn(inference_input)

eval_diff = (manual_eval - pytorch_eval).abs().max().item()
print(f'Max absolute difference (eval mode):     {eval_diff:.2e}')

import torch
import torch.nn as nn

# ─────────────────────────────────────────────────────────────────
# Demonstrate the silent accuracy difference between train() and eval().
# We'll train a tiny network, then compare inference accuracy in both modes.
# ─────────────────────────────────────────────────────────────────

torch.manual_seed(0)

class TinyClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(8, 32),
            nn.BatchNorm1d(32),
            nn.ReLU(),
            nn.Linear(32, 2)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)

model = TinyClassifier()
optimiser = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

# ── Simulate training on normally-distributed data ─────────────────
model.train()
for step in range(200):
    training_batch = torch.randn(64, 8)          # 64 samples, 8 features
    training_labels = torch.randint(0, 2, (64,))
    logits = model(training_batch)
    loss = loss_fn(logits, training_labels)
    optimiser.zero_grad()
    loss.backward()
    optimiser.step()

print('Training complete. Running stats are now well-calibrated.\n')

# ── Inspect the accumulated running stats ─────────────────────────
bn_layer = model.net[1]  # the BatchNorm1d layer
print(f'Running mean (first 4 features): {bn_layer.running_mean[:4].tolist()}')
print(f'Running var  (first 4 features): {bn_layer.running_var[:4].tolist()}\n')

# ── Inference with a single sample ────────────────────────────────
single_input = torch.randn(1, 8)

# CORRECT: eval mode uses running stats — deterministic and stable
model.eval()
with torch.no_grad():
    logits_eval = model(single_input)
    prob_eval = torch.softmax(logits_eval, dim=1)

print(f'eval()  mode probabilities: {prob_eval.tolist()}')

# WRONG: train mode tries to compute variance from a single sample
# PyTorch will actually raise a warning or silently use NaN variance
model.train()   # <─── forgot to call eval() — common bug
with torch.no_grad():
    try:
        logits_train = model(single_input)
        prob_train = torch.softmax(logits_train, dim=1)
        print(f'train() mode probabilities: {prob_train.tolist()}')
        nan_count = prob_train.isnan().sum().item()
        print(f'NaN values in output: {nan_count}  ← this is the bug')
    except Exception as e:
        print(f'Error in train mode with batch_size=1: {e}')

print('\n>>> Always call model.eval() before inference. No exceptions.')

import torch
import torch.nn as nn
import time

# ─────────────────────────────────────────────────────────────────────────
# Compare BatchNorm, LayerNorm, GroupNorm on the same activation tensor.
# We'll inspect output statistics and measure timing to surface real trade-offs.
# ─────────────────────────────────────────────────────────────────────────

torch.manual_seed(7)

# Simulated convolutional feature map: (batch, channels, height, width)
batch_size = 4
channels   = 64
height     = 28
width      = 28

# Deliberately skewed input: each channel has a different mean/variance
channel_means = torch.linspace(-3.0, 3.0, channels).view(1, channels, 1, 1)
channel_stds  = torch.linspace(0.5, 5.0, channels).view(1, channels, 1, 1)
feature_maps  = torch.randn(batch_size, channels, height, width) * channel_stds + channel_means

print(f'Input  — mean: {feature_maps.mean().item():.4f}, std: {feature_maps.std().item():.4f}')

# ── BatchNorm2d: normalises over (N, H, W) per channel ───────────────────
batch_norm  = nn.BatchNorm2d(num_features=channels)
batch_norm.eval()  # use running stats path (freshly init'd: mean=0, var=1)

# Note: fresh BatchNorm2d in eval mode has running_mean=0, running_var=1,
# gamma=1, beta=0 — so output ≈ input (running stats haven't been trained).
# To demonstrate normalisation behaviour fairly we use it in train mode:
batch_norm.train()
bn_output = batch_norm(feature_maps)
print(f'BatchNorm — mean: {bn_output.mean().item():.4f}, std: {bn_output.std().item():.4f}')

# ── LayerNorm: normalises over (C, H, W) per sample ─────────────────────
# normalised_shape must cover the dims you want to normalise over.
layer_norm = nn.LayerNorm(normalized_shape=[channels, height, width])
ln_output  = layer_norm(feature_maps)
print(f'LayerNorm — mean: {ln_output.mean().item():.4f}, std: {ln_output.std().item():.4f}')

# ── GroupNorm: normalises over (C/G, H, W) per sample per group ──────────
# num_groups=32 is a robust default for channels >= 32
group_norm = nn.GroupNorm(num_groups=32, num_channels=channels)
gn_output  = group_norm(feature_maps)
print(f'GroupNorm — mean: {gn_output.mean().item():.4f}, std: {gn_output.std().item():.4f}')

# ── Timing comparison at small vs large batch size ────────────────────────
print('\n--- Timing: 500 forward passes ---')

for test_batch_size in [2, 64]:
    test_input = torch.randn(test_batch_size, channels, height, width)

    # BatchNorm
    batch_norm.train()
    t0 = time.perf_counter()
    for _ in range(500):
        _ = batch_norm(test_input)
    bn_time = (time.perf_counter() - t0) * 1000

    # GroupNorm (no train/eval distinction — always uses per-sample stats)
    t0 = time.perf_counter()
    for _ in range(500):
        _ = group_norm(test_input)
    gn_time = (time.perf_counter() - t0) * 1000

    print(f'  Batch size {test_batch_size:>2} — BatchNorm: {bn_time:.1f}ms, GroupNorm: {gn_time:.1f}ms')

import torch
import torch.nn as nn

# ─────────────────────────────────────────────────────────────────────────
# Pattern 1: Convert all BatchNorm layers to SyncBatchNorm for multi-GPU.
# Call this BEFORE wrapping with DistributedDataParallel.
# ─────────────────────────────────────────────────────────────────────────

class ResidualBlock(nn.Module):
    def __init__(self, num_channels: int):
        super().__init__()
        self.conv1 = nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, bias=False)
        self.bn1   = nn.BatchNorm2d(num_channels)
        self.relu  = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(num_channels, num_channels, kernel_size=3, padding=1, bias=False)
        self.bn2   = nn.BatchNorm2d(num_channels)

    def forward(self, feature_map: torch.Tensor) -> torch.Tensor:
        residual = feature_map
        out = self.relu(self.bn1(self.conv1(feature_map)))
        out = self.bn2(self.conv2(out))
        return self.relu(out + residual)

model = ResidualBlock(num_channels=64)

print('Before conversion:')
for name, module in model.named_modules():
    if isinstance(module, (nn.BatchNorm2d, nn.SyncBatchNorm)):
        print(f'  {name}: {type(module).__name__}')

# One-liner to convert every BatchNorm layer in the model graph
synced_model = nn.SyncBatchNorm.convert_sync_batchnorm(model)

print('\nAfter convert_sync_batchnorm():')
for name, module in synced_model.named_modules():
    if isinstance(module, (nn.BatchNorm2d, nn.SyncBatchNorm)):
        print(f'  {name}: {type(module).__name__}')


# ─────────────────────────────────────────────────────────────────────────
# Pattern 2: Freeze BatchNorm during fine-tuning.
# Keeps running stats from the pre-trained model unchanged.
# This is the correct approach when your fine-tuning dataset is small.
# ─────────────────────────────────────────────────────────────────────────

def freeze_batchnorm_layers(model: nn.Module) -> None:
    """Set every BatchNorm layer to eval mode so they use fixed running stats.
       The rest of the model (Linear, Conv) remains in train mode."""
    for module in model.modules():
        if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
            module.eval()          # use running stats, not batch stats
            # Also freeze the learnable gamma and beta if you want full freeze:
            module.weight.requires_grad_(False)
            module.bias.requires_grad_(False)

fine_tune_model = ResidualBlock(num_channels=64)
fine_tune_model.train()   # whole model in train mode initially

print('\nBefore freeze — bn1 training:', fine_tune_model.bn1.training)
freeze_batchnorm_layers(fine_tune_model)
print('After freeze  — bn1 training:', fine_tune_model.bn1.training)
print('After freeze  — conv1 training:', fine_tune_model.conv1.training)  # still True
print('After freeze  — bn1 weight requires_grad:', fine_tune_model.bn1.weight.requires_grad)