Pix2pix + Gradient penaltyを用いた線画着色

Introduction

前に線画着色についての記事を書きましたが、個人的に不満足だったので書き直しの意味で今回の記事を書きます(前回の記事は削除します)。ただ、前回の内容をそのまま行うのも面白くないのでネットワーク構造と損失関数を変えて実験した結果を載せます。

Network Architecture

今回用いたネットワーク構造: Generatorの出力はチャンネル数3。Discriminatorの出力はチャンネル数1。

UNetベースで特段面白い構造ではないですが、GeneratorにResBlockを導入し、アップサンプリングとしてはNearest Neighbor Upsampling -> Convolutionを行っています。

Dataset

線画着色するにあたって、今回線画と着色画像を19000ペアを用意しました。

  • 着色画像はsafebooruからスクレイピングしました。なるべく学習すべき着色部分を減らしたいので、背景画像は白で統一する必要がありました。そこでsafebooruのwhite_backgroundタグでスクレイピングしました。集めた画像に対して、縦横短い方を256pixelに圧縮して、その後ランダムに256✕256サイズでクラップしました。
  • 線画は着色画像から抽出しています。抽出方法にはXDoGを用いています。計算する際に用いた論文中の各値は以下のようにしました。
    σ=0.4, k=4.5, p=19.0, ε=0.01, φ=10⁹

Experiment

Generator(UNet)とDiscriminatorの損失関数は以下のようにしています。

  • GeneratorはAdversarial lossとContent lossの重み付けした和としています。Content lossは生成画像と着色画像の平均絶対値誤差です。
    ※Content lossにVGGの出力層を用いたものも勿論考慮しましたが、何故かGeneratorのlossがあるiterationで急に大きくなったことが確認されたので、今回は採用しませんでした。ただ、VGGの出力層を用いること自体はよくある手法ですので、詳細は確認する次第です。
  • Discriminatorの損失関数としては以下の3つを考えました。
    (1) Adversarial lossのみ
    (2) Adversarial loss + zero-centered gradient penalty
    (3) Adversarial loss + one-centered gradient penalty(WGAN-GPで提案
     されたもの)
    Gradient Penaltyを求める際にDiscriminatorに入力するものは、生成分布と真のデータ分布からサンプリングしたものを混合させますが、この場合ではGeneratorの出力と着色画像の混合になります。
    また、Zero-centered gradient penaltyとはこちらで提案されたものです。この論文では、生成分布と真のデータ分布それぞれでGradient penaltyを計算していますが、ICLR2019で出る予定のこちらの論文によれば、それぞれで求めるとcollapseしてしまうため上記のように混同させた分布でGradient penaltyを求めるのが良い、とあったのでそのようにしています。

ハイパーパラメータ等詳細は以下に記します。

  • バッチサイズは4
  • DiscriminatorとGeneratorは交互に学習(Discriminatorが先)
  • 最適化手法はAdam(α=0.0001, β1=0.5, β2=0.9)
  • Generatorの損失関数において、Adversarial lossの重みは1.0、Content lossの重みは10.0
  • Discriminatorの損失関数において、Adversarial lossの重みは1.0、Gradient penaltyの重みは10.0

Result

まず、Discriminatorの損失関数は結局(1)~(3)のどれを用いるのが良いのか見るため、(1)~(3)の場合のGenerator lossを見てみました。Generatorの損失関数は(1)~(3)で共通のため値による比較が出来るはずです。結果を以下に示します。凡例は上から(1),(2),(3)です。各iteation毎に求めた(train時の)Generator lossの推移を表しています。
図から分かるように(2)のZero-centered gradient penaltyを用いたものが最もlossが小さくなっています。従って、今回はZero-centered gradient penaltyを用いることにしました。

それでは、着色結果を以下に示します。併せて(1)Adversarial lossのみと(2)Adversarial loss+1-centered gradient penaltyで着色した場合も乗せています。こう見ても0-centered gradient penaltyを用いた場合が自然に着色されていますね。

ただ、今回の手法でも問題点はあります。以下は0-centered gradient penaltyを用いた場合の着色画像ですが、黒を基調にした色になっています。これ以外にも黒になってしまうのは多く見られていて、過学習していることが分かります。着色としては自然ですが、色合いとしては微妙なので学習データを増やす等で対応していきたいです。

Summary

今回は、線画着色を行いました。様々なDiscriminatorの損失関数で試して0-centered gradient penaltyを用いることにしました。ただ、上記でも述べたように過学習していると考えられるので学習データを増やす等で対応していきたいと思います(そういう意味では必ずしも0-centered gradient penaltyを用いるのが良いとは限らないですね….)。

今回の試行はこちらの論文を参考にしました。画像生成タスクとしてGANを用いる際にgradient penaltyを用いるのは最早今となっては当たり前のように思われますが、今回の着色のようなStyle-Transferでも用いられており、その有効性が確認出来ました。これによってより安定した学習が出来ると良いですね。

Environment

OS : Ubuntu 16.04 LTS (64-bit)
CPU : Intel(R) Core(TM) i5–4590 CPU @ 3.30 GHZ
GPU : NVIDIA GTX970
メモリ : 8GB