diffusers版DreamboothをText Encoderまで学習する&Windows GPU VRAM 12GBで動かすための最低限の情報

Kohya S.

のちほどもう少し詳細に書きますが、取り急ぎまとめておきます。わかっている方向けの記事になります。
※無保証です。自己責任でお願いします。

10/3時点のShivamShrirao氏のdiffusers版Dreamboothを元にしています（当該バージョンはこちら）。
添付ファイルのライセンスは元ファイルに準拠します。ファイルを置き換えるときには念のため元ファイルを別名でコピーしておくことをお勧めします。

Text Encoderまで学習する

train_dreambooth.pyを変更してtext encoderを学習対象にします。以下の変更済みのtrain_dreambooth_mod.pyを使ってください。Colabでも使えるかもしれません（未検証です）。動作報告をいただきましたのでColabでも動作します。

（2022/10/4 元リポジトリの10/4の変更を取り込みファイルを変更しました。マルチGPUで分散学習するとき以外は学習に影響ありません。）

  • クラス画像（regularization image）はすべて生成済みの状態でしかテストしていません（num_class_images=1）

  • diffusers版の学習率は5e-6ですがXavierXiao版は1e-6です。

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

unetだけを学習対象としていたオリジナル版から、unet+text_encoderのwrapper modelを作りそれを学習対象とするように変更しています。

WindowsでVRAM 12GBで動かす

環境整備の注意

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

8-bit AdamW optimizerを入れる

pip install bitsandbytesでインストールできますが、Windows用のDLLが提供されていないためそのままでは動きません。こちらのissueを参考に、有志の方がコンパイルされたWindows用DLLファイル、libbitsandbytes_cuda116.dllをこちらのページから落としてきます（CUDA 11.6用ですが、私はCUDA 11.3環境で動きました）。DLLを仮想環境のbitsandbytesのディレクトリにコピーします。
また先のissueを参考にbitsandbytesのcuda_setup/main.pyと、cextension.pyを書き換えます。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を置き換えてください。

GPUメモリを開けるためコマンドプロンプト以外のすべてのアプリを終了し（タスクトレイのアイコンも含む）、ShivamShrirao氏のページを参考に適宜オプションを設定してtrain_dreambooth.pyまたは先に添付したtrain_dreambooth_mod.pyで学習してください。
GPUメモリ使用量は別窓のコマンドプロンプトのnvidia-smiで確認できます。
参考までに手元で使用量11.8GBで動作したコマンドラインを書いておきます。

SET MODEL_NAME=../../diffusers_wd_1_2
SET INSTANCE_DIR=../../../data/frog_train
SET CLASS_DIR=../../../data/frog_regularization
SET OUTPUT_DIR=db_model
accelerate launch --num_cpu_threads_per_process 8 train_dreambooth_mod.py --pretrained_model_name_or_path=%MODEL_NAME% --instance_data_dir=%INSTANCE_DIR% --class_data_dir=%CLASS_DIR% --output_dir=%OUTPUT_DIR% --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 --max_train_steps=800 --mixed_precision="fp16"

num_cpu_threads_per_processはCPUコア数を指定すればよいようです。

この記事が気に入ったらサポートをしてみませんか？