| from typing import Optional, Callable, Union | |
| from torch.nn import Module | |
| from tha3.module.module_factory import ModuleFactory | |
| from tha3.nn.init_function import create_init_function | |
| from tha3.nn.nonlinearity_factory import resolve_nonlinearity_factory | |
| from tha3.nn.normalization import NormalizationLayerFactory | |
| from tha3.nn.spectral_norm import apply_spectral_norm | |
| def wrap_conv_or_linear_module(module: Module, | |
| initialization_method: Union[str, Callable[[Module], Module]], | |
| use_spectral_norm: bool): | |
| if isinstance(initialization_method, str): | |
| init = create_init_function(initialization_method) | |
| else: | |
| init = initialization_method | |
| return apply_spectral_norm(init(module), use_spectral_norm) | |
| class BlockArgs: | |
| def __init__(self, | |
| initialization_method: Union[str, Callable[[Module], Module]] = 'he', | |
| use_spectral_norm: bool = False, | |
| normalization_layer_factory: Optional[NormalizationLayerFactory] = None, | |
| nonlinearity_factory: Optional[ModuleFactory] = None): | |
| self.nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) | |
| self.normalization_layer_factory = normalization_layer_factory | |
| self.use_spectral_norm = use_spectral_norm | |
| self.initialization_method = initialization_method | |
| def wrap_module(self, module: Module) -> Module: | |
| return wrap_conv_or_linear_module(module, self.get_init_func(), self.use_spectral_norm) | |
| def get_init_func(self) -> Callable[[Module], Module]: | |
| if isinstance(self.initialization_method, str): | |
| return create_init_function(self.initialization_method) | |
| else: | |
| return self.initialization_method | |