UnrolledGANをTensorFlowで実装した

別のGANの論文を読むつもりが間違えてUnrolledGANと呼ばれるものを印刷して読んでしまった
途中でこれ違う論文やんけと気づいたけど紙がもったいないのでちゃんと読んでUnrolledGANを実装した。

まちがえて読んだ論文はこれ

[1611.02163] Unrolled Generative Adversarial Networks

GANとDCGANでは学習の不安定さや生成される画像の多様性がしょぼい、ノイズがかなり含まれる、などが問題として挙げられてた
この論文で学習の不安定さの原因として識別器の方が生成器よりも学習が早いというのが書かれていた
先に識別器の方がめちゃ強くなってしまう。ある学習途中のまだ弱い生成器にとって、強力な識別器の学習信号は学習信号として役立たなくなるとも書いてある。まじかよ・・・

リアルで考えれば、AとBの二人がお互い助け合って数学の勉強をしているとして、Aの方が成長が早く、AがBに教えようとするとBはまだAと同じレベルにいないからAの言っていることが時々わからずBの成長が悪くなってしまう感じだろうか。GANに限るかわからないが二人三脚な感じで学習していくのがベストってことなのかな

こういったことが起こると途中で学習が収束(悪い意味で崩壊)してそのエポック以降は同じもの、またはそれに類似したものしか生成しなくなってしまうんだとか。
これには思い当たりがある。実際前に作ったDCGANでも途中で識別器と生成器の損失が変わらず、次のような画像しか生成されないことが何回かあった

識別器が強力になることが原因だったとは知らなんだ。DCGANで識別器にのみバッチ正規化を適用しなかったり、重み減衰を適用してうまくいったのはつまりそういうことだったのか

ここまで論文を読んで、じゃあ識別器の学習をのろまにさせる手法がUnrolledなのかなと思ったけどそうじゃなかった。逆に生成器の学習スピードが識別器に追いつけるようにする方法が書かれてあった。識別器を悪くさせる必要ないもんな・・・

今までのGANは次のように両モデルのパラメータを更新していった

θG←θG−ηdf(θG,θD)dθG${\displaystyle {\theta }_G\leftarrow {\theta }_G-\eta \frac{df({\theta }_G,{\theta }_D)}{d{\theta }_G}}$ θD←θD+ηdf(θG,θD)dθD${\displaystyle {\theta }_D\leftarrow {\theta }_D+\eta \frac{df({\theta }_G,{\theta }_D)}{d{\theta }_D}}$

η${\displaystyle \eta }$は学習係数で、 f(θG,θD)${\displaystyle f({\theta }_G,{\theta }_D)}$ は目的関数

Unrolled GANではまず次のように定義して

θ0D=θD${\displaystyle {\theta }_D^0={\theta }_D}$ θk+1D=θkD+ηkdf(θG,θkD)dθkD${\displaystyle {\theta }_D^{k+1}={\theta }_D^k+{\eta }^k\frac{df({\theta }_G,{\theta }_D^k)}{d{\theta }_D^k}}$ fk(θG,θD)=f(θG,θkD(θG,θD))${\displaystyle f_k({\theta }_G,{\theta }_D)=f({\theta }_G,{\theta }_D^k({\theta }_G,{\theta }_D))}$

以下のようにパラメータを更新する

θG←θG−ηdfk(θG,θD)dθG${\displaystyle {\theta }_G\leftarrow {\theta }_G-\eta \frac{d{f}_k({\theta }_G,{\theta }_D)}{d{\theta }_G}}$ θD←θD+ηdf(θG,θD)dθD${\displaystyle {\theta }_D\leftarrow {\theta }_D+\eta \frac{df({\theta }_G,{\theta }_D)}{d{\theta }_D}}$

GANと変わったところは赤で書いた。数式眺めてすぐに理解できるほど賢くないけど、生成器のパラメータ更新は k回パラメータを更新した識別器を用いて更新するということなんじゃなかろうか

もしk = 0なら1回だけ識別器を更新して、次に生成器を1回更新する今までのGANと同じと書いてあるからkが１以上の自然数で与えればいいんだな

こうすることで生成器は識別器がどのように反応するかを考慮してより本物を作れるように学習ができるという理由が定義式の後に書いてあった。
まとめると識別器の更新パラメータは最初の1回の更新パラメータを採用する。生成器の更新パラメータはk回更新したときの識別器のパラメータを用いて更新したものを採用するということ。
論文ではこの生成器の更新の説明で 'see into the future' とかっこよく書かれていた。この一文を見たとき「才気煥発の極みじゃん」とかくだらないこと思ってしまった。

GANとの変更点は上記のパラメータ更新あたりしかない。それだけで改善するんかほんま・・・と思いつつ論文の実験の部分を読んで次のような実験結果が載っていた

右端のTargetが二次元のガウス分布の入力データで、上の段がUnrolledGAN、下の段が通常のGAN(k=0)の学習の途中経過

UnrolledGANのがうまく多様性も含めて学習してるのがわかる。下の段でTargetの一部分しか学習できていない。それ以上にめちゃくちゃこの画像かっこいい。

かっこいいと思ったので同じような結果が得られるのか自分でUnrolledGANを実装して観察することにする

コードはここ

github.com

Unrolledな実装部分は流れとしてはこんな感じで作ることにする

めんどくさいのがUnrolledStepに入る前の識別器の重みの保存。生成器のパラメータ更新が終わった後に識別器のパラメータを θ1D${\displaystyle {\theta }_D^1}$ のものに復元しないといけない。

グラフをUnrolledする前にフリーズさせて終わった後にフリーズさせたグラフから θ1D${\displaystyle {\theta }_D^1}$ のパラメータ取れるかなとか思ったけど、Variableノードのassign使えば雑だけどできそうだったのでこっちでやった

TensorFlow: Mutating variables and control flow – metaflow-ai

困ったのがメモリーリークだった
最初次のように識別器の重みをforループ内でコピーした

copy_dis_w1 = sess.run(self.discrimitor.dis_w1)

そして復元については次のようにしてた

sess.run(tf.assign(self.discrimitor.copy_dis_w1,copy_dis_w1))

何がまずいってforループ内でこんなことすると学習するたびに1Epochの終わる時間が4秒、5秒、7秒、...と増えていった

さすがにおかしいと思って調べたら復元のところでtf.assign()の引数でcopy_dis_w1がTensorじゃないため、呼び出すたびにcopy_dis_w1を格納するtf.constant()もどきの変数を作っていたらしい。いかんでしょ

それに加えてtf.assignのオペレーションのノード自体も呼び出すたびに作られていた。アホやんけ・・・

めんどくさがって学習のループ内でオペレーションを呼び出すんじゃなかった。学習に入る前に事前にオペレーションを定義しておいた。
いちいちtf.constantで退避させた重みを入れるのも嫌なのでtf.placeholderを作っておいてそれに格納してメモリーリークの問題を解決した。

以下結果

入力データを分散が0.01くらいの正規分布から作って8方向にシフトさせる

Unrolledする回数として論文のコードを参考に k=5と設定した。また入力のノイズ変数の次元は100次元じゃなくて50次元に削減した。
自分で作ったUnrolledGANの生成結果としてはこんな感じ(Epoch : 1000)

途中経過
　

同じモデル構造で、k=0のGANだと次のような学習結果が得られた

論文と同じような結果が大体得られた。重み減衰入れなくてもめっちゃ安定している。

おまけでノイズ変数を20次元にしてk=3で学習させてみたら芸術的なものができた