BatchNorm2d — PyTorch 2.7 documentation (original) (raw)

class torch.nn.BatchNorm2d(num_features, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True, device=None, dtype=None)[source][source]

Applies Batch Normalization over a 4D input.

4D is a mini-batch of 2D inputs with additional channel dimension. Method described in the paperBatch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift .

y=x−E[x]Var[x]+ϵ∗γ+βy = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta

The mean and standard-deviation are calculated per-dimension over the mini-batches and γ\gamma and β\beta are learnable parameter vectors of size C (where C is the input size). By default, the elements of γ\gamma are set to 1 and the elements of β\beta are set to 0. At train time in the forward pass, the standard-deviation is calculated via the biased estimator, equivalent totorch.var(input, unbiased=False). However, the value stored in the moving average of the standard-deviation is calculated via the unbiased estimator, equivalent totorch.var(input, unbiased=True).

Also by default, during training this layer keeps running estimates of its computed mean and variance, which are then used for normalization during evaluation. The running estimates are kept with a default momentumof 0.1.

If track_running_stats is set to False, this layer then does not keep running estimates, and batch statistics are instead used during evaluation time as well.

Note

This momentum argument is different from one used in optimizer classes and the conventional notion of momentum. Mathematically, the update rule for running statistics here isx^new=(1−momentum)×x^+momentum×xt\hat{x}_\text{new} = (1 - \text{momentum}) \times \hat{x} + \text{momentum} \times x_t, where x^\hat{x} is the estimated statistic and xtx_t is the new observed value.

Because the Batch Normalization is done over the C dimension, computing statistics on (N, H, W) slices, it’s common terminology to call this Spatial Batch Normalization.

Parameters

Shape:

Examples:

With Learnable Parameters

m = nn.BatchNorm2d(100)

Without Learnable Parameters

m = nn.BatchNorm2d(100, affine=False) input = torch.randn(20, 100, 35, 45) output = m(input)