Create normalization layer.
(num_features: int, norm_type: NormLayerName, num_groups: int = 8)
| 16 | |
| 17 | |
| 18 | def norm_layer_2d(num_features: int, norm_type: NormLayerName, num_groups: int = 8) -> nn.Module: |
| 19 | """Create normalization layer.""" |
| 20 | if norm_type == "noop": |
| 21 | return nn.Identity() |
| 22 | elif norm_type == "batch_norm": |
| 23 | return nn.BatchNorm2d(num_features=num_features) |
| 24 | elif norm_type == "group_norm": |
| 25 | return nn.GroupNorm(num_channels=num_features, num_groups=num_groups) |
| 26 | elif norm_type == "instance_norm": |
| 27 | return nn.InstanceNorm2d(num_features=num_features) |
| 28 | else: |
| 29 | raise ValueError(f"Invalid normalization layer type: {norm_type}") |
| 30 | |
| 31 | |
| 32 | def upsampling_layer(upsampling_mode: UpsamplingMode, scale_factor: int, dim_in: int) -> nn.Module: |
no outgoing calls
no test coverage detected