torchlayers.upsample module

class torchlayers.upsample.ConvPixelShuffle(in_channels, out_channels, upscale_factor: int = 2, kernel_size: int = 3, stride: int = 1, padding: Union[Tuple[int, int], int, str] = 'same', dilation: int = 1, groups: int = 1, bias: bool = True, padding_mode: str = 'zeros', initializer: Callable[[torch.Tensor], torch.Tensor] = None)

Two dimensional convolution with ICNR initialization followed by PixelShuffle.

Increases height and width of input tensor by scale, acts like learnable upsampling. Due to ICNR weight initialization of convolution it has similar starting point to nearest neighbour upsampling.

kernel_size got a default value of 3, upscale_factor got a default value of 2

Note

Currently only 4D input is allowed ([batch, channels, height, width]), due to torch.nn.PixelShuffle not supporting 1D or 3D versions. See [this PyTorch PR](https://github.com/pytorch/pytorch/pull/6340/files) for example of dimension-agnostic implementation.

Parameters

• in_channels (int) – Number of channels in the input image

• out_channels (int) – Number of channels produced after PixelShuffle

• upscale_factor (int, optional) – Factor to increase spatial resolution by. Default: 2

• kernel_size (int or tuple, optional) – Size of the convolving kernel. Default: 3

• stride (int or tuple, optional) – Stride of the convolution. Default: 1

• padding (int or tuple, optional) – Zero-padding added to both sides of the input. Default: 0

• padding_mode (string, optional) – Accepted values zeros and circular Default: zeros

• dilation (int or tuple, optional) – Spacing between kernel elements. Default: 1

• groups (int, optional) – Number of blocked connections from input channels to output channels. Default: 1

• bias (bool, optional) – If True, adds a learnable bias to the output. Default: True

• initializer (typing.Callable[[torch.Tensor,], torch.Tensor], optional) – Initializer for ICNR initialization, can be a function from torch.nn.init. Gets and returns tensor after initialization. Default: torch.nn.init.kaiming_normal_

forward(inputs)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

icnr_initialization(tensor)

ICNR initializer for checkerboard artifact free sub pixel convolution.

Originally presented in Checkerboard artifact free sub-pixel convolution: A note on sub-pixel convolution, resize convolution and convolution resize Initializes convolutional layer prior to torch.nn.PixelShuffle. Weights are initialized according to initializer passed to to __init__.

Parameters

tensor (torch.Tensor) – Tensor to be initialized using ICNR init.

Returns

Tensor initialized using ICNR.

Return type

torch.Tensor

post_build()

Initialize weights after layer was built.