espnet2.gan_codec.shared.encoder.seanet.apply_parametrization_norm

source

espnet2.gan_codec.shared.encoder.seanet.apply_parametrization_norm(module: Module, norm: str = 'none') → Module

Applies a specified normalization technique to a given module.

This function checks the normalization method specified in the norm argument and applies the corresponding normalization to the input module. If the normalization type is not recognized, the original module is returned unchanged.

• Parameters:
  • module (nn.Module) – The neural network module to which the normalization will be applied.
  • norm (str , optional) – The type of normalization to apply. It can be one of the following: ‘none’, ‘weight_norm’, or ‘spectral_norm’. Defaults to ‘none’.
• Returns: The module with the specified normalization applied.
• Return type: nn.Module
• Raises:AssertionError – If the specified normalization type is not in the predefined set of allowed normalizations.

Examples

>>> import torch.nn as nn
>>> conv_layer = nn.Conv2d(3, 64, kernel_size=3)
>>> normalized_layer = apply_parametrization_norm(conv_layer, 'weight_norm')
>>> assert isinstance(normalized_layer, nn.utils.weight_norm)