Google Colab ではじめる LoRA

2022年12月18日 16:38

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

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