Masked loss を使用したLoRA学習（SDXL)

2024年4月12日 01:57

はじめに

kohya_ss/sd-scripts (https://github.com/kohya-ss/sd-scripts) に PR(Pre Release?) として Masked loss (https://github.com/kohya-ss/sd-scripts/tree/main?tab=readme-ov-file#about-masked-loss) の機能が追加されました。

どんな機能かというと

マスクロスについて

各学習スクリプトでマスクロスをサポートしました。マスクロスを有効にするには --masked_loss オプションを指定してください。

機能は完全にテストされていないため、不具合があるかもしれません。その場合は Issue を立てていただけると助かります。

マスクの指定には ControlNet データセットを使用します。マスク画像は RGB 画像である必要があります。R チャンネルのピクセル値 255 がロス計算対象、0 がロス計算対象外になります。0-255 の値は、0-1 の範囲に変換されます（つまりピクセル値 128 の部分はロスの重みが半分になります）。データセットの詳細は LLLite ドキュメントをご覧ください。

kohya-ss/
sd-scripts
マスクロスについて

なるほど。よくわかりません。
期待するのは、マスクしたところだけ学習してくれること。例えば、白背景とか、背景切り抜きのキャラ画像をマスクと共に学習させると、背景は学習させずにキャラクターだけ学習してくれるようなユースケースを想定しています。

とりあえず、やってみましょう。

マスクイメージの生成

元画像と同じ名前で、マスクイメージを作成します。同じ名前なので元画像を上書きしちゃわないように注意…

ちまちま手＋ペイントソフトで作成しようと思ったのですが、ずんだもん画像は12枚とはいえ正直面倒だったので、マスクイメージ生成の為の ComfyUI のワークフローを作成しました。

Make mask image workflow

マスクイメージと反転マスクイメージ（白黒反転）のイメージを作ってくれます。どっちを使うのかわからなかったので両方作りましたが学習対象を白とするマスクで良いようです。

これを元画像とは別のフォルダに入れておきます。

設定ファイル(toml形式)の作成

kohya_ss/sd-scripts (以下学習スクリプト)は、通常、教師画像を train\<繰り返し回数>_<学習対象名> フォルダ（前回の例だと、train\13_zundamon ）に入れて学習するのですが、この形式だと同じファイル名のマスクデータを配置する場所がありません。
そのため、toml 形式の設定ファイルに元画像の位置と、マスク画像の位置を記載して指定します。

[general]
    caption_extension = ".txt"
[[datasets]]
    resolution = 1024
    batch_size = 12
    num_repeats = 13
    [[datasets.subsets]]
        image_dir = "/mnt/c/workspaces_c/kohya_lora_gui/zundamon/train/13_zundamon"
        conditioning_data_dir = "/mnt/c/workspaces_c/kohya_lora_gui/zundamon/conditioning_data_dir/zundamon"

この toml ファイルは階層構造になっており、親の階層の設定は子要素は引き継ぐ形になっています。kohya-ss / sd-scripts のドキュメントのデータセットの準備では caption_extension = ".txt" の行は [[datasets.subsets]]の下に記載するようになっていますが、.txt ファイル以外のキャプションのファイルを利用することはないなとおもって [general]（全体への指示）に位置をずらしています。
[[datasets]] という項目の下に、画像の解像度、バッチサイズ、繰り返し回数を入れます。
更にその下の要素の datasets.subsets に今までの画像とキャプション（.txtファイル）があるフォルダを image_dir として指定して、マスクイメージを入れたフォルダを conditioning_data_dir として指定します。
私の場合は WSL を使用したのでパス名が WSL の Linux 形式になっていますが、 Windows環境で実行する場合には、 C:\workspaces_c\ ~ というようなWindowsのパス形式で記載すれば大丈夫だと思います（多分）。
詳しくは、kohya-ss / sd-scripts のドキュメントの docs/config_README-ja.md にてご確認ください。

学習開始！

いよいよお楽しみの学習開始です。
今までのパラメータではMasked loss を使った学習をしてくれないので、コマンドから以下を変更します。

--masked_loss パラメータを追加します。
--dataset_config パラメータとして作成したパラメータファイル (tomlファイル)を指定します
今までの元画像指定だった --train_data_dir "フォルダ名" の記述を削除します。

--masked_loss --dataset_config "/mnt/c/workspaces_c/kohya_lora_gui/zundamon/zundamon.toml"

追加パラメータ例

これでOK

併せて、TensorBoard ログの出力先とか、作成したLoRA ファイルとかが同じような名前ばかりでごちゃごちゃしてきてわかりにくくなったためにパラメータ指定するようにしたのですが…

export lora_lora_type=LoRA
export lora_learning_rate=1e-4
export lora_name=zundamon
export lora_optimizer_type=Lion
export lora_lr_scheduler=polynomial
export lora_polynomial_lr_scheduler_power=1
export lora_network_dim=16
export lora_network_alpha=16
export lora_version=mask_loss_v03
export lora_version_long=${lora_lora_type}_${lora_optimizer_type}_dim${lora_network_dim}_alpha${lora_network_alpha}_${lora_version}

accelerate launch --num_cpu_threads_per_process 2 sdxl_train_network.py \
--pretrained_model_name_or_path "/mnt/c/workspaces_c/stablediffusion_data/models/checkpoints/sd_xl_base_1.0_0.9vae.safetensors" \
--output_dir "/mnt/c/workspaces_c/kohya_lora_gui/output/${lora_name}/${lora_version}" --network_module "networks.lora" \
--dataset_config "/mnt/c/workspaces_c/kohya_lora_gui/${lora_name}/${lora_name}.toml"  \
--masked_loss \
--xformers --gradient_checkpointing --persistent_data_loader_workers --cache_latents --cache_latents_to_disk --max_data_loader_n_workers 2 --enable_bucket --save_model_as "safetensors" --mixed_precision "bf16" --resolution 1024 --train_batch_size 10 --max_train_epochs 22 --network_dim ${lora_network_dim} --network_alpha ${lora_network_alpha} --shuffle_caption --save_every_n_epochs 1 --save_precision "fp16" --min_bucket_reso 512 --max_bucket_reso 2048 --caption_extension ".txt" --seed 42 \
--optimizer_type "${lora_optimizer_type}" \
--learning_rate "${lora_learning_rate}" \
--lr_scheduler "${lora_lr_scheduler}" --lr_warmup_steps 1 --lr_scheduler_power "${lora_polynomial_lr_scheduler_power}" \
--output_name "SDXL_${lora_name}_${lora_version_long}_lr${lora_learning_rate}_${lora_lr_scheduler}_power${lora_polynomial_lr_scheduler_power}.safetensors" \
--logging_dir "/mnt/c/workspaces_c/kohya_lora_gui/tensorboard_log/20_${lora_name}/${lora_version_long}_lr${lora_learning_rate}_${lora_optimizer_type}_${lora_lr_scheduler}_power${lora_polynomial_lr_scheduler_power}"

ながっ！

一応、意図としては、指定した学習パラメータや LoRAの名称がファイル名やフォルダ名に入るようにしています。

元画像だっただけの学習とはかなり学習の傾向が変わるようで、前回利用したパラメータが使えませんでした。何度も発散して失敗しながらパラメータを選びました。

なんとか完走したものの学習量は微妙なため、その後も何度もlr と polynomial power値を変更しながら追加学習させ、5回の追加学習（計6回学習）させました。総ステップ数は2,000くらい（バッチ1換算だと x12で 24,000ステップ相当！）でしょうか。

学習結果の評価

学習させた結果、だいぶずんだもんが出るようになったので、前回作成したものと比較しました。
プロンプトは

zundamon, in the beach, masterpiece

検証プロンプト

を元として、 beach のところを city、 forest、 desert　に置き換えた4種で出力を試して背景が変わるかを試しました。

1段目がLoRAなし（sd_xl_base_1.0_0.9vae.safetensorsのみ)
2段目が前回作成した V1(zundamon_SDXL_LoRA)
3段目が今回作成した V2(zundamon_SDXK_LoRA_v02)

1段目 LoRAを適用していないので、ずんだもんは出力されません（装甲車っぽいのかっけぇ！）
2段目、ずんだもんしか出ないくらい学習させたので過学習気味だったにもかかわらず、色合いもあまく、プロンプトを変更していっても背景が変更されません。砂漠は出ている感じ？
3段目が今回のV2。背景が指定通りのものが出ています。ずんだもんも（多少形状が怪しいのもありますが）緑髪で出力されています。2段目よりいいのは明らかですね！

終わりに

今回の学習結果のLoRAも前回と同じ以下に配置しました。
zundamon_SDXL_LoRA_v02.safetensors
https://huggingface.co/suito-venus/zundamon_sdxl_lora/resolve/main/zundamon_SDXL_LoRA_v02.safetensors?download=true

kohya-ss / sd-scripts

ツールを作ってくださる方に感謝！

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