Google Colab ではじめる LoRA｜npaka #note
Google Colab で LoRA を試したのでまとめました。
1. LoRA
「LoRA」(Low-rank Adaptation)は、数枚の被写体画像と対応するテキストを元にファインチューニングを行うことで、Text-to-Imageモデルに新たな被写体を学習させる手法です。
特徴は、次のとおりです。
・Dreamboothより高速・VRAM 8GBでも動作・学習データだけ抽出して他モデルとマージできる・学習結果のサイズが小さい (Unet のみで3MB、Unet+Clipで6MB)・UnetとCLIPの両方をファインチューニング可能。
2. ファインチューニングの実行
Colabでのファインチューニングの実行手順は、次のとおりです。
(1) メニュー「編集→ノートブックの設定」で、「ハードウェアアクセラレータ」に「GPU」を選択。
(2) Googleドライブのマウント。
(3) 作業フォルダへの移動。
(4) パッケージのインストール。
(5) 学習画像を「instance_data」フォルダに配置。
以下の5枚の画像をGoogleドライブ経由で配置しました。
おすすめの「学習画像」の条件は、次のとおりです。
・学習画像の枚数は 10〜120枚 程度・単体オブジェクトが写ってるもの・多様なポーズ・多様な背景
必要な画像枚数は、モデルにどの程度の柔軟さを求めるかによります。似たポーズ、似た背景が多いと、それも学習してしまいます。
(6) ファインチューニングの実行。
インスタンスプロンプト「sks」、学習画像フォルダ「instance_data」で学習します。
これで、「sks」というプロンプトに「特定の猫」が学習されます。学習結果は、「output」フォルダに出力されます。
・lora_weight.pt : Untの重み・lora_weight.text_encoder.pt : テキストエンコーダーの重み
3. 推論の実行
Colabでの推論の実行手順は、次のとおりです。
(1) パイプラインの準備。
モデルにLoRAの重みを適用するには、monkeypatch_lora()を使います。
(2) 推論の実行。
モデルのLoRAの重みを調整するには、tune_lora_scale()を使います。
初期状態では月がでなかったので、0.8で学習結果を弱めています。
4. 正則化画像を利用したファインチューニングの実行
正則化画像を利用したファインチューニングの実行手順は、次のとおりです。
(1) 正則化画像を用意し、class_dataフォルダに配置。
今回は、500枚用意しました。
「sks cat」というインスタンスプロンプトで「sks」に「特定の猫の柄」を学習させます。正則化画像がないと、「sks」と「cat」の両方の単語の概念が変更されるため、「sks」にどのような内容が学習されるかわかりません。そこで、「cat」という単語を正則化画像で学習させることで、結果として「sks」 に「cat」以外の概念を学習させることができます。
おすすめの「正則化画像」の条件は、次のとおりです。
・正則化画像の枚数は 学習画像枚数x100 程度
(2) ファインチューニングの実行。
クラスプロンプト「cat」、正則化画像フォルダ「class_data」で学習します。
これで、「sks cat」というプロンプトに「特定の猫」、「sks」に「特定の猫の柄」が学習されます。
(3) 推論の実行。
「sks cat」で「特定の猫」、「sks dog」で「同じ柄の犬」、「sks rabbit」で「同じ柄のうさぎ」が生成されることを確認します。
見たところ「sks」に部屋の背景も学習されてしまってるようなので、多様な背景の学習画像を用意した方が良さそうです。
【おまけ】 train_lora_dreambooth.py のパラメータ
「train_lora_dreambooth.py」のパラメータは、次のとおりです。
◎ 基本-h, --help : ヘルプ--pretrained_model_name_or_path <学習済みモデル名>--pretrained_vae_name_or_path <学習済みVAE名>--revision <学習済みモデルのリビジョン>--tokenizer_name <トークナイザー名>--resolution <入力画像の解像度>--train_batch_size <学習バッチサイズ>--max_train_steps <学習ステップ数>◎ 入出力--instance_prompt <インスタンスプロンプト>--instance_data_dir <学習画像フォルダのパス> --class_prompt <クラスプロンプト>--class_data_dir <正則化画像フォルダのパス>--num_class_images <正則化画像数>--output_dir <出力フォルダのパス>◎ 学習率--learning_rate <学習率>--learning_rate_text <テキストエンコーダーの学習率>--lr_scheduler {"linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"}--lr_warmup_steps <ウォームアップステップ数>--scale_lr : 学習率をスケーリングするか◎ テキストエンコーダーの学習--train_text_encoder : テキストエンコーダーを学習するか◎ 学習再開--resume_unet <Unetの重みのパス>--resume_text_encoder <テキストエンコーダーの重みのパス>◎ その他--with_prior_preservation : prior preservation lossの追加--prior_loss_weight <PRIOR_LOSS_WEIGHT> : prior preservation lossの重み--seed <乱数シード>--center_crop : センタークロップするか--color_jitter : カラージッターを適用するか--sample_batch_size <サンプルバッチサイズ>--num_train_epochs <学習エポック数>--save_steps <保存ステップ数> --gradient_accumulation_steps <勾配累積ステップ数>--gradient_checkpointing : メモリ節約のため勾配チェックポイントを使用するか--lora_rank <LoRAランク>--use_8bit_adam : 8-bit Adamを使用するか--adam_beta1 <ADAM_BETA1> : Adamのパラメータ--adam_beta2 <ADAM_BETA2> : Adamのパラメータ--adam_weight_decay <ADAM_WEIGHT_DECAY> : Adamのパラメータ--adam_epsilon <ADAM_EPSILON> : Adamのパラメータ--max_grad_norm <最大勾配ノルム>--push_to_hub : モデルをHuggingface HubにPushするか--hub_token <HUB_TOKEN>--logging_dir <TensorBoardのログフォルダ>--mixed_precision {no,fp16,bf16} : 混合精度を使用するか--local_rank <LOCAL_RANK> : 分散学習のローカルランク
【おまけ】 正則化画像の生成スクリプト
正則化画像の生成スクリプトの例は、次のとおりです。
