espnet2.enh.layers.tcn.choose_norm

source

espnet2.enh.layers.tcn.choose_norm(norm_type, channel_size, shape='BDT')

Choose the appropriate normalization layer based on the norm type.

The input of normalization will be (M, C, K), where M is the batch size. C is the channel size and K is the sequence length.

• Parameters:
  • norm_type (str) – The type of normalization to use. Options include: “gLN” for Global Layer Normalization, “cLN” for Channel-wise Layer Normalization, “BN” for Batch Normalization, “GN” for Group Normalization.
  • channel_size (int) – The number of channels to normalize.
  • shape (str , optional) – The shape of the input tensor, either “BDT” (Batch, Depth, Time) or “BTD” (Batch, Time, Depth). Defaults to “BDT”.
• Returns: The selected normalization layer.
• Return type: nn.Module
• Raises:ValueError – If an unsupported normalization type is provided.

Examples

>>> norm_layer = choose_norm("gLN", 64)
>>> input_tensor = torch.randn(10, 64, 50)  # (M, C, K)
>>> output_tensor = norm_layer(input_tensor)