PyTorch×VAEGANによるソ連プロパガンダポスターの再構成

#概要

前回に引き続き、プログラミングの練習として、いろいろなものを実装しています。
今回はVAEGANを用いて、プロパガンダポスターを再構成しようと思います。
（突然なんでソ連？って話なんですけど、本当に単なる趣味なんです。政治的主張とか一切ございません。）

技術的なご指摘点などありましたらコメントいただけると幸いです。
当記事でご理解いただけるのは、VAEGANの再構成のレベル、GANの学習における留意点、がメインです。

再構成の流れとしては、以下の通りです。

1. プロパガンダポスター（学習データ）を集める
2. VAEGANで学習を行う
3. 学習済みモデルを使って、データの再構成を行う

なお、PyTorchのVAEGANの実装はこちらのを基本とさせていただいてます。ありがとうございます。

#VAEGANについて
Autoencoding beyond pixels using a learned similarity metric
通常の単純なGANでは、GeneratorとDiscriminatorの2つのmodelがあり、それらが互いに敵対的学習を行うことで、優れたGeneratorが学習できるというような仕組みでした。
しかしGANの問題点として、**Generatorの入力がランダムなノイズである点が挙げられます。**どういうことかというと、入力がランダムなので、明示的にこのデータが欲しい！と思っても、狙ったデータを生成することが通常のGANではできないのです。
一方VAEGANでは、明示的に欲しいデータを再構成できるように設計されています。どういうことかというと、GeneratorがVAEに入れ替わった形をしています。
これにより、入力データに近いデータを明示的に再構成可能となっているわけです。
今回はソ連風なポスターのデータを明示的に再構成させようと思います。

#学習データを集める
Image Downloaderを使いました。このchrome用の拡張機能は、学習データを集める際にめちゃくちゃ便利です。おすすめです。
そして以下のようなデータを300枚ほど収集しました。

いやーこのアバンギャルドな感じが素晴らしいですね！（あくまで趣味ですよ）

#VAEGANで学習を行う
実装は基本的に上記の方のを参考にさせていただきました。しかし、そのままではうまく学習ができない（再構成に失敗する）時もあったので、一部修正を行っています。大きく分けて、誤差関数とモデル定義の２つです。

誤差関数は以下のように、train.pyにあるEnc_lossのL_llikeを無しに変更しました。
また、Dec_lossのL_gan_fakeも無しに変更しています。
上記二つのL_llike、L_gan_fakeは、Discriminatorの特徴を誤差関数に組み込んでいるみたいです。
僕の環境では、オリジナルのverですと、残念ながらなかなか収束することができませんでした、、、

train.py

# Enc_loss ---
    Enc_loss = L_prior + L_recon 
    Enc_loss.backward()
    Enc_optimizer.step()
    Enc_running_loss += Enc_loss.item()

    # train Decoder ===
    Dec_optimizer.zero_grad()

    x_real = Variable(data)
    z_fake_p = Variable(torch.randn(opt.batch_size, opt.nz))
    if opt.use_gpu:
      x_real = x_real.cuda()
      z_fake_p = z_fake_p.cuda()
    x_fake, mu, logvar = G(x_real)

    # L_gan ---
    y_real_loss = bce_loss(D(x_real), t_fake)
    y_fake_loss = bce_loss(D(x_fake), t_real)
    y_fake_p_loss = bce_loss(D(G.decoder(z_fake_p)), t_real)
    L_gan_fake = (y_real_loss + y_fake_loss + y_fake_p_loss) / 3.0

    # L_llike ---
    L_recon = opt.gamma * l1_loss(x_fake, x_real)
    L_llike = l1_loss(D.feature(x_fake), D.feature(x_real))

    # Dec_loss ---
    Dec_loss = L_recon + L_llike
    Dec_loss.backward()
    Dec_optimizer.step()
    Dec_running_loss += Dec_loss.item()

モデル定義では、各モデルの活性化関数をLeakyReLUにした点、DiscriminatorのBatchNorm2dを無くした点が主な変更です。
また、入力のサイズを128×128とし、channelはカラー画像のため3としています。

#学習済みモデルを使って、データの再構成を行う
train.pyを回してlogsディレクトリに100epoch毎に再構成結果を保存しました。
epoch=1000、batchsize=16、潜在変数の次元=32として計算しました。
再構成結果は以下の通りです。なお、各画像の上段は入力画像、下段がVAEGANの出力結果です。

100epoch

300epoch

600epoch

1000epoch

#考察
VAEGANすごいな..と素直に思いました。100epoch目ではまた再構成結果としては、ぼやけていますが、1000epochも回せばほぼ本物（入力）と差がないほど完璧に再現できていると思います。当記事の最初に挙げたVAEGANの論文でも、VAEとの比較がなされており、極めてきれいに再構成できることはわかってはいましたが、今回の検証で改めて実感できました。
しかしGANの学習において留意すべきことは多くあります。
まずそもそもGANの学習は複雑で安定しにくい傾向があるため、Discriminatorとそのほかの各モデルのLossの推移に注意しなくてはなりません。
うまくいかないケースとして、DiscriminatorのLossがものすごい勢いで収束してしまうことや、一向にほかのLossが収束せずむしろ発散することがあります。
今回はDiscriminatorがやたらに収束するケースが発生したので、DiscriminatorのBatchNorm2dを無くしました。
また、最近のGANモデルにおいては、ReLUよりもLeakyReLUが使われているようです。
そのほか、GANの学習がうまくいかない状況に対する対処として、こちらも参考になります。