espnet2.asr.state_spaces.components.Normalization

source

class espnet2.asr.state_spaces.components.Normalization(d, transposed=False, _name_='layer', **kwargs)

Bases: Module

Normalization module for various normalization techniques.

This class implements several normalization techniques, including LayerNorm, InstanceNorm, BatchNorm, GroupNorm, and Identity normalization. The choice of normalization method is determined by the _name_ parameter during initialization.

transposed

Indicates if the normalization is applied on transposed input (length dimension is -1 or -2).

• Type: bool

_name_

The type of normalization to apply. Options include: ‘layer’, ‘instance’, ‘batch’, ‘group’, ‘none’.

• Type: str

norm

The normalization layer based on the selected method.

• Type: nn.Module

• Parameters:

  • d (int) – The number of features (channels) for normalization.
  • transposed (bool , optional) – Whether to apply normalization to transposed input. Default is False.
  • _name_ (str , optional) – The name of the normalization type. Default is ‘layer’.
  • **kwargs – Additional keyword arguments for specific normalization methods.

• Raises:NotImplementedError – If an unsupported normalization type is specified in _name_.

######### Examples

>>> layer_norm = Normalization(d=64, _name_='layer')
>>> instance_norm = Normalization(d=64, _name_='instance')
>>> batch_norm = Normalization(d=64, _name_='batch')
>>> group_norm = Normalization(d=64, _name_='group')
>>> identity_norm = Normalization(d=64, _name_='none')

NOTE

The transposed option is particularly useful for working with sequences or multi-dimensional data where the last dimensions represent time or length.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Normalization module for different normalization techniques.

This module implements various normalization techniques such as LayerNorm, InstanceNorm, BatchNorm, GroupNorm, and an identity operation. It allows for normalization across different dimensions and can be configured to use transposed dimensions.

transposed

Indicates whether the normalization should be applied on transposed dimensions.

• Type: bool

_name_

Specifies the type of normalization to apply.

• Type: str

norm

The actual normalization layer based on the specified normalization type.

• Type: nn.Module

• Parameters:

  • d (int) – The number of features in the input tensor.
  • transposed (bool) – If True, applies normalization on the last dimension. Default is False.
  • _name_ (str) – The type of normalization to apply. Options are: “layer”, “instance”, “batch”, “group”, or “none”. Default is “layer”.
  • **kwargs – Additional keyword arguments to pass to the normalization layer.

• Returns: The normalized tensor.

• Return type: Tensor

######### Examples

Using LayerNorm

layer_norm = Normalization(d=64, _name_=”layer”) output = layer_norm(input_tensor)

Using BatchNorm

batch_norm = Normalization(d=64, _name_=”batch”) output = batch_norm(input_tensor)

NOTE

The input tensor is expected to be of shape (B, D, L) or (B, L, D) depending on the transposed configuration.

• Raises:NotImplementedError – If an unsupported normalization type is specified.

step(x, **kwargs)

Apply the normalization step to the input tensor.

This method applies the normalization operation to the input tensor x based on the specified normalization type (layer, instance, batch, group, or none). It handles the transposed state of the tensor as needed and reshapes the input to apply normalization correctly. The result is returned in the same shape as the input tensor.

• Parameters:
  • x (Tensor) – The input tensor to be normalized.
  • **kwargs – Additional keyword arguments to be passed to the normalization layer.
• Returns: The normalized tensor with the same shape as the input.
• Return type: Tensor
• Raises:AssertionError – If the normalization type (name) is not one of the supported types (“layer”, “instance”, “batch”, “group”, “none”).

######### Examples

>>> norm_layer = Normalization(d=64, _name_='layer')
>>> input_tensor = torch.randn(32, 64)  # Batch of 32 samples, 64 features
>>> output_tensor = norm_layer.step(input_tensor)
>>> output_tensor.shape
torch.Size([32, 64])

>>> group_norm_layer = Normalization(d=64, _name_='group')
>>> input_tensor = torch.randn(32, 64)  # Batch of 32 samples, 64 features
>>> output_tensor = group_norm_layer.step(input_tensor)
>>> output_tensor.shape
torch.Size([32, 64])