kohya版のLora(DreamBooth)による追加学習をsd-scriptsのとしあきWikiのBatファイルを使用して行う際、どの程度のバッチサイズ・解像度(resolution)までできるのかを検証しました。
各解像度での最大バッチサイズでの学習速度s/itも参考値として載せました。
当方のVRAM24GB・RTX3090の場合ですが、何か参考になれば幸いです。
導入等はこちらを参照してください。
検証対象の箇所
https://wikiwiki.jp/sd_toshiaki/LoRA#x5a23f01のバッチファイルを使わせていただいていますが、ポップアップGUI版でも限界は同じです。
あくまでもLoRAを実行するにあたって、各種設定を入力するのがGUIかバッチファイルに書き込むのかのか違うだけで、実行対象となるハードウェアやプログラム自体は共通なため。
関係があるかわかりませんが一応書いておくと、学習設定は
学習元画像12枚
正則化画像12枚透明PNG
繰り返し5
epoch1
dim=64
です。
最大解像度とバッチサイズが知りたいだけなので、epoch1でご勘弁ください。
また、今後のアプデなどで今の結果が通用しなくなったり、試行回数が1回ゆえのハズレ値を取り上げたりしている可能性をご了承ください。
rem =========================================================================
rem VRAMなど環境に影響されるパラメーター 余裕があれば増やすと速度や精度が改善
rem CPUスレッド数 CPUのコア数がいいらしい
rem 増やすと生成速度が上がる代わりにメインメモリの消費が増えるので
rem メモリ32GBの人はCPUのコア数よりも減らした方が安定するかも
set cpu_thread=20
rem データローダーのワーカー数
rem デフォルトは8、減らすとメインメモリの使用量が減り、学習時間が増加
set workers=8
rem バッチサイズ: 増やすと計算が早く終わるがVRAM消費が増える
set train_batch_size=6
rem 学習素材の解像度: 大きくすると細部まで学習するが消費VRAMが増える
rem VRAM12GBなら768くらいまで増やせる
set resolution=768,768
赤文字の部分を変えていきます。
とはいえresolutionは基本縦横同じなので、512×512と768×768の2通り。
その2つに対してtrain_batch_sizeがいくつまで行けるか検証です。
RuntimeError: CUDA out of memory. Tried to allocate 340.00 MiB (GPU 0; 24.00 GiB total capacity; 22.73 GiB already allocated; 0 bytes free; 22.89 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
バッチサイズが大きすぎると上のようなエラーがでるので、それが出ないギリギリの数値を探るのが今回の検証です。
解像度(resolution)512×512の場合:LoRAの速度と最大バッチサイズは?
512×512ならBatch Sizeは16が限界です。
速度は[01:28<00:00, 9.83s/it, loss=0.116]。

かなり高速ですね。
キャプションの妥当性を検討する際などは、512で試行を爆速で繰り返すのがおすすめです。
解像度(resolution)768×768の場合:LoRAの速度と最大バッチサイズは?
768×768ならBatch Sizeは7が限界です。
速度は[01:51<00:00, 5.88s/it, loss=0.0624]。

Batch Size 6の時は22GBだったので、単純な掛け算でVRAM使用量が増えるわけではなさそうです。
768×768では、画素数が512×512のときの2.25倍で
最大バッチサイズは16÷7=2.285 ということで大体計算通りです。

速度と品質を両立したいときは768くらいがいいですね。
追加検証:解像度(resolution)1024×1024の場合:LoRAの速度と最大バッチサイズは?
1024×1024ならBatch Sizeは3が限界です。
速度は[03:17<00:00, 4.70s/it, loss=0.0636]。
画素数が768×768時の1.8倍程度のため、理論上最大バッチサイズは7÷1.8=3.89です。
案の定バッチサイズ4ではメモリ不足でした。

単位がデカいのでBatch Sizeに乗り切らないのは仕方ないとはいえ、1024だとVRAMがあまりがちなので、少々もったいないですね。
ぶっちゃけSD v2.0以降モデルをベースにしていて、超高品質を追い求める場合以外使い勝手が悪そうです。残念ながら当然遅い。
現状の大多数のモデルは512ベースで一部の新しいモデルが768な程度のため、1024というデカい画像で学習しても、時間に見合う効果を期待できないかもしれない。
仮に1024でやるなら512で試行錯誤した結果を元に、どっしりと腰を据えて深夜ぶん回す場合だけにしましょう。
まとめ LoRAの解像度ごとの最大バッチサイズと速度(s/it)一覧
- 512×512 : batch size 16, 9.83s/it
- 768×768 : batch size 7, 5.88s/it
- 1024×1024 : batch size 3, 4.70s/it
各解像度(resolution)での速度と最大バッチサイズ(train_batch_size)は上記の通りです。
VRAM24GB環境での参考値としてお考えください。
その他検証記事などもご覧ください。