見出し画像

Waifu diffusionの追加学習をColabのT4で行う。（使用VRAM13.6GB）※dreamboothじゃないよ

gcem156

　先日Waifu diffusionの追加学習についての記事を投稿しましたが、ColabのT4でできたのでやり方を紹介します。

追記：haruさんがAMPに対応できるようにソースコードを修正していたので、やり方をちょっと変えたよ。リポジトリの内容も下の記事に合わせて変えちゃうから同じようにはもうできないよ。4日もcloneされてなかったしもうええやろ。

追加学習の仕方

　Colabの標準GPU（T4）、ランタイムハイメモリで動作確認しています。

　まずはNVIDIA apexをインストール。

%cd /content
!git clone https://github.com/NVIDIA/apex
%cd apex
!pip install -v --no-cache-dir ./

　ドライブにマウントしてgit cloneします。ドライブは無料版だと容量が足りないかもしれません。ドライブを使わなくてもできるとは思いますが、完成した学習済みモデルを保存するのが面倒になります。またドライブを利用しないとColabの起動のたびにダウンロードが必要になります。

%cd /content/drive/MyDrive/
!git clone https://github.com/laksjdjf/waifu-diffusion.git

　必要なライブラリをインストール。

%cd /content/drive/MyDrive/waifu-diffusion
!pip install -r requirements.txt

　ランタイムを再起動せよという文章が出てくるので、RESTART RUNTIMEを押します。カレントディレクトリが元に戻るので再度cd。

%cd /content/drive/MyDrive/waifu-diffusion

　logs/originalディレクトリを作って、Waifu diffusion v1.3の学習済みモデルをそこにダウンロードします。

!mkdir -p logs/original
!wget -P logs/original https://huggingface.co/hakurei/waifu-diffusion-v1-3/resolve/main/wd-v1-3-float16.ckpt

 このままだとエラーが(多分)でるのでちょっと書き換えます。

import torch
ckpt = torch.load("logs/original/wd-v1-3-float16.ckpt")
del ckpt["state_dict"]["model_ema.decay"]
del ckpt["state_dict"]["model_ema.num_updates"]
ckpt["optimizer_states"] = []
torch.save(ckpt,"logs/original/wd-v1-3-float16-del.ckpt")
del ckpt

　学習用データを用意します。danbooruの画像をスクレイピングしたい人は、山岡さんのブログを参考にしてください。

学習用データについての説明はここ。

　学習の準備が整ったので学習開始です。
※configファイルにflip_p: 0.0という設定を入れていますが、左右対称のキャラクターを学習させたい場合、必要ないので消してください。

!python main.py -n "aesthetic" --resume_from_checkpoint logs/original/wd-v1-3-float16-del.ckpt --base ./configs/stable-diffusion/v1-finetune-4gpu.yaml -t --no-test --seed 25 --scale_lr False --gpus 0,

　モデルは logs/inputなんちゃら/checkpointsに入っています。
以下の操作でいらないデータを消して2GBくらいにできます。

!python scripts/prune.py -I <ckptのパス>

こんな感じで画像生成できます。

!python scripts/txt2img.py --ckpt <ckptのパス> --prompt "プロンプト" --n_iter 4

1エポックやってみた結果。2646ステップで46分かかりました。

　一応追加学習はできてそうですね。

VRAM節約方法の解説

　解説といっても私自身完全に理解しているわけではなく、何だかよく分からないけど動いたからヨシ！という状態なので参考程度にしてください。具体的な変更点はcommitsを見ればわかると思います。
　ちなみにほとんどは以下のdreamboothでVRAM節約するやつでもやっていることです。

EMAを無効化

　EMAとは重みの指数平滑移動平均のことです。株価みたいにギザギザしている時系列データをなだらかにして、全体的なトレンドを読み取るために使われるやつです。モデルの学習も株価のように振動して安定しないので、画像生成するときにはEMAを使います。有効にするとそれぞれの重みに対して、EMAのデータを入れておく必要があるので、必要データ量が増えます。

8bit Adamを使う

　Adamとはモデルを最適化してくれる関数の一種です。以下の記事がスーパーわかりやすいです。

　この記事にある通り、Adamでは勾配のEMAと勾配の二乗のEMAの2つが必要なので、重みの数に対して2倍のデータ量が必要です。8bit Adamはようするに精度を8bitにすることでデータ量を節約するということです。

Mixed Precisionを使う

　デフォルトではデータを32bitで扱っていますが、一部のデータを自動的に16bitにしてくれるNVIDIAさんの便利な機能、AMPを利用しました。

ddpを無効

　複数GPUを利用するとき使う手法らしいんですが、デフォルトで有効になっています。1GPUで有効にしているとどうなるかよく分からないのですが、無効にしたらcuda out of memoryが起こらなくなったのでヨシ！

benchmarkを無効、move_metrics_to_cpuを有効にする

　Trainerのドキュメント見てVRAMの節約に関係ありそうな説明文があったのでやってみました。意味は何となくしかわかりません。

おわりに

　男の娘にちゃんとち○ち○を生やすプロンプトを発見できた方がいらっしゃれば、教えてください。
　

追記：

amp_levelをO3にすると11GBでできるみたいです。最初は9.2GBくらいだけどモデルを一回セーブすると11GBになります。出来上がったモデルがfp32で保存されるのが原因のような気がする。