名古屋出身ソフトウェアエンジニアのブログ

プリキュア StyleGAN を作った話 〜 其の弐

公開:
更新:

前回記事（執筆のやる気がなさすぎて二年近くも開いている）では StyleGAN 1.0 でプリキュアの顔画像を生成しましたが、FID 80 というなんとも微妙な結果で終わりました。今回は StyleGAN 2.0 と Adaptive Discriminator Augmentation (ADA) を導入して、FID を 17.39 まで改善した話です。

curegit/precure-stylegan-ada: Yet another StyleGAN 2.0 implementation using Chainer with Adaptive Discriminator Augmentationto to synthesize specific Precure (Cure Beauty) images
Yet another StyleGAN 2.0 implementation using Chainer with Adaptive Discriminator Augmentationto to synthesize specific Precure (Cure Beauty) images - curegit/precure-stylegan-ada
目次

先に結果を貼ります。

完成版モデル (Cure Beauty v2, ψ = 0.8) の生成例

完成版モデル (Cure Beauty v2, ψ = 0.8) の生成例

前回のものと比べると、だいぶプロフィールアイコン等に使えそうな見た目になりました。以下、何を入れたらこうなったのかを実装レベルで書いていきます。

StyleGAN 2.0 での変更点

StyleGAN 2.01 は、1.0 のアーキテクチャと正則化を作り直して、画像品質と学習の扱いやすさを改善したものです。このプロジェクト的に効果があったのは、Progressive Growing の廃止（運用が楽になる）と Weight Demodulation（アーティファクトが消える）でした。

Progressive Growing がなくなった

1.0 の実装では、低解像度から順に stage を上げながら学習していました。学習 CLI プログラムは -s で現在の stage を、-a で新解像度をブレンドする係数 α\alpha を、-t で α\alpha の 1 イテレーションあたりの増分を指定する形です。

# StyleGAN 1.0 実装での学習（stage ごとに起動し直す）
python3 train.py -s 5 -a 0.0 -t 0.00005 -x 7 -b 16 -e 100 DATASET_DIR

つまり、解像度を上げるたびに人間が学習を止め、次の stage の引数を組み直して再開する必要がありました。各 stage に何エポックかけるか、 α\alpha をどこまで伸ばしてから次に行くか、といった判断も手作業です（本当は自動化もできるけど、データセット的にも学習の進みが安定するか怪しかったので、見ながらの調整だった）。出力されるスナップショットのファイル名に stage とイテレーション数が入っている（gen_5_12000.hdf5 のような命名）のも、この運用の名残です。

2.0 ではこれが丸ごと不要になります。最初から最終解像度のネットワークを組んで、ひたすら回すだけです。

# StyleGAN 2.0 実装での学習（これ一発でよい）
python3 train.py -e 100 -x 7 -b 16 -k 8 --ada -o results DATASET_DIR

Progressive Growing をやめられたのは、代わりに multi-scale な出力経路が導入されたからです（MSG-GAN2 の影響が見えます）。 Generator は各解像度ブロックが torgb を持ち、下位ブロックの RGB 出力をアップサンプルして足し込む skip 構成になっています。

def __call__(self, x, y, w, noise=1.0, fixed=None):
	h1 = self.up(x)
	h2 = self.wmconv1(h1, self.style1(w))
	h3 = self.noise1(h2, coefficient=noise, fixed=fixed)
	h4 = self.act1(h3)
	h5 = self.wmconv2(h4, self.style2(w))
	h6 = self.noise2(h5, coefficient=noise, fixed=fixed)
	h7 = self.act2(h6)
	return h7, self.skip(y) + self.torgb(h7, self.style3(w))

Discriminator 側は residual ブロックの積み重ねです。

def __call__(self, x):
	h = self.pool(self.act2(self.conv2(self.act1(self.conv1(x)))))
	skip = self.down(self.conv3(x))
	return (h + skip) / root(2)

どちらも全解像度の経路が最初から使えるので、学習の途中でネットワーク構造を切り替える必要がありません。これは運用が楽になるだけでなく、Progressive Growing が持っていた副作用の解消にもなっています。各解像度が一時的に最終出力層として振る舞うせいで、細部の特徴が画像内の特定位置に固執する（顔の向きが変わっても歯や目の位置が動かない）という現象が知られていました。

副産物として、スナップショットからの再開や転移学習の実装も素直になりました。 stage の概念がないので、どのレベルから上の重みを移植するかという選択だけで話が済みます。

AdaIN から Weight Demodulation へ

1.0 では、スタイルの適用に AdaIN を使っていました。特徴マップを正規化してから、スタイルベクトル由来のスケールとバイアスを掛け直す操作です。

2.0 ではこれをやめて、畳み込みの重みそのものをスタイルで変調し、変調後の重みのノルムで正規化 (demodulation) します。

def __call__(self, x, y):
	out_channels = self.b.shape[0]
	batch, in_channels, height, width = x.shape
	modulated_w = self.w * y.reshape(batch, 1, in_channels, 1, 1)
	w = modulated_w / sqrt(sum(modulated_w ** 2, axis=(2, 3, 4), keepdims=True) + 1e-08) if self.demod else modulated_w
	grouped_w = w.reshape(batch * out_channels, in_channels, self.ksize, self.ksize)
	grouped_x = x.reshape(1, batch * in_channels, height, width)
	padded_grouped_x = pad(grouped_x, ((0, 0), (0, 0), (self.pad, self.pad), (self.pad, self.pad)), mode="edge") if self.pad else grouped_x
	h = convolution_2d(padded_grouped_x, grouped_w, stride=1, pad=0, groups=batch)
	return h.reshape(batch, out_channels, height, width) + self.b.reshape(1, out_channels, 1, 1)

実装上のポイントは、変調後の重みがミニバッチの要素ごとに異なるという点です。素朴に書くとバッチ内でループを回すことになるので、バッチ次元をチャンネル次元に潰して groups=batch のグループ化畳み込みにするお決まりのトリックで回避します。

AdaIN による特徴マップの正規化は、Generator が「特定チャンネルにわざと強い信号を作って正規化を実質的に無効化する」抜け道を許してしまい、これが水滴状のアーティファクトの原因でした。 Weight Demodulation はデータの統計に関係なく重み側だけで正規化するので、この抜け道がなくなります。

なお、パディングは edge（端の値を複製）にしています。ゼロパディングだと生成画像の縁に不自然な部分が出やすいためです。

Path Length Regularization

スタイル空間 W\mathcal{W} での移動量と、生成画像の変化量が比例していると嬉しい、という考え方があります。 Path Length Regularization は、 W\mathcal{W} に一定の摂動を与えたときの画像の変化量が一定になるよう、ヤコビアンのノルムを目標値へ寄せる正則化です。

@staticmethod
def path_length(ws, x, mask):
	levels, batch, size = len(ws), *(ws[0].shape)
	gradients = grad([x * mask], ws, enable_double_backprop=True)
	gradient = stack(gradients).transpose(1, 0, 2).reshape(batch * levels, size)
	path_lengths = batch_l2_norm_squared(gradient).reshape(batch, levels)
	return sqrt(mean(path_lengths, axis=1))

ヤコビアンを陽に作るのは無理なので、画像側にランダムな方向 mask を掛けてから W\mathcal{W} へ逆伝播させ、ヤコビアン・ベクトル積として求めています。目標値は過去のパス長の指数移動平均で、学習中に更新していきます。

潜在空間が滑らかになるので、Style Mixing や補間アニメーションの品質に貢献します。

Lazy Regularization

R1 勾配ペナルティも Path Length Regularization も二階微分が必要で高価です。毎イテレーション計算する必要はないので、数回に 1 回だけ計算します。

def r1_regularization(self):
	return self.r1_regularization_interval and self.iteration % self.r1_regularization_interval == 0

間隔を空けた分だけ係数を掛けて、平均的な強さを揃えておくのを忘れないようにします。

if self.r1_regularization():
	weight = self.r1_regularization_interval * self.r1_regularization_gamma
	penalty = weight * CustomUpdater.gradient_penalty(x_real, y_real) / self.batch_size

この実装のデフォルトは R1 が 16 イテレーションごと、Path Length が 8 イテレーションごとです。

Adaptive Discriminator Augmentation

ここからが品質改善の本題です。 StyleGAN 2.0 化よりも ADA3 の効果のほうが圧倒的に大きかったです。

少数データでは Discriminator が過学習する

前回の結果が微妙だった原因は、データセットが（水平反転で水増しして）2000 枚程度しかなかったことです。このサイズだと Discriminator が訓練画像を覚えてしまい、本物と偽物を完璧に見分けられるようになります。そうなると Generator に返る勾配が意味を失い、学習がろくに進まなくなります。

普通の識別タスクなら、データ拡張して過学習を防ぐところです。しかし GAN で素朴に拡張画像をデータセットへ混ぜると、Generator は拡張後の分布を学んでしまいます。前回、水平反転画像を足しておいて害がなかったのは、反転しても許容できる顔画像であるという性質にたまたま助けられていたからです。同じノリで色をいじる拡張を混ぜると、色がおかしい画像を生成するようになります（論文では拡張が leak すると言っている）。

拡張を leak させない

ADA の解決策は、拡張をデータセットではなく Discriminator の入力に掛けることです。本物/偽物のどちらにも同じ確率 pp で同じ拡張パイプラインを通し、Discriminator には拡張後の画像しか見せません。

y_real = self.discriminator(self.augumentation_pipeline(x_real), c_real)
...
y_fake = self.discriminator(self.augumentation_pipeline(x_fake), c)

前回記事のデータフロー図でいうと、Discriminator の手前に拡張が挟まる形になります。

  ---
title: ADA での Discriminator の訓練時
---
graph LR
  N(Random)
  G[Generator]
  A1[Augment]
  A2[Augment]
  D[Discriminator]
  R(Real)
  F(Fake)
  L(Loss)
  P(Prediction)
  T(Ground Truth)
  N --> G
  G --> F
  F --> A1
  R --> A2
  A1 --> D
  A2 --> D
  D --> P
  P --> L
  T --> L
  linkStyle 6,7 stroke:#31bec5,stroke-width:4px

Generator を訓練するときも、生成画像は同じパイプラインを通してから Discriminator へ入れます。

ws, x_fake = self.generator(z, c, random_mix=mix)
y_fake = self.discriminator(self.augumentation_pipeline(x_fake), c)

こちらのデータの流れは以下のようになります。

  ---
title: ADA での Generator の訓練時
---
graph LR
  N(Random)
  G[Generator]
  A[Augment]
  D[Discriminator]
  F(Fake)
  L(Loss)
  P(Prediction)
  N --> G
  G --> F
  F --> A
  A --> D
  D --> P
  P --> L
  linkStyle 1,2,3,4,5 stroke:#31bec5,stroke-width:4px

図の青緑線の通り、Generator への逆伝播は拡張パイプラインを通過します。したがって、拡張の実装は微分可能でなければなりません。リサンプリングや FFT を含む処理まで全部 Chainer の関数で書いているのはそのためです。

こうすると、Generator は拡張後の分布が本物の拡張後の分布と一致するように学習します。拡張が分布として可逆であれば、そこから拡張前の分布、つまり本来欲しい分布が復元されるという理屈です。ただし pp が 1 に近づくと Discriminator がほぼ拡張画像しか見なくなり、可逆性が崩れて leak します。この実装では上限をデフォルトで 0.8 にしています。

拡張確率を自動で決める

pp を手で調整するのは面倒なので、Discriminator の過学習度合いを測って自動調整します。論文の指標は、本物データに対する Discriminator 出力の符号の期待値です。

rt=E[sign(D(Xreal))] r_t = \mathbb{E}\left[\mathrm{sign}(D(X_{real}))\right]

rt=0r_t = 0 なら本物を五分五分に判定している（過学習なし）、 rt=1r_t = 1 なら本物を全部本物と言い当てている（完全に過学習）という見方になります。

accumulated_rt += sum(sign(y_real - 0.5 if self.lsgan else y_real)).item() / self.batch_size

これを目標値（デフォルト 0.6）と比べて、 pp を上下させるだけです。

def adapt_augumentation(self, overfitting):
	delta = self.batch_size / self.augumentation_delta_images
	direction = sgn(overfitting - self.overfitting_target)
	probability = self.augumentation_probability + delta * direction
	self.augumentation_probability = clamp(0.0, probability, self.augumentation_limit)

更新幅は「画像 50 万枚で pp が 0 から 1 まで動く」速さとして正規化されているので、バッチサイズを変えても挙動が揃います。学習の序盤は p=0p = 0 から始まり、Discriminator が強くなってくると自動的に拡張が強まっていきます。

拡張パイプライン

拡張は固定の順序で並べ、それぞれを確率 pp で適用します。

self.manipulations = ChainList(
	Mirror(probability_multiplier=pixel),
	Rotation(probability_multiplier=pixel),
	Shift(probability_multiplier=pixel),
	AffineTransformation(probability_multiplier=geometric),
	ColorAffineTransformation(probability_multiplier=color),
	Filtering(probability_multiplier=filtering),
	AdditiveNoise(probability_multiplier=noise),
	Cutout(probability_multiplier=noise),
)

中身は以下の通りで、論文のカテゴリにおおよそ対応します。

  • Mirror, Rotation, Shift: 水平反転・90 度単位の回転・整数ピクセルのシフト（ピクセル単位の操作なので情報を落とさない）
  • AffineTransformation: 任意角回転・等方/非等方スケーリング・平行移動
  • ColorAffineTransformation: 明度・コントラスト・色相回転・彩度の調整
  • Filtering: FFT して 4 つの周波数帯のゲインをランダムに変える（リファレンス実装だとウェーブレット変換を使っていたが、よう分からんかった）
  • AdditiveNoise, Cutout: ガウスノイズの加算と矩形の塗り潰し

幾何変換と色変換は、個々の操作を順番に適用するのではなく、全部を 1 つのアフィン行列（色の方は RGB 空間の 4×44 \times 4 行列）へ合成してから一度だけ適用するようにしています。補間を何度も通すと誤差が累積するので、リサンプリングは 1 回で済ませたいわけです（学習速度もそのほうが速い）。

カテゴリごとの適用率倍率はコマンドラインから指定できます (-I, -G, -C, -F, -N)。理屈の上では leak しない範囲であれば人間目線で不自然な拡張でも構わないのですが、実際の効き方はデータセット依存なので、特定カテゴリは無効化したい場合も出てきます。

結果

FID は 80.43 から 17.39 まで下がりました。

モデルFID条件
StyleGAN 1.0（前回）80.43ADA なし
Cure Beauty v121.98ADA あり, batch = 16, R1 gamma = 100
Cure Beauty v217.39ADA あり, batch = 12, R1 gamma = 75

水平反転は ADA のパイプライン (Mirror) に含まれるので、前回のように事前に反転画像を足しておく必要はなくなりました。反転なしの実データは追加収集して 3000 枚程度になっていますが、論文が扱う規模と比べれば依然として桁違いに小さいままです。

v1 から v2 の差は、バッチサイズを 16 から 12 に、R1 の係数を 100 から 75 に下げただけです。これだけで品質と多様性の両方が改善しました。

R1 の係数がデフォルトの 10 より一桁大きいのは、データセットが小さいためです。少数データでは Discriminator が暴れやすく、勾配ペナルティを強く掛けたほうが安定しました。とはいえ強すぎると Discriminator が鈍って多様性が落ちるので、この辺りは回して試すしかありません。

学習曲線はこんな感じになります。

学習初期の曲線例

学習初期の曲線例

GAN の損失は下がれば下がるほど良いというものではないので、曲線からは発散していないかくらいしか読めません（とはいえ、Generator の損失がじわじわ増えるような場合は経験的に失敗している）。むしろ overfitting ( rtr_t ) が目標値付近に落ち着いているかを見ていました。ここが 1 に張り付くようなら、 pp が上限に達しても過学習を抑えきれていないということなので、データセットを増やすしかありません。

転移学習

転移学習とファインチューニングにも対応させました。実装は Mo et al.4 を参考にしていて、-t で移植元スナップショットと移植するレベル範囲を、-Z で更新を止める（凍結する）レベル範囲を指定します。 Generator は指定レベル以上（高解像度側）、Discriminator は指定レベル以下（入力側）を対象にする作りで、Discriminator 側の凍結は同論文の FreezeD に相当します。

以下は Cure Beauty v2 を出発点に、300 枚程度のキュアトゥインクルへ転移した例です。さすがに枚数が少なすぎて厳しいようです。

キュアトゥインクル（Cure Beauty v2 から転移）

キュアトゥインクル（Cure Beauty v2 から転移）

まとめ

  • 品質改善の主役は ADA で、少数データでも Discriminator の過学習を抑えられる
  • FID は 80.43 から 17.39 まで改善したが、論文級の数値を狙うならやはりデータ量が必要

関係ない所感

LLM によるテキスト生成技術活用はもはや当たり前といった感じですが、これを作った当時はこれ全部コード自分で書いてたんだなと感慨深い気持ちになりました。

なお、本編に載せきれなかった実装の細部は次回にまとめました。

  1. Karras et al. Analyzing and Improving the Image Quality of StyleGAN (2020). https://doi.org/10.48550/arXiv.1912.04958 ↩︎

  2. Karnewar and Wang. MSG-GAN: Multi-Scale Gradients for Generative Adversarial Networks (2020). https://doi.org/10.48550/arXiv.1903.06048 ↩︎

  3. Karras et al. Training Generative Adversarial Networks with Limited Data (2020). https://doi.org/10.48550/arXiv.2006.06676 ↩︎

  4. Mo et al. On Leveraging Pretrained GANs for Generation with Limited Data (2020). https://doi.org/10.48550/arXiv.2002.11810 ↩︎

Effective Python 第3版
まちカドまぞく2丁目 コンパクト・コレクション
低レイヤの世界を探検するテクニック89選
メダリスト Season2 vol.1 [Blu-ray]