Source code for torchlayers.upsample

import typing

import torch

from . import convolution

[docs]class ConvPixelShuffle(torch.nn.Module): """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]( 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_` """ def __init__( self, in_channels, out_channels, upscale_factor: int = 2, kernel_size: int = 3, stride: int = 1, padding: typing.Union[typing.Tuple[int, int], int, str] = "same", dilation: int = 1, groups: int = 1, bias: bool = True, padding_mode: str = "zeros", initializer: typing.Callable[[torch.Tensor,], torch.Tensor] = None, ): super().__init__() self.convolution = convolution.Conv( in_channels, out_channels * upscale_factor * upscale_factor, kernel_size, stride, padding, dilation, groups, bias, padding_mode, ) self.upsample = torch.nn.PixelShuffle(upscale_factor) if initializer is None: self.initializer = torch.nn.init.kaiming_normal_ else: self.initializer = initializer
[docs] def post_build(self): """Initialize weights after layer was built.""" self.icnr_initialization(
[docs] def icnr_initialization(self, 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 ------- torch.Tensor Tensor initialized using ICNR. """ if self.upsample.upscale_factor == 1: return self.initializer(tensor) new_shape = [int(tensor.shape[0] / (self.upsample.upscale_factor ** 2))] + list( tensor.shape[1:] ) subkernel = self.initializer(torch.zeros(new_shape)).transpose(0, 1) kernel = subkernel.reshape(subkernel.shape[0], subkernel.shape[1], -1).repeat( 1, 1, self.upsample.upscale_factor ** 2 ) return kernel.reshape([-1, tensor.shape[0]] + list(tensor.shape[2:])).transpose( 0, 1 )
[docs] def forward(self, inputs): return self.upsample(self.convolution(inputs))