Residual blocks and normalization state

Define the addition contract

A residual block returns a transformed feature map plus a shortcut. Addition requires identical batch, channel, height and width dimensions. If the main branch changes width or downsamples, use a projection shortcut with matching stride and output channels; silently broadcasting a singleton channel is a bug, not a substitute. Assert branch shapes at the addition point. Convolution geometry predicts the new grid before the first optimizer step.

Place nonlinearities deliberately

The block below applies convolution, batch normalization and activation on the main path, then adds the shortcut and activates the sum. Other residual layouts exist, but changing the order changes the function and checkpoint compatibility. A zeroed main branch should still pass shortcut information through the final activation. The shortcut makes a short route for signals and gradients, yet it does not guarantee that an oversized learning rate or bad labels will train successfully. Compare training curves against a plain stack under a fixed data split.

Treat running statistics as state

BatchNorm has learned affine parameters and running mean/variance buffers. Setting requires_grad false freezes the affine parameters but does not stop running buffers from changing in train mode. During low-label transfer, a backbone can be frozen for gradients while its feature distribution drifts because those buffers keep updating. Evaluate and serve with eval mode; if the plan is to freeze the backbone entirely, keep its BatchNorm modules in eval mode even when the new head trains. The train/eval contract controls this state.

Handle small and uneven batches

Batch statistics become noisy with very small per-device batches. Gradient accumulation does not create a larger BatchNorm batch, because each forward pass still normalizes its own microbatch. A final batch of one can fail for a feature map with one spatial position. Options include a frozen pretrained normalization state, a different normalization layer, or a data loader that avoids the problematic batch; each changes the model behavior and needs validation. Accumulation changes update cadence, not normalization scope.

Test the transition to inference

Capture output for a fixed feature tensor in evaluation mode, save weights and buffers, reload into a fresh model, and compare output. Compare one train-mode pass separately to expose running-stat updates. Check projection and identity variants, odd spatial sizes and a batch of one. A checkpoint containing only trainable parameters is incomplete when nontrainable buffers affect inference; record a whole state dictionary and a model revision in the artifact.

Implementation

import torch
from torch import nn

class ReceiptResidualBlock(nn.Module):
    def __init__(self, input_channels: int, output_channels: int, stride: int = 1):
        super().__init__()
        self.main = nn.Sequential(
            nn.Conv2d(input_channels, output_channels, 3, stride, 1, bias=False),
            nn.BatchNorm2d(output_channels), nn.ReLU(),
            nn.Conv2d(output_channels, output_channels, 3, 1, 1, bias=False),
            nn.BatchNorm2d(output_channels))
        self.shortcut = (nn.Identity() if stride == 1 and input_channels == output_channels
                         else nn.Sequential(nn.Conv2d(input_channels, output_channels, 1,
                                                      stride, bias=False),
                                            nn.BatchNorm2d(output_channels)))
        self.activation = nn.ReLU()

    def forward(self, receipt_features: torch.Tensor) -> torch.Tensor:
        transformed = self.main(receipt_features)
        shortcut = self.shortcut(receipt_features)
        if transformed.shape != shortcut.shape:
            raise RuntimeError("residual branches have different shapes")
        return self.activation(transformed + shortcut)

block = ReceiptResidualBlock(12, 24, stride=2).eval()
with torch.inference_mode():
    result = block(torch.randn(2, 12, 47, 83))
assert result.shape == (2, 24, 24, 42)

Performance and operating cost

The two main 3 by 3 convolutions and optional 1 by 1 projection dominate arithmetic; the addition itself is linear in the output feature-map size. BatchNorm keeps channel-wise buffers and reads activation statistics during training. Projection blocks add parameters and compute, while identity shortcuts are cheap. Memory still grows with stored activations, so residual connections are not a substitute for a measured training budget.

Common Mistakes

• Do not add branches with mismatched spatial or channel dimensions.
• Do not equate frozen gradients with frozen BatchNorm running statistics.
• Do not omit buffers from a checkpoint intended for inference.

Read next

• Convolution output geometry and receptive fields
• Training and validation modes: measure the model you will serve
• Checkpoint recovery: save optimizer state and the run boundary
• Gradient accumulation and effective batch accounting
• Project: localize receipt defects with a residual CNN

Continue the workflow: Monte Carlo dropout, evaluation mode and batch-normalization boundaries.