Kohya版のLoRAをDockerで訓練する
LoRAのデファクトスタンダードとなりつつあるKohya版のLoRA学習スクリプトをDockerで動かしてみました。Diffusersとの連携も可能で、DiffusersオフィシャルのLoRAよりも高品質な学習が期待できます。PyTorch2の環境でも動かすことができました。
目次
はじめに
Civitaiなどで配布されているLoRAでデファクトスタンダード形式となりつつあるKohya版のLoRAをWebUIなしで訓練します。ベースのモデルはDiffusers形式でもOKで、Diffusersの連携もできそうです。
LoRAの訓練コードの違い
この世界には画像生成モデルのLoRA訓練コードが多数溢れています。いくつか例を挙げますと、
- Hugging FaceのDiffusersライブラリに登録されている訓練コード
- dreamboothとtext_to_imageの2種類
- Hugging FaceのPEFTライブラリに登録されている訓練コード
- 基本はDiffusers版の古いバージョンが入っているだけです
- Kohya版のLoRA訓練コード
- WebUI界隈でデファクトスタンダードとなっているもの
PEFTはDiffusersの古い版だったので除外してもいいと思います。DiffusersかKohya版かのどちらかになると思いますが、結論からいうとKohya版のほうが高度な訓練しています。わかりやすい点をいくつか挙げると、
- Diffusers版はData Augmentationが簡単なフリップ程度であるのに対し、Kohya版はData Augmentationの専用のライブラリのAlbumentationsを使っている
- bitsandbytesというライブラリを使うことで、AdamWや最近出たLionなどなどの発展的なオプティマイザの8bit訓練ができる
- 公式解説が日本語記事
- safetensorsで吐き出しできる
などです。ただデメリットとして、使っているライブラリが多い分依存関係がややこしいというのがあります。そこを整理してDockerfileにしたというのが今回の記事の目的です。
ディレクトリ構成
- docker
- Dockerfile
- requirements.txt
- workdir
- data
- lora_output
- dataset.toml
このような構成にします。dockerフォルダにはDockerビルドに必要なもの、workdirはDockerイメージにマウントさせる作業ディレクトリです。データ類もこちらに入れます。
Dockerイメージ
公式の環境構築だとPyTorch1系統で書いていますが、PyTorch2にしてもLoRAの訓練程度だと動きました。
Dockerfile
最後のGit cloneはブランチを指定していて、バージョン0.6.4のものをCloneするようにしています。ここが不要でしたら-b以降をとってください。
ベースのイメージはPyTorch2.0.0としました。timmの関係上、現在(2023/6/3)はPyTorch2.0.1は対応していませんでした。
requirements.txt
requirements.txtのバージョンは公式と一部変えています。キャプション生成のためだけにTensorFlow入れたくなかったので、外しました。必要だったら入れてください。
普通にWSLなどでDockerをビルドすればOKです。カレントディレクトリは「docker」で
データセットの準備
東北ずん子のLoRAの学習データを例にとります。これは最初からキャプションデータが用意されているので簡単です。
https://drive.google.com/drive/folders/1NIZcBRvr5i8YfPsPYwvVMC7SoH-cWLIk
今回は「01_LoRA学習用データA氏提供版背景白」から「kiritan」を対象とします。
ダウンロードするとkiritanフォルダ直下にpngとtxtが
- kiri(1).png
- kiri(1).txt
- kiri(2).png
- kiri(2).txt
のように入っています。これを「data」フォルダ直下に移します。
次にKohya版LoRAの訓練に必要なtomlを記述します。これはyamlの拡張版で実際に記述するものはただのテキストです。
このように記述しました。「image_dir」の部分だけ注意が必要で、Dockerコンテナ内でのパスになります。このケースでは、ローカルで「workdir」というフォルダを、「/workdir」というパスにマウントさせるので、「/workdir/…」と書きます。
コンテナの起動
ビルドが終わったらコンテナを起動します。カレントディレクトリを「docker」の一個上にして、
次にコンテナ内に移動するので、スクリプトのパスに移動します。
accelerateの設定をします。面倒だったらローカルで定義したyamlをマウントしてもいいかもしれません。
このようにしました。GPU IDの指定はallでもいいのですが、2個GPUがあって2個目のみ使いたかったので「1」と指定しました。
LoRAの訓練コマンド
このようにします。改行するときはバックスラッシュをお忘れなく!(これを忘れて設定が読み込まれずに2時間ぐらいハマりました)
オプティマイザは公式例にあるような「AdamW8bit」だとエラーになってしまったので、「Lion8bit」にしました。こっちのほうが高度なオプティマイザなのでこれでいいかなと思います。AdamWを使いたい場合は8bitやめて「AdamW」だと動きました。
ここがKohya版の強いところなのですが、「pretrained_model_name_or_path」はDiffusersのモデルIDを指定しても動きます。しかもそれはローカルにある必要はなく、Hugging Faceからよしなにダウンロードしてくれます。訓練コードのモデルの読み込みがDiffusersなので、このへんは得意です。
2080Tiで40分ぐらいで終わりました。もっとエポック数減らしてもいいかなと思います。
このようによく見慣れたLoRAができあがりました。Dockerでディレクトリごとマウントしているので、Dockerで吐き出したファイルを、ローカルからも使えます(永続化されます)。
推論してみる
safetensorsのLoRAと扱いは同じになりますので、オンザフライでマージさせます。やり方はこの記事参照。
αをかなり落とさないといけなかったので、もしかしたらかなり過学習してしまったかもしれません。エポック数をもっと減らしても良かったかもしれませんね。
推論結果
推論時のプロンプトは「kiritan, look at viewer, 1girl, best quality」です。
SD1.5で訓練し、SD1.5で推論
α=0.2
元がアニメモデルでないのでこんな感じかなという具合です
SD1.5で訓練し、SomethingV2で推論
α=0.2
アニメモデルのSomethingV2_2で推論した結果です。結果はかなりうまくいっていて、SD1.5で訓練して別モデルに移植してもうまくいっているのが確認できます。
SomethingV2で訓練し、SD1.5で推論
α=0.3
こちらはLoRAの学習をアニメモデル(SomethingV2)で行い、SD1.5で推論した場合です。アニメ→きりたんだけでは、アニメ方向へのタスクの学習が行われず、LoRAの適用結果は実写のままでした。αを大きくすると画像が破綻しました。
SomethingV2で訓練し、SomethingV2で推論
こちらはLoRAの学習モデルと推論モデルが一致した結果です。これは分布シフトがないので一番きれいな結果になります。ここまでできるKohya版のLoRAすごいですね。
Shikoan's ML Blogの中の人が運営しているサークル「じゅ~しぃ~すくりぷと」の本のご案内
技術書コーナー