Diffusionモデルで猫を生成しようとした話

Diffusionモデルを自分で実装し、ノイズから猫を生成しようとした時の備忘録。

1~10日目: VAEが難しい

Diffusionモデルの解説記事を読むと、VAEという他の画像生成アルゴリズムが関係あると書いてあったので、そっちから入ることにした。このサイトの解説が詳しかったので参考にした。ただベイズ自体がかなり怪しくなっていたので、実際には同サイトのEMアルゴリズムや変分ベイズから始めることになった。
ベイズは抽象的な話が続いてパニックに陥ることも多かったので、このあたりの練習問題的なものを随時解いていた。

11日目: Diffusionモデルは簡単

Diffusionモデルの解説記事を読んだが、VAEに似ていたので簡単だった。Langevin dynamicsだけはよくわからなかったが、わからなくてもあまり問題はなかった。

12日目: 実装用の論文を読んだ

実装するにあたり、細かい情報も欲しいのでこの論文を参考にした。出発地点としては良い選択であったが、後で改良の為に別の論文も参考にした。

13~20日目: MNIST攻略

実装はかなり簡単で、MNISTから数字を生成できるようになった。余談ではあるが、MNISTのデータを取ってくる綺麗な方法がなくて、結局TensorFlowから取ってきた。トレーニング時間はかなり短く、おそらく数十分~数時間くらいだった気がする。

21日目: Open Imagesは難しい

上で実装したコードをOpen Imagesに適応したが、何かしらの意味ある画像を生成することはなかった。Unetの実装が悪かったのか。NNのサイズが小さかったのか。とりあえず、そのあたりを改善していくことにした。

22~25日目: Unetの実装を変えた

実装の参考にする論文をこっちに変えた。おそらくその論文に使われたであろうコードも見つけたので、Unetの実装はそこを参照した。

26~30日目: PyTorchの最適化をした

参考にした論文では、そこそこ大きいモデルを使っていたらしい。正直思ったよりもかなり小さかったが、それでもそこそこ大きいので工夫が必要だった。
試行錯誤の末、最終的には16GBのGPUを2つ使い、分割されたモデルの半分を一方のGPUに置き、もう半分をもう一方に置くことにした。それでもメモリ使用量がすごかったので、Gradient checkpointingを使った。
トレーニングはパイプライン方式を採用した。Unetをパイプラインでトレーニングするのは難しいかと思ったが、インプットからアウトプットまで一気に計算させずに、ネットワークの前半部分だけを全バッチに対してパイプラインで計算して、その後に後半部分をパイプラインで計算すればどうにかなった。

31~60日目: ターゲットを猫に絞る

モデルや訓練をいくら最適化してもOpen Imagesで動く気配はなかった。そこで、もしかしたら画像の種類が多いのかなと思った。確かに言われてみれば、GAN系の論文はベッドルームばかり生成していた。あれはベッドルームを生成するのが比較的簡単だからだったからなのかなと思った。
ということで猫の画像にターゲットを絞ることにした。結果は上澄みを集めて以下の感じ。

ただ、80%以上の確率で黒い画像やモヤみたいな物を生成していた。見栄えの観点から、画像サイズを64x64に設定していたが、訓練に数日かかってキツかった。32x32にしたり、ベッドルームの生成に切り替えてたら、もっと安定して綺麗に生成できていたのかもしれない。バグってる可能性を排除できないのもキツかった。
ここから改良しようとすると、あとはパラメーターを変えたり重課金するくらいしか思いつかず、学びが少ないかなと思い、ここで一旦終了することにした。
文字列から画像を生成するのも流行っているが、そちらもかなり似たアルゴリズムで実現可能らしいので、いずれ機会があれば実装してみたい。