メインコンテンツへスキップ
見出し画像

Google Colab ではじめる LoRA

    npaka

    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ドライブのマウント。

    # Googleドライブのマウント
from google.colab import drive
drive.mount('/content/drive')

    (3) 作業フォルダへの移動。

    # 作業フォルダへの移動
import os
os.makedirs("/content/drive/My Drive/work", exist_ok=True)
%cd '/content/drive/My Drive/work'

    (4) パッケージのインストール。

    # パッケージのインストール
!git clone https://github.com/cloneofsimo/lora.git
!pip install accelerate bitsandbytes
%cd lora
!pip install .

# 入出力フォルダの準備
os.makedirs("./instance_data", exist_ok=True)
os.makedirs("./output", exist_ok=True)

    (5) 学習画像を「instance_data」フォルダに配置。
    以下の5枚の画像をGoogleドライブ経由で配置しました。

    画像

    おすすめの「学習画像」の条件は、次のとおりです。

    ・学習画像の枚数は 10〜120枚 程度
    ・単体オブジェクトが写ってるもの
    ・多様なポーズ
    ・多様な背景

    必要な画像枚数は、モデルにどの程度の柔軟さを求めるかによります。似たポーズ、似た背景が多いと、それも学習してしまいます。

    (6) ファインチューニングの実行。
    インスタンスプロンプト「sks」、学習画像フォルダ「instance_data」で学習します。

    # ファインチューニングの実行
!accelerate launch   \
  --num_processes=1  \
  --num_machines=1  \
  --mixed_precision="fp16"  \
  --dynamo_backend="no"  \
  train_lora_dreambooth.py \
  --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1-base"  \
  --instance_prompt="sks" \
  --instance_data_dir="./instance_data" \
  --output_dir="./output" \
  --resolution=512 \
  --train_batch_size=1 \
  --color_jitter \
  --learning_rate=1e-4 \
  --learning_rate_text=5e-5 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --gradient_accumulation_steps=1 \
  --max_train_steps=3000 \
  --train_text_encoder

    これで、「sks」というプロンプトに「特定の猫」が学習されます。学習結果は、「output」フォルダに出力されます。

    ・lora_weight.pt : Untの重み
    ・lora_weight.text_encoder.pt : テキストエンコーダーの重み

    3. 推論の実行

    Colabでの推論の実行手順は、次のとおりです。

    (1) パイプラインの準備。

    import torch
from lora_diffusion import monkeypatch_lora, tune_lora_scale
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler


# パイプラインの準備
model_id = "stabilityai/stable-diffusion-2-1-base"
pipe = StableDiffusionPipeline.from_pretrained(
    model_id, 
    scheduler=EulerDiscreteScheduler.from_pretrained(
        model_id, 
        subfolder="scheduler"
    ),     
    torch_dtype=torch.float16
).to("cuda")

# LoRAの重みの適用
monkeypatch_lora(pipe.unet, 
    torch.load(os.path.join("./output", "lora_weight.pt")))
monkeypatch_lora(pipe.text_encoder, 
    torch.load(os.path.join("./output", "lora_weight.text_encoder.pt")), 
    target_replace_module=["CLIPAttention"])

    モデルにLoRAの重みを適用するには、monkeypatch_lora()を使います。

    (2) 推論の実行。
    モデルのLoRAの重みを調整するには、tune_lora_scale()を使います。
    初期状態では月がでなかったので、0.8で学習結果を弱めています。

    # LoRAの重みの調整
tune_lora_scale(pipe.unet, 0.8)
tune_lora_scale(pipe.text_encoder, 0.8)

# 推論の実行
image = pipe(
    "sks on the moon", 
    num_inference_steps=50, 
    guidance_scale=7
).images[0]
image
    画像

    4. 正則化画像を利用したファインチューニングの実行

    正則化画像を利用したファインチューニングの実行手順は、次のとおりです。

    (1) 正則化画像を用意し、class_dataフォルダに配置。
    今回は、500枚用意しました。

    画像

    「sks cat」というインスタンスプロンプトで「sks」に「特定の猫の柄」を学習させます。正則化画像がないと、「sks」と「cat」の両方の単語の概念が変更されるため、「sks」にどのような内容が学習されるかわかりません。そこで、「cat」という単語を正則化画像で学習させることで、結果として「sks」 に「cat」以外の概念を学習させることができます。

    おすすめの「正則化画像」の条件は、次のとおりです。

    ・正則化画像の枚数は 学習画像枚数x100 程度

    (2) ファインチューニングの実行。
    クラスプロンプト「cat」、正則化画像フォルダ「class_data」で学習します。

    !accelerate launch   \
  --num_processes=1  \
  --num_machines=1  \
  --mixed_precision="fp16"  \
  --dynamo_backend="no"  \
  train_lora_dreambooth.py \
  --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-1-base"  \
  --instance_prompt="sks cat" \
  --instance_data_dir="./instance_data" \
  --class_prompt="cat" \
  --class_data_dir="./class_data" \
  --num_class_images=500 \
  --output_dir="./output" \
  --resolution=512 \
  --train_batch_size=1 \
  --with_prior_preservation \
  --learning_rate=1e-4 \
  --learning_rate_text=5e-5 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --train_text_encoder \
  --gradient_accumulation_steps=1 \
  --max_train_steps=3000 \
  --train_text_encoder

    これで、「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> : 分散学習のローカルランク

    【おまけ】 正則化画像の生成スクリプト

    正則化画像の生成スクリプトの例は、次のとおりです。

    # Googleドライブのマウント
from google.colab import drive
drive.mount('/content/drive')

# 作業フォルダへの移動
import os
os.makedirs("/content/drive/My Drive/work", exist_ok=True)
%cd '/content/drive/My Drive/work'

# パッケージのインストール
!pip install diffusers[torch] transformers safetensors accelerate
    from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
import torch

# パイプラインの準備
model_id = "stabilityai/stable-diffusion-2-1-base"
pipe = StableDiffusionPipeline.from_pretrained(
    model_id, 
    scheduler=EulerDiscreteScheduler.from_pretrained(
        model_id, 
        subfolder="scheduler"
    ), 
    torch_dtype=torch.float16
).to("cuda")
    # パラメータ
class_prompt = "cat, photo, full body shot"
num_class_images = 500

# 出力フォルダの準備
os.makedirs("./output", exist_ok=True)

# 推論の実行
num = 0
while num < num_class_images:
    images = pipe(class_prompt, 
        width=512,
        height=512,
        num_inference_steps=20,
        guidance_scale=7.5,
        num_images_per_prompt=10).images
    for image in images:
        image.save("./output/image" + str(num) + ".png")
        num += 1



     
     

    npaka

     
     
    プログラマー。iPhone / Android / Unity / ROS / AI / AR / VR / RasPi / ロボット / ガジェット。年2冊ペースで技術書を執筆。アニソン / カラオケ / ギター / 猫 twitter : @npaka123

    あなたへのおすすめ