Group Normalization (GroupNorm)#

API#

class warp_nn.modules.layers.GroupNorm(
num_groups: int,
num_channels: int,
*,
eps: float = 1e-05,
affine: bool = True,
initialize_parameters: bool = True,
requires_grad: bool = True,
)[source]#

Bases: Module

Apply Group Normalization over the channels of the input.

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

The channels are separated into num_groups groups, each containing num_channels / num_groups consecutive channels. The mean and the (biased) variance are computed separately for each sample over the channels (and all the remaining dimensions) of each group.

Learnable parameters:

Name

Shape

Description

\(\gamma\)

weight

(num_channels,)

Per-channel weights. Only if affine is true

\(\beta\)

bias

(num_channels,)

Per-channel bias. Only if affine is true

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

Parameters:

  • num_groups – The number of groups to separate the channels into.

  • num_channels – The number of channels expected in the input. It must be divisible by num_groups.

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

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

  • 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.

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

Raises:

ValueError – If num_channels is not divisible by num_groups.

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

Forward pass of the module.

Parameters:

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

Returns:

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

Raises:

ValueError – If the input array’s shape is not supported.