ICNR: Sub-Pixel Conv使用時のcheckerboard artifactを防ぐ初期化

f:id:ejinote:20190729202002p:plain

ニューラルネットワークで画像を拡大(アップサンプリング)する際、Transposed ConvolutionやSub-Pixel Convolutionという手法を使います。

しかし、これらの手法を用いた場合、拡大した画像に格子状の模様(checkerboard artifact)が発生することが知られています。

以前のLTのスライドも併せてご覧いただけると良いかと思います。

checkerboard artifactの主要な原因として挙げられているのは2つです。

1つ目はfilter overlapです。filter overlapとは、kernel_sizeやstrideの設定によっては出力に寄与する画素数が異なってしまう(フィルタが重なってしまう)現象のことです。特にTransposed Convを利用した場合に発生します。

Sub-Pixel Convolutionの場合、原理上出力に寄与する画素数が異なることはありませんが、kernelの初期化によってcheckerboard artifactが発生することが指摘されています。これが2つ目の原因です。

Sub-Pixel Convolution使用時、初期化によるcheckerboard artifactを防ぐ方法として、ICNRという初期化方法が提案されています *1。 ICNRを利用することで超解像タスクにおいて(向上幅はわずかですが)綺麗な画像が生成されることが報告されています。

今回はPyTorchでICNRを実装し、ICNRの動きについて簡単な可視化を行いました。

ICNRの挙動

ICNRの基本的な発想は「Sub-Pixel ConvolutionでNearest Neighborを再現する」です。

Nearest Neighborで拡大した画像にはcheckerboard artifactは出ないので、それを初期値として学習を始めることでcheckerboard artifactを抑えて学習をすすめることができると期待されます。

Sub-Pixel ConvolutionはConv2dとPixelShuffleからなります。Conv2dの初期化を下図のように行うことで、PixelShuffleした後の画像がNearest Neighborと同様になります。

f:id:ejinote:20190728165735p:plain

直感的ではないと思いますので、これを実験によって確かめてみます。

実験

設定

32x32のRGB(3チャネル)の画像を nn.Conv2d で畳み込み、12チャネルに増やした後 nn.PixelShuffle で並べ替えて64x64のRGB画像を作るという問題設定にします。

f:id:ejinote:20190728154213p:plain — 元画像

元画像は以下のようなコードでPyTorchで読み込み可能な形式に変換します。

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image

import torch
import torch.nn.functional as F
from torch import nn

image = Image.open('dog.jpg')
image = np.asarray(image, np.float32) / 255

x = image[np.newaxis, :, :, :].transpose(0, 3, 1, 2)
x = torch.from_numpy(x)

通常の初期化を行った場合

out = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=3, stride=1, padding=1)(x)
out = nn.PixelShuffle(2)(out)

f:id:ejinote:20190728155037p:plain — 通常のConv2d→PixelShuffle

ご覧のとおり、それはそれはキレイなcheckerboard artifactが出ていることが分かるかと思います。もちろんこれは全く学習を行っていない状態なので、学習を進めていけば徐々に消えてはいきます。

ICNRで初期化した場合

ICNRの実装例は以下のとおりです。2倍に拡大する場合は出力チャネル数を1/4で初期化したあと、チャネル方向に4倍に引き伸ばしたものをConv2dの初期値として使います。

通常であればreshapeとpermuteを組み合わせて行列演算だけで書くのですが、初期化処理は何度も呼ばれるものではないので可読性重視でfor文を使っています。

def ICNR(tensor, scale_factor=2, initializer=nn.init.kaiming_normal_):
    OUT, IN, H, W = tensor.shape
    sub = torch.zeros(OUT//scale_factor**2, IN, H, W)
    sub = initializer(sub)
    
    kernel = torch.zeros_like(tensor)
    for i in range(OUT):
        kernel[i] = sub[i//scale_factor**2]
        
    return kernel

使用例です。conv.weight.data.copy_を使用してICNRで生成した初期値を与えています。

conv = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=3, stride=1, padding=1)
kernel = ICNR(conv.weight)
conv.weight.data.copy_(kernel)

out = conv(x)
out = nn.PixelShuffle(2)(out)