ICNR initialization for 2D/3D kernels adapted from Aitken et al.,2017 , "Checkerboard artifact free sub-pixel convolution".
(conv, upsample_factor, init=nn.init.kaiming_normal_)
| 348 | |
| 349 | |
| 350 | def icnr_init(conv, upsample_factor, init=nn.init.kaiming_normal_): |
| 351 | """ |
| 352 | ICNR initialization for 2D/3D kernels adapted from Aitken et al.,2017 , "Checkerboard artifact free |
| 353 | sub-pixel convolution". |
| 354 | """ |
| 355 | out_channels, in_channels, *dims = conv.weight.shape |
| 356 | scale_factor = upsample_factor ** len(dims) |
| 357 | |
| 358 | oc2 = int(out_channels / scale_factor) |
| 359 | |
| 360 | kernel = torch.zeros([oc2, in_channels] + dims) |
| 361 | kernel = init(kernel) |
| 362 | kernel = kernel.transpose(0, 1) |
| 363 | kernel = kernel.reshape(oc2, in_channels, -1) |
| 364 | kernel = kernel.repeat(1, 1, scale_factor) |
| 365 | kernel = kernel.reshape([in_channels, out_channels] + dims) |
| 366 | kernel = kernel.transpose(0, 1) |
| 367 | conv.weight.data.copy_(kernel) |
| 368 | |
| 369 | |
| 370 | def pixelshuffle(x: torch.Tensor, spatial_dims: int, scale_factor: int) -> torch.Tensor: |