xnn.cnn.models.steerable.SteerableBatchNorm#

class xnn.cnn.models.steerable.SteerableBatchNorm(fields, eps=1e-05, momentum=0.1, affine=True)[source]#

Bases: Module

Equivariant batch normalization (supplement Sec. 1.2, eq 17).

Scalar fields get ordinary batch normalization (mean and variance over the batch and the grid). A non-scalar field is only rescaled, by the inverse root of the batch- and grid-averaged squared norm of its components, which commutes with every rotation. Running statistics are kept for evaluation; the affine transform has a gain per field and a bias for scalar fields only.

Parameters:

  • fields (sequence of int) – Multiplicities of the normalized stack.

  • eps (float, optional) – Added to the variances, by default 1e-5.

  • momentum (float, optional) – Running-statistics momentum, by default 0.1.

  • affine (bool, optional) – Learn the gains and scalar biases, by default True.

forward(x)[source]#

Normalize (B, field_dim(fields), X, Y, Z) fields.

Parameters:

x (Tensor)

Return type:

Tensor