AdaINを実装して試す

Image style transferに興味を持ったのでArbitrary Style Transfer in Real-time with Adaptive Instance Normalizationを実装して試してみました．

実装 github.com

論文

Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization （ICCV, 2019） https://arxiv.org/abs/1703.06868

NNベースのImage style transferで著名なGatysらの手法[1]はスタイル転送したい画像ペアごとに学習が必要なのがネックでした．この論文ではスタイルの特徴をコンテンツ画像とミックスするためのAdaptive Instance Nomalization（AdaIN）レイヤを導入し，一度の学習で任意のコンテンツ-スタイル変換を実現するアーキテクチャを提案しています．

AdaIN

AdaINはInstance Normalizationの拡張したものという位置づけです． コンテンツとスタイルの特徴量を入力とし，コンテンツの特徴の平均と分散がスタイルの特徴が持つ平均と分散のスケールに合うよう線形変換する処理です．このとき特徴のチャネルごとに変換をかけます． コンテンツ画像をx，スタイル画像をyとするとき，AdaINは次のような演算になります．

Pytorchで実装すると以下のようになるでしょうか．

class AdaptiveInstanceNorm2d(nn.Module):
    def __init__(self):
        super(AdaptiveInstanceNorm2d, self).__init__()
        self.eps = 1e-5

    def forward(self, x, y):
        mean_x, mean_y = torch.mean(x, dim=(2, 3), keepdim=True), torch.mean(y, dim=(2, 3), keepdim=True)
        std_x, std_y = torch.std(x, dim=(2, 3), keepdim=True) + self.eps, torch.std(y, dim=(2, 3), keepdim=True) + self.eps
        return std_y * (x - mean_x) / std_x + mean_y

AdaINがスタイル転送に有効な理由としては，例えばスタイル画像が明確なテクスチャの変化を持つときにはスタイル特徴も強い活性を示すはずなので，これによる平均，分散などの統計量の変化をコンテンツ特徴に反映すれば，スタイルに順応しやすくかつコンテンツの持つ幾何的な情報を保持した出力ができるからだとしています．

実験

コンテンツ，スタイルともに5000枚ずつ用意し，40エポックだけ学習しました．オプティマイザはデフォルトのAdamです．

結果は以下．

これ以外にも他の写真も貼り付けようと思っていたのですがなぜか横長の画像が貼り付けられませんでした．ただ他のはあまりうまく行っていません．まだ実装に間違いがあるかも．

参考文献

[1] L. A. Gatys, et al., "Image Style Transfer Using Convolutional Neural Networks," In CVPR, 2016