spd_learn.modules.BatchReNorm

class spd_learn.modules.BatchReNorm(num_features, momentum=0.9, rebias=True, renorm=False, device=None, dtype=None)

Bases: Module

Batch Re-Normalization.

This class implements Batch Re-Normalization, which is a variant of batch normalization that can be used in recurrent neural networks or as a Euclidean baseline for comparison with Riemannian batch normalization methods.

Parameters:

• num_features (int) – The number of features in the input.

• momentum (float, default=0.9) – The momentum for the running mean and variance.

• rebias (bool, default=True) – If True, the layer has a learnable bias parameter.

• renorm (bool, default=False) – If True, the layer uses re-normalization.

forward(input)

Forward pass of the Batch Re-Normalization layer.

Parameters:

input (torch.Tensor) – Input tensor of shape (batch_size, [h,] d), where d is the number of features.

Returns:

Normalized tensor of the same shape as the input.

Return type:

torch.Tensor

reset_parameters() → None

reset_running_stats() → None