Implementation of ICNR with PyTorch

ICNR.md

ICNR

ICNR is an initialization method for sub-pixel convolution.

References

Papar

• Checkerboard artifact free sub-pixel convolution: A note on sub-pixel convolution, resize convolution and convolution resize

Github or Gist

• ICNR inizialization by catta202000 · Pull Request #5429 · pytorch/pytorch

• pytorch/PixelShuffle.cpp at 517c7c98610402e2746586c78987c64c28e024aa · pytorch/pytorch

• kostyaev/ICNR: Convolution NN resize initialization for subpixel convolutions

• Sub-Pixel Conv with ICNR

Requirements

Python 3.6 or later
PyTorch 1.1.0 or later

Usage

For example, input Tensor with the sizes (64, 8, 32, 32) would be output as (64, 8, 64, 64) if you set upscale_factor to 2.

import torch
import torch.nn as nn

from icnr import ICNR


upscale_factor = 2
input = torch.randn(64, 3, 32, 32)
conv = nn.Conv2d(3, 3 * (upscale_factor ** 2), 3, 1, 1, bias=0)
pixelshuffle = nn.PixelShuffle(upscale_factor)
weight = ICNR(conv.weight, initializer=nn.init.kaiming_normal_,
              upscale_factor=upscale_factor)
conv.weight.data.copy_(weight)   # initialize conv.weight
output = conv(input)  # (64, 12, 32, 32)
output = pixelshuffle(output)  # (64, 3, 64, 64)

Visualization

See visualization.ipynb

icnr.py

import torch


def ICNR(tensor, initializer, upscale_factor=2, *args, **kwargs):
    "tensor: the 2-dimensional Tensor or more"
    upscale_factor_squared = upscale_factor * upscale_factor
    assert tensor.shape[0] % upscale_factor_squared == 0, \
        ("The size of the first dimension: "
         f"tensor.shape[0] = {tensor.shape[0]}"
         " is not divisible by square of upscale_factor: "
         f"upscale_factor = {upscale_factor}")
    sub_kernel = torch.empty(tensor.shape[0] // upscale_factor_squared,
                             *tensor.shape[1:])
    sub_kernel = initializer(sub_kernel, *args, **kwargs)
    return sub_kernel.repeat_interleave(upscale_factor_squared, dim=0)

visualization.ipynb