Spaces:
Runtime error
Runtime error
| from typing import Optional, Union, Callable | |
| from torch.nn import Conv2d, Module, Sequential, ConvTranspose2d | |
| from tha3.module.module_factory import ModuleFactory | |
| from tha3.nn.nonlinearity_factory import resolve_nonlinearity_factory | |
| from tha3.nn.normalization import NormalizationLayerFactory | |
| from tha3.nn.util import wrap_conv_or_linear_module, BlockArgs | |
| def create_conv7(in_channels: int, out_channels: int, | |
| bias: bool = False, | |
| initialization_method: Union[str, Callable[[Module], Module]] = 'he', | |
| use_spectral_norm: bool = False) -> Module: | |
| return wrap_conv_or_linear_module( | |
| Conv2d(in_channels, out_channels, kernel_size=7, stride=1, padding=3, bias=bias), | |
| initialization_method, | |
| use_spectral_norm) | |
| def create_conv7_from_block_args(in_channels: int, | |
| out_channels: int, | |
| bias: bool = False, | |
| block_args: Optional[BlockArgs] = None) -> Module: | |
| if block_args is None: | |
| block_args = BlockArgs() | |
| return create_conv7( | |
| in_channels, out_channels, bias, | |
| block_args.initialization_method, | |
| block_args.use_spectral_norm) | |
| def create_conv3(in_channels: int, | |
| out_channels: int, | |
| bias: bool = False, | |
| initialization_method: Union[str, Callable[[Module], Module]] = 'he', | |
| use_spectral_norm: bool = False) -> Module: | |
| return wrap_conv_or_linear_module( | |
| Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=bias), | |
| initialization_method, | |
| use_spectral_norm) | |
| def create_conv3_from_block_args(in_channels: int, out_channels: int, | |
| bias: bool = False, | |
| block_args: Optional[BlockArgs] = None): | |
| if block_args is None: | |
| block_args = BlockArgs() | |
| return create_conv3(in_channels, out_channels, bias, | |
| block_args.initialization_method, | |
| block_args.use_spectral_norm) | |
| def create_conv1(in_channels: int, out_channels: int, | |
| initialization_method: Union[str, Callable[[Module], Module]] = 'he', | |
| bias: bool = False, | |
| use_spectral_norm: bool = False) -> Module: | |
| return wrap_conv_or_linear_module( | |
| Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=bias), | |
| initialization_method, | |
| use_spectral_norm) | |
| def create_conv1_from_block_args(in_channels: int, | |
| out_channels: int, | |
| bias: bool = False, | |
| block_args: Optional[BlockArgs] = None) -> Module: | |
| if block_args is None: | |
| block_args = BlockArgs() | |
| return create_conv1( | |
| in_channels=in_channels, | |
| out_channels=out_channels, | |
| initialization_method=block_args.initialization_method, | |
| bias=bias, | |
| use_spectral_norm=block_args.use_spectral_norm) | |
| def create_conv7_block(in_channels: int, out_channels: int, | |
| initialization_method: Union[str, Callable[[Module], Module]] = 'he', | |
| nonlinearity_factory: Optional[ModuleFactory] = None, | |
| normalization_layer_factory: Optional[NormalizationLayerFactory] = None, | |
| use_spectral_norm: bool = False) -> Module: | |
| nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) | |
| return Sequential( | |
| create_conv7(in_channels, out_channels, | |
| bias=False, initialization_method=initialization_method, use_spectral_norm=use_spectral_norm), | |
| NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True), | |
| resolve_nonlinearity_factory(nonlinearity_factory).create()) | |
| def create_conv7_block_from_block_args( | |
| in_channels: int, out_channels: int, | |
| block_args: Optional[BlockArgs] = None) -> Module: | |
| if block_args is None: | |
| block_args = BlockArgs() | |
| return create_conv7_block(in_channels, out_channels, | |
| block_args.initialization_method, | |
| block_args.nonlinearity_factory, | |
| block_args.normalization_layer_factory, | |
| block_args.use_spectral_norm) | |
| def create_conv3_block(in_channels: int, out_channels: int, | |
| initialization_method: Union[str, Callable[[Module], Module]] = 'he', | |
| nonlinearity_factory: Optional[ModuleFactory] = None, | |
| normalization_layer_factory: Optional[NormalizationLayerFactory] = None, | |
| use_spectral_norm: bool = False) -> Module: | |
| nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) | |
| return Sequential( | |
| create_conv3(in_channels, out_channels, | |
| bias=False, initialization_method=initialization_method, use_spectral_norm=use_spectral_norm), | |
| NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True), | |
| resolve_nonlinearity_factory(nonlinearity_factory).create()) | |
| def create_conv3_block_from_block_args( | |
| in_channels: int, out_channels: int, block_args: Optional[BlockArgs] = None): | |
| if block_args is None: | |
| block_args = BlockArgs() | |
| return create_conv3_block(in_channels, out_channels, | |
| block_args.initialization_method, | |
| block_args.nonlinearity_factory, | |
| block_args.normalization_layer_factory, | |
| block_args.use_spectral_norm) | |
| def create_downsample_block(in_channels: int, out_channels: int, | |
| is_output_1x1: bool = False, | |
| initialization_method: Union[str, Callable[[Module], Module]] = 'he', | |
| nonlinearity_factory: Optional[ModuleFactory] = None, | |
| normalization_layer_factory: Optional[NormalizationLayerFactory] = None, | |
| use_spectral_norm: bool = False) -> Module: | |
| if is_output_1x1: | |
| return Sequential( | |
| wrap_conv_or_linear_module( | |
| Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False), | |
| initialization_method, | |
| use_spectral_norm), | |
| resolve_nonlinearity_factory(nonlinearity_factory).create()) | |
| else: | |
| return Sequential( | |
| wrap_conv_or_linear_module( | |
| Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False), | |
| initialization_method, | |
| use_spectral_norm), | |
| NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True), | |
| resolve_nonlinearity_factory(nonlinearity_factory).create()) | |
| def create_downsample_block_from_block_args(in_channels: int, out_channels: int, | |
| is_output_1x1: bool = False, | |
| block_args: Optional[BlockArgs] = None): | |
| if block_args is None: | |
| block_args = BlockArgs() | |
| return create_downsample_block( | |
| in_channels, out_channels, | |
| is_output_1x1, | |
| block_args.initialization_method, | |
| block_args.nonlinearity_factory, | |
| block_args.normalization_layer_factory, | |
| block_args.use_spectral_norm) | |
| def create_upsample_block(in_channels: int, | |
| out_channels: int, | |
| initialization_method: Union[str, Callable[[Module], Module]] = 'he', | |
| nonlinearity_factory: Optional[ModuleFactory] = None, | |
| normalization_layer_factory: Optional[NormalizationLayerFactory] = None, | |
| use_spectral_norm: bool = False) -> Module: | |
| nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) | |
| return Sequential( | |
| wrap_conv_or_linear_module( | |
| ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False), | |
| initialization_method, | |
| use_spectral_norm), | |
| NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True), | |
| resolve_nonlinearity_factory(nonlinearity_factory).create()) | |
| def create_upsample_block_from_block_args(in_channels: int, | |
| out_channels: int, | |
| block_args: Optional[BlockArgs] = None) -> Module: | |
| if block_args is None: | |
| block_args = BlockArgs() | |
| return create_upsample_block(in_channels, out_channels, | |
| block_args.initialization_method, | |
| block_args.nonlinearity_factory, | |
| block_args.normalization_layer_factory, | |
| block_args.use_spectral_norm) | |