Batch Normalization (BatchNorm)#

API#

class warp_nn.modules.layers.BatchNorm(
num_features: int,
*,
eps: float = 1e-05,
momentum: float = 0.1,
affine: bool = True,
track_running_stats: bool = True,
initialize_parameters: bool = True,
requires_grad: bool = True,
)[source]#

Bases: Module

Apply Batch Normalization over the channels of the input.

\[\text{BatchNorm}(x) = \frac{x - \text{E}[x]}{\sqrt{\text{Var}[x] + \epsilon}} \, \gamma + \beta\]

The mean and the (biased) variance are computed separately for each channel over the batch (and all the remaining dimensions).

In training mode (see train()), the statistics are computed from the input data and, if track_running_stats is true, their running estimates are updated (using the unbiased variance) as follows:

\[\hat{x}_{\text{running}} = (1 - \text{momentum}) \, \hat{x}_{\text{running}} + \text{momentum} \, x_t\]

In evaluation mode, the running estimates are used for normalization if track_running_stats is true. Otherwise, the statistics are computed from the input data.

Learnable parameters:

Name

Shape

Description

\(\gamma\)

weight

(num_features,)

Per-channel weights. Only if affine is true

\(\beta\)

bias

(num_features,)

Per-channel bias. Only if affine is true

The weights are initialized to ones and the bias to zeros.

Buffers (only if track_running_stats is true):

Name

Shape

Description

running_mean

(num_features,)

Running mean, initialized to zeros

running_var

(num_features,)

Running (unbiased) variance, initialized to ones

Parameters:

  • num_features – The number of channels expected in the input.

  • eps – The value added to the variance for numerical stability.

  • momentum – The weight of the current statistics in the running estimates update.

  • affine – Whether to include learnable per-channel affine parameters.

  • track_running_stats – Whether to track running estimates of the mean and variance.

  • initialize_parameters – Whether to initialize the parameters with their default/initial values. If false, the parameters are left as uninitialized memory, and reading them (e.g. during a forward pass) before loading their values (e.g. from a state dictionary) is undefined behavior. The buffers are always initialized.

  • requires_grad – Whether the parameters and the cached output arrays of the module require gradients.

__call__(
input: array,
) → array[source]#

Forward pass of the module.

Parameters:

input – The input array, with shape (batch_size, num_features, *).

Returns:

The output array, with same shape as the input array.

Raises:

ValueError – If the input array’s shape is not supported, or if there is only one value per channel in training mode.