Diffusers版DreamboothをVRAM 12GBのWindows PCで動かす

Kohya S.

10/29追記・さらに精度を上げたバージョンがVRAM 12GBで動作しますので、そちらをご覧ください。

概要

diffusers版のShivamShiriao氏のDreamboothは10/5時点で使用するVRAM容量が9.92GBまで削減されていますが、依存ライブラリの関係で残念ながらWindowsでは動きません。この記事ではWindowsでなるべく省メモリで動作させる方法を簡単に解説します。
Pythonで仮想環境を構築できるくらいの方を対象にしています。また細かいところは省略していますのでご容赦ください。

※タスクマネージャ等でGPU使用メモリが12GB以下に収まることは確認していますが、12GBの実機での確認はしておりません。動作報告をいただけると助かります。

※使用に当たっては自己責任でお願いいたします。

添付ファイルのライセンスは元ファイルに準拠します。ファイルを置き換えるときには念のため元ファイルを別名でコピーしておくことをお勧めします。

環境構築

適当な作業ディレクトリを用意して仮想環境を作ります。以下はvenvを使う例です。pipをバージョンアップしておきます。

E:\Workに作りました
python -m venv --system-site-packages venv_dbwin
venv_dbwin\Scripts\activate
python.exe -m pip install --upgrade pip

PyTorchとTorchvisionを入れる

diffusersを入れる前にPyTorchとTorchvisionを入れます（diffusersをインストールするとPyTorchもインストールされますがCPU版になるようです）。
PyTorch、TorchVisionの、お使いのCUDAのバージョンにあわせた最新版を入れてください（私は1.12.1 CUDA 11.3で検証）。以前のバージョンよりもわずかにメモリ使用量が少ないようです。

pip3 install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu113

diffusersと依存ライブラリを入れる

作業ディレクトリにShivamShiriro氏のdiffusersをcloneするか、zipでダウンロードしてきて展開します。

※将来的に当該リポジトリのバージョンアップでこの記事の方法では対応できなくなる可能性があります。
私は10/4のこちらのバージョンを使いました（requirements.txtにaccelerate==0.12.0の修正が必要です）。

こんな感じです

氏のDreamboothのページに書いてある「pip install git+https: //github.com/ShivamShrirao/diffusers.git」を行わず（スクリプトを変更したいため）、展開したdiffusersディレクトリでeditableモードでインストールします。

cd diffusers
pip install -e .

その後、examples/dreamboothディレクトリに移り、依存ライブラリをインストールします。

cd examples\dreambooth
pip install -U -r requirements.txt

その他のライブラリ、Windowsで動かすためのライブラリを入れる

StableDiffusionのcheckpointからdiffusersのモデルへの変換に必要なOmegaConfとpytorch_lightning、省メモリAttentionに必要なeinops、8bit Adamのbitsandbytesをインストールします。

pip install OmegaConf
pip install pytorch_lightning
pip install einops
pip install bitsandbytes

bitsandbytesはWindows用のDLLが提供されていないためそのままでは動きません。こちらのissueを参考に、有志の方がコンパイルされたWindows用DLLファイル、libbitsandbytes_cuda116.dllをこちらのページから落としてきます（CUDA 11.6用ですが、私はCUDA 11.3環境で動きました）。DLLを仮想環境のbitsandbytesのディレクトリにコピーします。

DLLを仮想環境内など適切なところにコピー

また先のissueを参考にbitsandbytesのcuda_setup/main.pyと、cextension.pyを書き換えます（書き換え済みファイルも添付してあります）。

main.pyを修正
main.pyの$$$のある行を追加
cextention.pyを修正
cextention.pyの$$$の前の行を、$$$のある行のように、str()を追加する

Pythonを起動し、import bitsandbytesとしてエラーが出なければOKです。

OK

0.34.0に対応した両ファイルを添付しておきます。

Attentionを省メモリ版に入れ替える

ShivamShrirao氏のDreamboothではxformersのmemory efficient flash attentionを使い省メモリ化を図っていますが、CUDAを直接使用しているためWindowsでは現時点では動きません。そのため、Memory Efficient Attention Pytorchのリポジトリにあるflash attentionのpure PyTorch実装を使います（速度はxformers版に比べるとかなり落ちるようです）。
そちらを組み込んだattention.pyを添付しますのでdiffusersのsrc/diffusers/modelsのattention.pyを置き換えてください。

ここに置きます

accelerateを設定する

diffusers版Dreamboothのページにあるようにacclerate configを実行します。

accelerate config
In which compute environment are you running? ([0] This machine, [1] AWS (Amazon SageMaker)): 0
Which type of machine are you using? ([0] No distributed training, [1] multi-CPU, [2] multi-GPU, [3] TPU [4] MPS): 0
Do you want to run your training on CPU only (even if a GPU is available)? [yes/NO]:
Do you want to use DeepSpeed? [yes/NO]:
Do you wish to use FP16 or BF16 (mixed precision)? [NO/fp16/bf16]: fp16

GPU 1枚で学習する場合、0、0、そのままエンター、そのままエンターでOKでしょう。
mixed precision設定は最も省メモリなのはfp16になります（bf16は動きませんでした）。

このままでShivamShrirao氏のtrain_dreambooth.pyが動くと思いますが、さらに省メモリ化し、またText Encoderも学習対象とした以下の改造版スクリプトtrain_dreambooth_mod2.pyをexamples/dreamboothに入れてください（前回の記事から省メモリ化しています）。

ここに入れます

※省メモリ化のため、ステップ当たりの学習回数が元のtrain_dreambooth.pyの半分になっています（対象の画像と正則化画像を同一のバッチではなく別のバッチに分割して学習するため）。元のdiffusers版やXavierXiao氏版とほぼ同じ学習を行うには、ステップ数を倍にしてください。
（shuffle=Trueのため厳密にはデータの順番が変わってしまいますが、学習には大きな影響はないと思います。）
※class_image（正則化画像）を使わない学習には対応していません。

以上で環境構築は完了です。

StableDiffusionのcheckpointのdiffusersモデルへの変換

Dreamboothを学習させる元モデルがStableDiffusionのckpt形式の場合、diffusersモデル形式に変換します。diffusersディレクトリで、以下のように入力すると変換できます。
checkpoint_pathにckptファイルを、original_config_fileにStableDiffusionのv1-inference.yamlを、dump_pathにdiffusersモデルの出力先ディレクトリを指定します。

こんな感じに置くと例えば次のようなコマンドラインになります
cd ..\..
python scripts\convert_original_stable_diffusion_to_diffusers.py --checkpoint_path ../models/model.ckpt --original_config_file ../models/v1-inference.yaml --scheduler_type ddim --dump_path ../models/diffusers_model
diffusers_modelに出力された

データの用意

学習させたい物体等の画像を数枚から数十枚、用意します（適切な枚数、内容についてはいろいろと研究されているようです）。サンプルにダ鳥獣ギ画さんのカエル、15枚を使います。
とりあえず作業ディレクトリのdata/frog_trainに置いた例です。

こんな感じです

正則化画像を用意します（正則化画像は学習させたい物体のクラス（カエル、少女など）の、学習させたい画像の特徴を持たない画像です）。StableDiffusionを用いて生成した画像（たとえば"frog"や"photo of frog"で）で構わないようです。
XavierXiao氏のDreambooth-Stable-Diffusionのページによると「生成される画像が非現実的な場合（たとえばクラスがmanやwomanの場合）はネットからダウンロードしてくるといい」そうです。
学習スクリプト内で自動生成されますが先に生成しておくといいでしょう。
data/frog_regularizationに90枚置きました。ShivamShriraoのページによると「たいていの場合は200~300枚でうまくいく」とのことです。

生成した画像です

正則化画像は学習用画像の整数倍をお勧めします。

学習の実行

ShivamShrirao氏のDreamboothのページを参考に、examples/dreamboothディレクトリでスクリプトを実行します。

GPUメモリを開けるためコマンドプロンプト以外のすべてのアプリを終了してください（タスクトレイのアイコンも含む）。タスクマネージャで使用メモリ量0.5GBを切るくらいになると、学習時のメモリ量が12GBに収まると思います。
さらにタスクマネージャも終了し、GPUメモリ使用量を別窓のコマンドプロンプトのnvidia-smiで確認するとよいです。

コマンドラインは、たとえば先のようにdiffusersのモデルを変換して用意した場合、以下のようになります。

cd examples\dreambooth

accelerate launch --num_cpu_threads_per_process 8 train_dreambooth_mod2.py 
  --pretrained_model_name_or_path=../../../models/diffusers_model 
  --instance_data_dir=../../../data/frog_train 
  --class_data_dir=../../../data/frog_regularization
  --output_dir=../../../models/dreambooth_diffusers_model 
  --with_prior_preservation --prior_loss_weight=1.0 
  --instance_prompt="sks frog" 
  --class_prompt="frog" 
  --resolution=512 
  --train_batch_size=1 
  --gradient_accumulation_steps=1 --gradient_checkpointing 
  --use_8bit_adam 
  --learning_rate=1e-6 
  --lr_scheduler="constant" 
  --lr_warmup_steps=0 
  --num_class_images=1 
  --mixed_precision="fp16"
  --max_train_steps=1600

（accelerate以降は改行を入れず一行で入力してください。エディタやメモ帳に書いてコピペするのをお勧めします。）

出力先にはmodels/dreambooth_diffusers_modelを指定しています。num_cpu_threads_per_processにはCPUコア数を指定するとよいようです。

class_imageは生成済みのため、num_class_images=1としています。スクリプト内で生成する場合、その枚数を指定します。

省メモリ化のためmixed_precision="fp16"、およびgradient_checkpointing を指定します。
また前述のとおりmax_train_stepsを1600とします。
学習率learning_rateは、diffusers版では5e-6ですがXavierXiao版は1e-6ですのでここでは1e-6を指定しています。

instance_promptは"sks XXX"、class_promptは"XXX"が学習結果が良いようです。

学習が終わると指定されたディレクトリにdiffusers形式のモデルが出力されます。

学習完了

diffusersモデルのStableDiffusionへの変換

(10/6) ShivamShrirao氏のリポジトリに追加されました。

ちょうど今日(10/5)に当該の変換を行うスクリプトがdiffusersのリポジトリに追加されましたが、ShivamShrirao氏のリポジトリには取り込まれていないため、https://github.com/huggingface/diffusers/blob/main/scripts/convert_diffusers_to_original_stable_diffusion.pyから落としてきて、diffusers/src/scriptsに入れます。

やっと変換できるようになったぞ

（--helpを指定すると表示されますが）model_pathに変換元（つまりDreamboothで学習した）diffusersモデル（ディレクトリ）を、checkpoint_pathに出力先のStableDiffusionのcheckpointファイルを指定します。

cd ..\..

python scripts\convert_diffusers_to_original_stable_diffusion.py 
    --model_path ../models/dreambooth_diffusers_model 
    --checkpoint_path ../models/dreambooth_frog.ckpt

（python以降は一行で入力します。）

指定したcheckpointファイルが出力されているはずです。

できた(･∀･)

あとはStableDiffusionで"sks frog"などとして画像を生成してみてください。

McDonald's感がないがまあまあ

メモリ使用量と速度、精度の向上

（学習データ、学習率、promptなどのほうが影響は大きいですが）メモリに余裕があればオプションを変えることで精度向上が期待できます。また速度も大きく変わります。

まずmixed precisionではなくfloat32で学習したほうが、データ等にもよりますが良くなる可能性があります。ただしメモリ使用量は増加し、学習速度も低下します。
accelerate configをもう一度実行して最後の質問でNOを選んでください。また学習用スクリプトのmixed precisionの設定を削除してデフォルト（float32）としてください。
一般的にはかなり省メモリに効くのですが、なぜかこのスクリプトではほぼ減らないので、NO指定はわりとお勧めです。

次にバッチサイズを増やすと一般的に精度が良くなるようです。train_batch_sizeを増やしてください。ただしリニアにメモリ使用量が増えます。
実際に学習に使用されるデータの数は「バッチサイズ×ステップ数」なのでバッチサイズを倍にしたらステップ数は半分でほぼ同じだけ学習が進みます（実際にはちょっと+αするといいかもしれません1->8なら1600->300など）。

gradient_accumulation_stepsはメモリ使用量をさほど増やさず、バッチサイズを増やすのとほぼ同じ効果がありますが、速度は低下します。またステップ数への影響はないはずです（元のステップ数と同じだけ必要）。

gradient_checkpointingを外すと速度が大きく向上しますが、メモリ使用量は増えます。

use_8bit_adamを外し、通常のoptimizerを使うと精度が向上する可能性がありますが、速度は低下し、メモリ使用量も大きく増えます。

（VRAM 24GBの場合、fp16なし、gradient_checkpointingなし、8bit Adamでバッチサイズ8まで増やせます。）

いいなと思ったら応援しよう！