Create upsampling layer.
(upsampling_mode: UpsamplingMode, scale_factor: int, dim_in: int)
| 30 | |
| 31 | |
| 32 | def upsampling_layer(upsampling_mode: UpsamplingMode, scale_factor: int, dim_in: int) -> nn.Module: |
| 33 | """Create upsampling layer.""" |
| 34 | if upsampling_mode == "transposed_conv": |
| 35 | return nn.ConvTranspose2d( |
| 36 | in_channels=dim_in, |
| 37 | out_channels=dim_in, |
| 38 | kernel_size=scale_factor, |
| 39 | stride=scale_factor, |
| 40 | padding=0, |
| 41 | bias=False, |
| 42 | ) |
| 43 | elif upsampling_mode in ("nearest", "bilinear"): |
| 44 | return nn.Upsample(scale_factor=scale_factor, mode=upsampling_mode) |
| 45 | else: |
| 46 | raise ValueError(f"Invalid upsampling mode {upsampling_mode}.") |
| 47 | |
| 48 | |
| 49 | class ResidualBlock(nn.Module): |