TadaoYamaokaの開発日記

個人開発しているスマホアプリや将棋AIの開発ネタを中心に書いていきます。

Latent Diffusionの学習を試す その2

昨日の記事の続きで、CelebA-HQデータセットでLatent Diffusionの学習を続けた場合に生成画像の質が上がるか確認した。

昨日の記事で、5000バッチ(51エポック)時点のモデルで、潜在空間からランダムにサンプリングして生成した画像は以下の通りであった。

学習を続けると、9649バッチ(98エポック)で評価損失が最小となり、15349バッチ(156エポック)あたりで訓練損失が発散して学習できなくなった。


※グラフの色が2色なのは途中から再開したためで特に意味はない

潜在空間からランダムにサンプリングした画像

5000バッチおきに評価された際に出力された画像は以下の通り。

1000バッチ(102エポック)時点

1500バッチ(153エポック)時点

2000バッチ(204エポック)時点

1000バッチ(102エポック)時点で、昨日の出力結果より質が上がっていることが確認できる。
2000バッチ(204エポック)時点では、損失が発散しておりノイズだけの画像になっている。

ベストモデルで生成

評価損失が最小となった9649バッチ(98エポック)時点のモデルで生成を行った。

学習済みモデルから生成するには、scripts/sample_diffusion.pyを使用する。

チェックポイントが保存されたディレクトリ(例:logs/2022-09-09T11-02-28_celebahq-ldm-vq-4/checkpoints)にconfig.yamlが必要になるため、configs/latent-diffusion/celebahq-ldm-vq-4.yamlからシンボリックリンクを張る。

以下のようなコマンドで生成を行う。

python scripts/sample_diffusion.py -r logs/2022-09-09T11-02-28_celebahq-ldm-vq-4/checkpoints/epoch=000098.ckpt -n 10

-nには、生成する枚数を指定する。
DDIMステップ数はデフォルトで50になっている。

DDIM steps=50

-cオプションでDDIMステップ数を500にすると以下のような画像になる。

DDIM steps=500

DDIMステップ数を増やすと、よりはっきりした画像になる。

白っぽい画像や、真っ黒な画像も含まれおり、潜在空間上のランダムにサンプリングした点からは生成が必ず上手くいくとも限らないようだ。

まとめ

CelebA-HQデータセットで、Latent Diffusionを条件なしで学習を行ったところ、98エポックで評価損失が最小となるモデルが学習できた。
また、156エポック以上学習すると訓練損失が発散し、生成される画像はノイズのみの画像となった。

98エポック学習したモデルで生成した画像は、生成がうまく行ったものは質の良いが生成できている。
潜在空間をランダムにサンプリングした点では、白っぽい画像や真っ黒な画像が含まれる。

学習を試せたので、実行しながらソースコードを理解していけるようになった。
論文と照らし合わせながら仕組みの理解を進めていきたい。
また、独自のデータセットでの学習も試してみたい。