話題のStable Diffusionは、 LAION-5Bデータセットを用いて学習された、CLIPのテキスト埋め込みを条件としたLatent Diffusionモデルである。
Stable Diffusionの仕組みを理解したいと思っており、そのために、まずは比較的小さいデータセットでLatent Diffusionの学習を試してみた。
GitHubのLatent Diffusionでは、CelebA-HQ、LSUN、ImageNetの3つのデータセットでの学習設定が用意されている。
ここでは、画像数が比較的小さいCelebA-HQで学習を試してみた。
環境構築
NVIDIAのPyTorchのイメージを使用する。
docker run --gpus all --shm-size=32g -it nvcr.io/nvidia/pytorch:21.10-py3
ライブラリインストール
必要なライブラリをインストールする。
pip install transformers omegaconf pytorch_lightning==1.4.2 torchmetrics==0.6.0 opencv-python==4.5.5.64 test-tube einops albumentations
バージョンを指定しているライブラリはこのバージョンでないと動かない。
なお、ImageNetの学習を行う場合は、albumentationsを使用してTorrentからデータセットがダウンロードされるが、ここでインストールされるバージョンだとエラーが起きる。
albumentationsを最新にすると、pytorch_lightningと競合するため、pytorch_lightningのバージョン指定を外すとインストールできる。
だだし、学習時にpytorch_lightningのバージョン不整合でエラーになるため、ダウンロードだけ行ってバージョンを戻すとよい。
もしくは、ImageNetのダウンロードは別の方法で行いキャッシュディレクトリ(~/.cache/autoencoders/data/ILSVRC2012_trainとILSVRC2012_validation)にシンボリックリンクを張ると良い。キャッシュディレクトリには.readyファイルが必要である。
albumentationsをインストールするとopencv-python-headlessがインストールされてしまうので、一旦アンインストールしてopencv-pythonを入れ直す。
pip uninstall opencv-python opencv-python-headless pip install opencv-python==4.5.5.64
OpenCVの実行に必要なUbuntuのパッケージをインストールする。
apt update apt install libgl1-mesa-dev
taming-transformersインストール
taming-transformersをGitHubからcloneしてインストールする。
git clone https://github.com/CompVis/taming-transformers.git cd taming-transformers pip install -e .
CelebA-HQデータセットの準備
CelebAデータセットと、CelebA-HQの差分データをダウンロードする。
CelebAは、公式ページのGoogle Driveのリンクからダウンロードする。
必要なのは、Img/img_celeba.7z/*.7zと、Anno/list_landmarks_celeba.txtである。
CelebA-HQは、taming-transformersのREADMEに説明がある通り、 PGGANのREADMEに従って準備する。
PGGANのREADMEにあるGoogle Driveのリンクからimage_list.txtと、deltas*.zipをダウンロードする。
zipを直接ダウンロードするとサードパーティーcookieのエラーが起きる場合があるので、一度自分のGoogleドライブにコピーを作成するとよい。
なお、手動でダウンロードすると時間がかかるのでrcloneを使用した。
.npy形式に変換
Latent Diffusionは、CelebA-HQのデータセットが.npy形式になっていることを要求する。
PGGANの説明に従って準備するように書いているが、PGGANのREADMEには.npyにする方法は書いていない。
ネットでいろいろ調べたら、.npyに変換するコードを公開しているレポジトリが見つかった。
https://github.com/mazzzystar/make-CelebA-HQ
こちらで公開されているスクリプトを使用して変換を行った。
ダウンロードしたデータセットを格納するディレクトリ構成は、READMEに従う。
python make_HQ_images.py ./
Latent Diffusionの準備
GitHubからLatent Diffusionをcloneして準備を行う。
/workにcloneする例)
cd /work git clone https://github.com/CompVis/latent-diffusion.git cd latent-diffusion pip install -e .
CelebA-HQの学習に必要なファイルをダウンロードする。
cd data wget https://github.com/CompVis/taming-transformers/raw/master/data/celebahqtrain.txt wget https://github.com/CompVis/taming-transformers/raw/master/data/celebahqvalidation.txt
.npy形式に変換したCelebA-HQデータセットのディレクトリへのシンボリックリンクを作成する。
.npyが/work/celeba/celebahqに格納されている場合の例)
ln -s /work/celeba/celebahq celebahq
学習
CelebA-HQのLatent Diffusionモデルの学習を行う。
Latent Diffusionの学習は、オートエンコーダーと、拡散モデルの2つのモデルの学習に分かれている。
オートエンコーダーの学習済みモデルダウンロード
オートエンコーダーの学習方法も記載されているが、ここでは、オートエンコーダーは学習済みモデルを使うことにする。
cd /work/latent-diffusion wget https://ommer-lab.com/files/latent-diffusion/vq-f4.zip unzip vq-f4.zip mv model.ckpt models/first_stage_models/vq-f4/
設定変更
デフォルトの設定だと、GPUがV100だとGPUメモリが足りず、CUDA out of memoryが発生する。
そこで、バッチサイズを変更する。
/work/latent-diffusion/configs/latent-diffusion/celebahq-ldm-vq-4.yaml
の以下の2箇所を変更する。
変更前：batch_size: 48
変更後：batch_size: 32
変更前：max_images: 8
変更後：max_images: 6
学習実行
V100 8GPUで学習する場合の例)
python main.py --base configs/latent-diffusion/celebahq-ldm-vq-4.yaml -t --gpus 0,1,2,3,4,5,6,7
学習結果
5000バッチ(51エポック)ごとに、評価が行われ、結果がlogsに出力される。
5000バッチ(51エポック)学習時点で出力された画像は以下の通り。
ノイズ画像から再構成
Diffusion過程でノイズを加えた画像
再構成した画像
潜在空間からランダムにサンプリングした画像
ノイズ除去過程
Inpainting(修復)
入力
マスク
修復画像
学習時間
V100 8GPUを使用して、5000バッチ(51エポック)の学習に3時間23分かかった。
51エポックではまだ画像品質は高くない。
