プリキュア StyleGAN を作った話 〜 其の参
前回記事では StyleGAN 2.0 と ADA で FID を 17.39 まで改善した話を書きました。今回は本編に載せきれなかった実装の比較的どうでもいい細部の話です。
補間フィルタ
特徴マップのリサイズにおいて、元論文実装はバイリニア系ですが、この実装では Generator のアップサンプルに bicubic、Discriminator の skip 経路のダウンサンプルに Lanczos を使いました。 RGB の skip 経路だけは、よりローパスな B-spline 相当のパラメータ ( ) にしています。
self.up = BicubicUpsampler()
...
self.skip = BicubicUpsampler(1, 0)
BicubicUpsampler() の既定は (
) で、いわゆる Catmull-Rom です。少しシャープめになることを狙っています。
RGB の skip は解像度をまたいで足し込むので、リンギングが積み上がらないようぼかし寄りのカーネルにしました。
Discriminator の residual ブロックは、本線が平均プーリング、skip 側だけ Lanczos です。
self.pool = Downsampler(lanczos=False)
...
self.down = Downsampler()
これら bicubic や Lanczos のような非二項フィルタを使った効果が、目に見えてあるかどうかは正直微妙です。
勾配累積
StyleGAN 2.0 は 1.0 より全体的に重く、そこらへんの GPU では VRAM が足りません。バッチサイズは結果に影響するハイパーパラメータなので、VRAM の都合であまり下げたくはありません。
そこで -k でミニバッチを分割して勾配を溜められるようにしました。バッチサイズ 16 のまま、実際には 8 枚ずつ 2 回に分けて逆伝播する、といったことができます。
損失は分割後の枚数ではなく、もともとのバッチサイズで割っています。なので、更新 1 回あたりの勾配のスケールは、分割しない場合と揃います。
ミニバッチ標準偏差のグループサイズは -g で別に指定できるので、分割の仕方と独立に決められます。ただし、順伝播に乗る枚数より大きなグループは取れません。-k 8 なら -g の上限も 8 です。
Generator の重み平均
推論フェーズで生成に使うのは、学習中の重みそのものではなく指数移動平均を取ったコピーの重みです。学習が始まると Generator を 2 つ持ち、Adam が更新するのは元のほうだけで、もう一方には毎ステップ指数移動平均で混ぜ込んでいきます。
def average_generator(self):
decay = 0.5 ** (self.batch_size / self.averaging_images)
for raw, averaged in zip(self.generator.params(), self.averaged_generator.params()):
averaged.copydata(lerp(raw, averaged, decay))
なので、中身はよく見る EMA そのものです。
はバッチサイズ、 は半減期に相当する画像枚数(デフォルト 1 万枚)です。バッチサイズ 16 なら で、おおよそ 1 万枚見るごとに、古い平均の寄与が半分になります。
減衰率をイテレーション単位ではなく「画像 1 万枚で半減」として定義しているので、バッチサイズを変えても平均化の定数が変わりません。こういうオプションを画像枚数で正規化する流儀は ADA の更新幅と同じで、インターフェースからハイパーパラメータの再調整を減らすのに役立ちます。半減期は -E で変えられます。
なぜ指数移動平均を取るときれいになるのか
GAN の学習は、普通の損失最小化ではありません。 Generator と Discriminator は互いに相手を動かす鞍点問題です。重みは均衡点へ一直線に収束するというより、その周りを回り続けます。学習中のスナップショットは、その周回軌道上の一点です。たまたま Discriminator に過剰反応した位相だと、画像が荒れたり、モードが偏ったりします。
指数移動平均は、この振動を抑える操作です。均衡点のまわりの振幅が縮まった重みを取り出せるので、単発の更新よりも安定してきれいな画像が出ます。 Yazıcı et al.1 は、簡単な双線形ゲームでも EMA が振動の振幅を小さくすることを示しています。 Progressive Growing の頃から StyleGAN 系が推論専用に平均 Generator を持つのは、この効果を狙っています。
前回測った FID も、この平均 Generator に対する値です。生の重みで同じ指標を取ると、同じ学習でも実際に数字は悪く出ます。
ちなみにこの手法は、セグメンテーションや超解像のように画素単位で教師信号を与える普通の CNN でも使えます。
条件付き生成
おまけでクラス条件付きにも対応しました。 Generator 側はラベルの埋め込みを潜在ベクトルと連結して Mapper に入れ、Discriminator 側は最終特徴とラベル埋め込みの内積を取る projection 方式2です。
return h.reshape(batch) if c is None else sum(h * c1, axis=1) / root(channels)
クラス数が多いと多様性が死にやすかったので、Kuzushiji-49(49 クラス)では Mode Seeking Regularization3 も併用しています。潜在ベクトルが近いのに生成画像も近いという潰れ方を罰して、同じクラスの中でもモードを散らす正則化です。 AFHQ のようなクラス数が少ない条件付きモデルにはこの正則化は不要でした。
Chainer を最近の NumPy で動かす
Chainer は 2019 年末に開発が終了している4ので、最近の NumPy と組み合わせようとすると、そもそもインストール時点でこけたりします。np.bool や np.float といったエイリアス(1.24 で削除)、np.sctypes や numpy.distutils(2.0 で削除)などが原因です。
そこで、学習済み資産を保つために、起動時に足りないものを生やすだけのモジュールを用意して、スクリプトの先頭で当てるというゲリラパッチを実施しています。
import cure
cure.patch()
自宅 GPU をいたわる
学習は 1〜2 週間ぶっ通しになるので、夏場は室温とファンの音が気になってきます5。一定イテレーションごとに sleep を挟むオプションを付けました。
def sleep(self):
if self.sleep_interval and self.iteration % self.sleep_interval == 0:
time.sleep(self.sleep_seconds)
-Q 1 100 で、100 イテレーションごとに 1 秒休みます。
Yazıcı et al. The Unusual Effectiveness of Averaging in GAN Training (2019). https://doi.org/10.48550/arXiv.1806.04498 ↩︎
Miyato and Koyama. cGANs with Projection Discriminator (2018). https://doi.org/10.48550/arXiv.1802.05637 ↩︎
Mao et al. Mode Seeking Generative Adversarial Networks for Diverse Image Synthesis (2019). https://doi.org/10.48550/arXiv.1903.05628 ↩︎
ユーザーは深い悲しみに包まれた ↩︎
冬場はただの良い暖房なので問題なし ↩︎



