LoRA

これは実験的な機能であり、将来APIが変更される可能性があります。

LoRA (Low-Rank Adaptation of Large Language Models)は、学習可能なパラメータ数を大幅に削減する、人気のある軽量な学習手法です。モデル内に少量の新しい重みを挿入し、それらの重みのみを学習することで機能します。これにより、LoRAを用いた学習ははるかに高速かつメモリ効率が高くなり、また、保存や共有が容易な小さなモデル重み（数百MB程度）が生成されます。LoRAは、DreamBoothなどの他の学習手法と組み合わせることで、学習の高速化も可能です。

LoRAは非常に汎用性が高く、DreamBooth、Kandinsky 2.2、Stable Diffusion XL、およびtext-to-imageに対応しています。

This guide will explore the train_text_to_image_lora.py script to help you become more familiar with it, and how you can adapt it for your own use-case.

スクリプトを実行する前に、ソースからライブラリをインストールしていることを確認してください。

git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .

トレーニングスクリプトを含むサンプルフォルダに移動し、使用しているスクリプトに必要な依存関係をインストールします:

cd examples/text_to_image
pip install -r requirements.txt

🤗 Accelerate は、複数 GPU/TPU での学習や混合精度学習を支援するためのライブラリです。ハードウェアや環境に基づいて、学習環境を自動的に設定します。詳細については、🤗 Accelerate の クイックツアー をご覧ください。

🤗 Accelerate の環境を初期化します:

accelerate config

設定を選択せずにデフォルトの🤗 Accelerate環境を設定するには:

accelerate config default

または、ノートブックのようにインタラクティブシェルをサポートしていない環境では、以下を使用できます。

from accelerate.utils import write_basic_config

write_basic_config()

最後に、独自のデータセットでモデルをトレーニングしたい場合は、トレーニング用のデータセットを作成するガイドを参照して、トレーニングスクリプトと互換性のあるデータセットの作成方法を確認してください。

以下のセクションでは、スクリプトの修正方法を理解するために重要な部分に焦点を当てていますが、スクリプトのすべての側面を詳細に説明するものではありません。さらに詳しく知りたい場合は、スクリプトをぜひお読みいただき、ご質問やご懸念事項があればお知らせください。

スクリプトのパラメータ

トレーニングスクリプトには、トレーニング実行をカスタマイズするための多くのパラメータがあります。すべてのパラメータとその説明は、parse_args()関数に記載されています。ほとんどのパラメータには、かなり良好に動作するデフォルト値が設定されていますが、必要に応じてトレーニングコマンドで独自の値を設定することもできます。

例えば、学習のエポック数を増加させるには:

accelerate launch train_text_to_image_lora.py \
  --num_train_epochs=150 \

基本的かつ重要なパラメータの多くは、Text-to-imageのトレーニングガイドで説明されているため、本ガイドではLoRAに関連するパラメータに焦点を当てます。

• --rank: 学習対象の低ランク行列の内部次元。ランクが高いほど、学習可能なパラメータ数が増えます
• --learning_rate: デフォルトの学習率は 1e-4 ですが、LoRA を使用する場合、より高い学習率を設定できます

学習スクリプト

データセットの前処理コードと学習ループはmain()関数にあり、学習スクリプトを調整する必要がある場合は、ここで変更を加えます。

スクリプトのパラメータと同様に、Text-to-imageトレーニングガイドでは、トレーニングスクリプトの解説が提供されています。本ガイドでは、スクリプトのうちLoRAに関連する部分に焦点を当てて解説します。

unet_lora_config = LoraConfig(
    r=args.rank,
    lora_alpha=args.rank,
    init_lora_weights="gaussian",
    target_modules=["to_k", "to_q", "to_v", "to_out.0"],
)

unet.add_adapter(unet_lora_config)
lora_layers = filter(lambda p: p.requires_grad, unet.parameters())

text_lora_config = LoraConfig(
    r=args.rank,
    lora_alpha=args.rank,
    init_lora_weights="gaussian",
    target_modules=["q_proj", "k_proj", "v_proj", "out_proj"],
)

text_encoder_one.add_adapter(text_lora_config)
text_encoder_two.add_adapter(text_lora_config)
text_lora_parameters_one = list(filter(lambda p: p.requires_grad, text_encoder_one.parameters()))
text_lora_parameters_two = list(filter(lambda p: p.requires_grad, text_encoder_two.parameters()))

lora_layers が最適化対象の重みのみであるため、optimizer は lora_layers を用いて初期化されます:

optimizer = optimizer_cls(
    lora_layers,
    lr=args.learning_rate,
    betas=(args.adam_beta1, args.adam_beta2),
    weight_decay=args.adam_weight_decay,
    eps=args.adam_epsilon,
)

Aside from setting up the LoRA layers, the training script is more or less the same as train_text_to_image.py!

スクリプトを実行します

すべての変更が完了したか、デフォルトの設定で問題ない場合は、トレーニングスクリプトの起動準備が整いました! 🚀

Naruto BLIP captions データセットで学習し、独自のナルトキャラクターを生成しましょう。環境変数 MODEL_NAME と DATASET_NAME に、それぞれモデルとデータセットを指定してください。また、OUTPUT_DIR にモデルの保存先を、HUB_MODEL_ID に Hub へ保存するモデル名を指定する必要があります。このスクリプトは、リポジトリに以下のファイルを作成して保存します。

• 保存されたモデルのチェックポイント
• pytorch_lora_weights.safetensors（学習済みの LoRA 重み）

複数の GPU で学習を行う場合は、accelerate launch コマンドに --multi_gpu パラメータを追加してください。

2080 Ti GPU（VRAM 11GB）での完全なトレーニング実行には、約5時間かかります。

export MODEL_NAME="stable-diffusion-v1-5/stable-diffusion-v1-5"
export OUTPUT_DIR="/sddata/finetune/lora/naruto"
export HUB_MODEL_ID="naruto-lora"
export DATASET_NAME="lambdalabs/naruto-blip-captions"

accelerate launch --mixed_precision="fp16"  train_text_to_image_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$DATASET_NAME \
  --dataloader_num_workers=8 \
  --resolution=512 \
  --center_crop \
  --random_flip \
  --train_batch_size=1 \
  --gradient_accumulation_steps=4 \
  --max_train_steps=15000 \
  --learning_rate=1e-04 \
  --max_grad_norm=1 \
  --lr_scheduler="cosine" \
  --lr_warmup_steps=0 \
  --output_dir=${OUTPUT_DIR} \
  --push_to_hub \
  --hub_model_id=${HUB_MODEL_ID} \
  --report_to=wandb \
  --checkpointing_steps=500 \
  --validation_prompt="A naruto with blue eyes." \
  --seed=1337

トレーニングが完了したら、モデルを推論に使用できます。

from diffusers import AutoPipelineForText2Image
import torch

pipeline = AutoPipelineForText2Image.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5", dtype=torch.float16).to("cuda")  # or "mps", "xpu", "cpu"
pipeline.load_lora_weights("path/to/lora/model", weight_name="pytorch_lora_weights.safetensors")
image = pipeline("A naruto with blue eyes").images[0]

次のステップ

LoRAで新しいモデルの学習を完了したことをお祝いします。新しいモデルの使用方法について詳しく知るために、以下のガイドが役立つ可能性があります。

• KohyaやTheLastBenなどのコミュニティ製トレーナーで学習したさまざまなLoRA形式を読み込む方法を学びましょう。
• PEFT を用いて推論時に複数の LoRA を使用し、それらを組み合わせる方法を学びましょう。