Args: spatial_dims: number of spatial dimensions of the input image. in_channels: number of channels of the input image. out_channels: optional number of channels of the output image. scale_factor: multiplier for spatial size. Defaults to 2.
(
self,
spatial_dims: int,
in_channels: int | None,
out_channels: int | None = None,
scale_factor: int = 2,
conv_block: nn.Module | str | None = "default",
apply_pad_pool: bool = True,
bias: bool = True,
)
| 215 | """ |
| 216 | |
| 217 | def __init__( |
| 218 | self, |
| 219 | spatial_dims: int, |
| 220 | in_channels: int | None, |
| 221 | out_channels: int | None = None, |
| 222 | scale_factor: int = 2, |
| 223 | conv_block: nn.Module | str | None = "default", |
| 224 | apply_pad_pool: bool = True, |
| 225 | bias: bool = True, |
| 226 | ) -> None: |
| 227 | """ |
| 228 | Args: |
| 229 | spatial_dims: number of spatial dimensions of the input image. |
| 230 | in_channels: number of channels of the input image. |
| 231 | out_channels: optional number of channels of the output image. |
| 232 | scale_factor: multiplier for spatial size. Defaults to 2. |
| 233 | conv_block: a conv block to extract feature maps before upsampling. Defaults to None. |
| 234 | |
| 235 | - When ``conv_block`` is ``"default"``, one reserved conv layer will be utilized. |
| 236 | - When ``conv_block`` is an ``nn.module``, |
| 237 | please ensure the output number of channels is divisible ``(scale_factor ** dimensions)``. |
| 238 | |
| 239 | apply_pad_pool: if True the upsampled tensor is padded then average pooling is applied with a kernel the |
| 240 | size of `scale_factor` with a stride of 1. This implements the nearest neighbour resize convolution |
| 241 | component of subpixel convolutions described in Aitken et al. |
| 242 | bias: whether to have a bias term in the default conv_block. Defaults to True. |
| 243 | |
| 244 | """ |
| 245 | super().__init__() |
| 246 | |
| 247 | if scale_factor <= 0: |
| 248 | raise ValueError(f"The `scale_factor` multiplier must be an integer greater than 0, got {scale_factor}.") |
| 249 | |
| 250 | self.dimensions = spatial_dims |
| 251 | self.scale_factor = scale_factor |
| 252 | |
| 253 | if conv_block == "default": |
| 254 | out_channels = out_channels or in_channels |
| 255 | if not out_channels: |
| 256 | raise ValueError("in_channels need to be specified.") |
| 257 | conv_out_channels = out_channels * (scale_factor**self.dimensions) |
| 258 | self.conv_block = Conv[Conv.CONV, self.dimensions]( |
| 259 | in_channels=in_channels, out_channels=conv_out_channels, kernel_size=3, stride=1, padding=1, bias=bias |
| 260 | ) |
| 261 | |
| 262 | icnr_init(self.conv_block, self.scale_factor) |
| 263 | elif conv_block is None: |
| 264 | self.conv_block = nn.Identity() |
| 265 | else: |
| 266 | self.conv_block = conv_block |
| 267 | |
| 268 | self.pad_pool: nn.Module = nn.Identity() |
| 269 | |
| 270 | if apply_pad_pool: |
| 271 | pool_type = Pool[Pool.AVG, self.dimensions] |
| 272 | pad_type = Pad[Pad.CONSTANTPAD, self.dimensions] |
| 273 | |
| 274 | self.pad_pool = nn.Sequential( |