画像の超解像度化: ESPCN の pytorch 実装 / 学習

画像の超解像度化シリーズ第二弾です。 第一弾 では NN を使ったモデルの中では、もっとも初期 2015年に提案された SRCNN を実装しました。今回はそれから一年後 2016年に提案された ESPCN を実装して学習させてみたよ、という話です。

ESPCN は Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network という論文で提案された手法です。こちらも SRCNN 同様画像の特徴抽出機構として CNN を使っていますが若干の変更点があります。順に紹介していきます。

ESPCN 概要

SRCNN のモデルは入力と出力の画像サイズが同じようなモデルでした。そのため特定の画像を $r$ 倍の画像へ超解像度化したい場合には

1. 元の画像を古典的手法で $r$ 倍して拡大
2. 拡大された画像をネットワークへ入力

という流れをとっていました。一方で ESPCN ではロス関数は同じく MSE ですが入力画像は直接拡大されるという点が異なります。実際にネットワークを表したのが以下の図になります。

Figure1 ESPCN のネットワーク概要図 (提案論文より引用)

入力された $W \times H$ サイズの画像は3層のCNNにより $r^2$ のチャネルを持つ3次元テンソルへと変換されます(この時画像の高さと幅の大きさは変化させないように padding なりを用意します。) その後 PixelShuffle と呼ばれる機構によって $rW \times rH$ の大きさへと変換されます。

pytorch のネットワークで表現すると以下のようになります。

class ESPCN(AbstractNet):
    only_luminance = True
    input_upscale = False

    def __init__(self, upscale=2):
        super(ESPCN, self).__init__()
        self.upscale = upscale
        self.conv1 = nn.Conv2d(1, 64, kernel_size=5, stride=1, padding=2)
        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1)
        self.conv4 = nn.Conv2d(32, upscale ** 2, kernel_size=3, stride=1, padding=1)
        self.pixel_shuffle = nn.PixelShuffle(self.upscale)

    def _initialize_weights(self):
        weights_with_relu = [
            self.conv1.weight,
            self.conv2.weight,
            self.conv3.weight
        ]

        for w in weights_with_relu:
            nn.init.orthogonal_(w, nn.init.calculate_gain('relu'))

        nn.init.orthogonal_(self.conv4.weight)

    def forward(self, x):
        h = F.tanh(self.conv1(x))
        h = F.tanh(self.conv2(h))
        h = F.tanh(self.conv3(h))
        return self.pixel_shuffle(self.conv4(h))

ESPCN の肝: PixelShuffle

この最終層で用いられている PixelShuffle はこの論文で提案された新しい画像の拡大方法です。それまでは画像の拡大には DeConvolution を用いられることが多かったのですが、この演算は単に CNN を逆方向に行うだけなので画像に周期的ができるような現象が観測されることがありました。 Deconvolution の問題点は Deconvolution and Checkerboard Artifacts などが詳しいです。

PixelShuffle では拡大したい倍率を $r \in \mathbb{N}$ とした時 $r \times r$ のチャネルを直前の層で用意して、それらを1チャネルの上にタイル状に並べることで画像の拡大を行います。

例えば拡大率を 3 として元画像サイズが 1x1 だとします。このとき出力層の前の層では 9 チャネルを持つ 1x1 のテンソルを用意しておきます。これを $x \in \mathbb{R}^{1 \times 1 \times 9} $ と置きましょう。この時 PixelShuffle の出力 $y \in \mathbb{R}^{3 \times 3}$ は

$$ y_{i,j} = x_{1, 1, i + (j - 1) \times 3} $$

になります。特徴は出力がすべて元のテンソル成分によって決まっておりかつ元のテンソル要素が出力に寄与する部分がひとつしかない (i.e. overlap がない)という点です。これによって出力がゆがんでしまうことを防ぐことが出来ます。

その他変更点

その他の重要な変更点としてはネットワークに通す画像は RGB の 3チャネルの画像ではなく輝度情報のみの1チャネルの画像のみを用いるという点があります。これは人間が画像のクオリティを判定する際には輝度情報を重要視する、という部分から来ているようです。 他の2色の情報は古典的な手法で拡大し、最終的にマージして RGB に変換する、という処理を行います。

この処理は PIL.Image インスタンスに対してならば以下のように行なえます。

def get_luminance(img):
    x, _, _ = img.convert('YCbCr').split()
    return x

また PixelShuffle の構造上ネットワークの出力は元の画像の $r$ 倍でないといけない、という制約がある点も実装を行う上では注意する必要がある点です。要するに復元する画像は縮小された画像の $r$ 倍である必要があるため割り算できない大きさの画像は扱えないのです。そのため画像の大きさから拡大率で割り算できる最大の整数を返すような関数を定義しました。

def calculate_original_img_size(origin_size: int, upscale_factor: int) -> int:
    """
    元の画像サイズを縮小拡大したいときに元の画像をどの大きさに
    resize する必要があるかを返す関数

    例えば 202 px の画像を 1/3 に縮小することは出来ない(i.e. 3の倍数ではない)ので
    事前に 201 px に縮小しておく必要がありこの関数はその計算を行う
    すなわち

    calculate_original_img_size(202, 3) -> 201

    となる

    Args:
        origin_size:
        upscale_factor:

    Returns:
    """
    return origin_size - (origin_size % upscale_factor)

実験

今回も数値実験を行ってみます。実装のコードは以下からアクセス出来ます。

github.com

実験条件は以下のとおりです

環境

• Intel(R) Core(TM) i7-7700K CPU @ 4.20GHz
• GEFORCE GTX 1080
• Ubuntu 16:04

学習/検証 データセット

BSDS3000 データセットを用いました. 含まれる画像が300枚と少ないので random clip / horizontal flipping を行い水増ししています。1epochの定義は 100000 images と定義しています。 拡大のスケーリングは3倍としました。ですので、学習時にはもとの画像を 1/3 に縮小した画像を入力とし、出力を元の画像と合わせるように学習を行います。

検証用データセットには前回同様に Set5 を用いました。

最適化

最適化手法は Adam/SGD/Adabound 各々でパラメータの組み合わせでグリッドサーチを行いました。
ここですべての Optimizer に共通な条件として epoch=45 / 15 epoch ごとに 0.1 倍の lr decay, weight_decay=1e-8 と設定しています。

以下で結果に示すモデルは Validation Loss が最も良かった Adam (lr=1e-4, weight_decay=1e-8) を用いて学習されたものです。

一回の学習はだいたい 30 分程度です。 データの加工部分に画像を都度 random clip して縮小するというまあまあ重たい処理が入っているので, workers=4 程度だとGPU 律速ではなく CPU 律速になっているように見受けられました。逐次水増しせず事前に切り出しておくなどしておくべきだったかもしれません。

結果

実際に学習されたネットワークを用いて超解像度化を行ってみます。Set5 のなかではもっとも復元が難しい(bicubicでのRMSEが最も低い) butterfly を用います。実際に超解像度化を行ったのが以下の図です。

butterfly x3. 左から入力画像(元の画像の1/3の画像), 古典手法(BICUBIC)での拡大, ESPCN での拡大

BICUBIC に比べると ESPCN では黒い部分の境界が鮮明に復元できていることがわかります。

おまけ

学習/Validation データのロス/PSNR の遷移は以下の様な感じ。ほぼ同じような動きをしているので学習はできていることがわかります。

まとめと感想

SRCNN の次世代アーキテクチャ ESPCN を紹介しました。この手法の肝は PixelShuffle で、入力画像が直接拡大される、輝度情報以外はネットワークでは扱わないのが変わった点です。

本当は SRCNN の結果と比べたい、のですがこれを書いている時に SRCNN の実験を x2 で行っていたことに気がついてしまい乗せられませんでした。今から実験します…

個人的に面白いなと思ったのは ESPCN では adam が一番よかった、という点です。 SRCNN では一番良かったのは SGD + Nesterov + Momentum で, adam は training loss は一番下がるが汎化せず(Validation Data のロスが全く下がらない)過学習に陥っていました。おそらく原因はネットワークの構造が単純すぎるため、すぐに最適解にたどり着けてしまう adam では局所解にハマってぬけだせなくなるからかなーと推察していました。

よりネットワークが複雑な ESPCN では adam が一番汎化することからも、上記の仮説はあっているような気がしています。単純すぎるときはゆっくり最適化するほうがいいのは面白いですね。

次やりたいこと

2016年まで来た。もうちょっと未来にいく。次はGANをつかった人間が見た時に自然な復元を試みている手法を実装したい。

参考文献

• Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network
• leftthomas/ESPCN
  • pytorch espcn でググった時に一番上に出てくる実装。最終的な出力に sigmoid を掛ける処理が入っているのですが論文中にその記述が見つからなかったので僕の実装ではそのままを出力にしています。