HuggingFace Diffusers 0.12 : 訓練 : DreamBooth 再調整 (翻訳/解説)

翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 02/18/2023 (v0.12.1)

* 本ページは、HuggingFace Diffusers の以下のドキュメントを翻訳した上で適宜、補足説明したものです：

• Training : DreamBooth fine-tuning example

* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。

 

 

HuggingFace Diffusers 0.12 : 訓練 : DreamBooth 再調整

DreamBooth は主題の少ない (3~5) 画像だけが与えられたときに stable diffusion のようなテキスト-to-画像変換モデルをパーソナライズ (個人向けにカスタマイズ) する手法です。

 
Dreambooth 訓練スクリプト は、この訓練手続きを事前訓練済み Stable Diffusion モデルで実装する方法を示します。

Dreambooth fine-tuning is very sensitive to hyperparameters and easy to overfit. We recommend you take a look at our in-depth analysis with recommended settings for different subjects, and go from there.

 

ローカルで訓練する

依存関係のインストール

スクリプトを実行する前に、ライブラリの訓練依存関係を確実にインストールしてください。diffusers を main github ブランチからインストールすることも勧めます。

pip install git+https://github.com/huggingface/diffusers
pip install -U -r diffusers/examples/dreambooth/requirements.txt

xFormers は訓練要件の一部ではありませんが、できればそれをインストールすることを勧めます。それは訓練を高速にしてメモリ負荷を少なくします。

すべての依存関係がセットアップされた後、以下で 🤗 Accelerate 環境を configure できます :

accelerate config

この例ではモデルバージョン v1-4 を使用していますので、進む前に そのカード にアクセスしてライセンスを注意深く読んでください。

下のコマンドはモデル重みをハブからダウンロードしてキャッシュします、モデルの Hub id CompVis/stable-diffusion-v1-4 を使用するからです。レポジトリをローカルで複製してチェックアウトがセーブされたローカルパスを使用することもできます。

 

ドッグ toy サンプル

このサンプルでは これらの画像 を使用して、Dreambooth プロセスを使用して新しいコンセプトを Stable Diffusion に追加します。これらは訓練データになります。それらをダウンロードして貴方のシステムのどこかに配置してください。

それから以下を使用して訓練スクリプトを起動できます :

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export INSTANCE_DIR="path_to_training_images"
export OUTPUT_DIR="path_to_saved_model"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=400

 

prior-preserving 損失による訓練

過剰適合と言語ドリフトを回避するために prior preservation (事前保存) が使用されます。関心があれば、それの詳細を学習するために論文を参照してください。事前保存のために、同じクラスの別の画像を訓練プロセスの一部として使用します。素敵なことに、Stable Diffusion 自体を使用してそれらの画像を生成することができます！訓練スクリプトは生成された画像を指定したローカルパスにセーブします。

論文によれば、事前保存のために num_epochs * num_samples 画像を生成することが推奨されています。200-300 が殆どの場合に対して上手く動作します。

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export INSTANCE_DIR="path_to_training_images"
export CLASS_DIR="path_to_class_images"
export OUTPUT_DIR="path_to_saved_model"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --class_data_dir=$CLASS_DIR \
  --output_dir=$OUTPUT_DIR \
  --with_prior_preservation --prior_loss_weight=1.0 \
  --instance_prompt="a photo of sks dog" \
  --class_prompt="a photo of dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --num_class_images=200 \
  --max_train_steps=800

 

訓練中にチェック・ポイントをセーブする

Dreambooth による訓練中は簡単に過剰適合しますので、そのプロセスの間規則的にチェック・ポイントをセーブすることは有用な場合があります。中間的なチェック・ポイントの一つが最終的なモデルよりも良い場合があります！この機能を使用するには訓練スクリプトに以下の引数を渡す必要があります :

  --checkpointing_steps=500

これは output_dir のサブフォルダに完全な訓練状態をセーブします。サブフォルダ名はプレフィックス checkpoint- で始まり、次にそこまでに実行されたステップ数となります ; 例えば : checkpoint-1500 は 1500 訓練ステップ後にセーブされたチェック・ポイントです。

 

セーブされたチェックポイントから訓練を再開する

セーブされたチェック・ポイントのいずれかから訓練を再開したい場合、引数 –resume_from_checkpoint を渡してから使用したいチェック・ポイントの名前を指定できます。セーブされた最新のチェック・ポイント (i.e, 最大のステップ数を持つもの) から再開するために特殊文字列 “latest” を使用することもできます。例えば、以下は 1500 ステップ後にセーブされたチェック・ポイントから訓練を再開します :

  --resume_from_checkpoint="checkpoint-1500"

This would be a good opportunity to tweak some of your hyperparameters if you wish.

 

セーブされたチェック・ポイントを使用して推論を実行する

セーブされたチェック・ポイントは訓練を再開するために適した形式でストアされています。それらはモデル重みを含むだけでなく、optimizer の状態, データローダと学習率も含みます。

推論のためにチェック・ポイントを使用できますが、最初にそれを推論パイプラインに変換する必要があります。This is how you could do it:

from accelerate import Accelerator
from diffusers import DiffusionPipeline

# Load the pipeline with the same arguments (model, revision) that were used for training
model_id = "CompVis/stable-diffusion-v1-4"
pipeline = DiffusionPipeline.from_pretrained(model_id)

accelerator = Accelerator()

# Use text_encoder if `--train_text_encoder` was used for the initial training
unet, text_encoder = accelerator.prepare(pipeline.unet, pipeline.text_encoder)

# Restore state from a checkpoint path. You have to use the absolute path here.
accelerator.load_state("/sddata/dreambooth/daruma-v2-1/checkpoint-100")

# Rebuild the pipeline with the unwrapped models (assignment to .unet and .text_encoder should work too)
pipeline = DiffusionPipeline.from_pretrained(
    model_id,
    unet=accelerator.unwrap_model(unet),
    text_encoder=accelerator.unwrap_model(text_encoder),
)

# Perform inference, or save, or push to the hub
pipeline.save_pretrained("dreambooth-pipeline")

 

16GB GPU での訓練

勾配チェックポインティングと bitsandbytes の 8-bit optimizer の支援により、16GB GPU で dreambooth を訓練することができます。

pip install bitsandbytes

Then pass the –use_8bit_adam option to the training script.

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export INSTANCE_DIR="path_to_training_images"
export CLASS_DIR="path_to_class_images"
export OUTPUT_DIR="path_to_saved_model"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --class_data_dir=$CLASS_DIR \
  --output_dir=$OUTPUT_DIR \
  --with_prior_preservation --prior_loss_weight=1.0 \
  --instance_prompt="a photo of sks dog" \
  --class_prompt="a photo of dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=2 --gradient_checkpointing \
  --use_8bit_adam \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --num_class_images=200 \
  --max_train_steps=800

 

UNet に加えてテキストエンコーダを再調整する

スクリプトはまた unet と共に text_encoder を再調整することも可能にします。これは特に顔について、遥かに良い結果を与えることが実験で観察されています。詳細は ブログ を参照してください。

このオプションを有効にするには、訓練スクリプトに –train_text_encoder 引数を渡します。

Training the text encoder requires additional memory, so training won’t fit on a 16GB GPU. You’ll need at least 24GB VRAM to use this option.

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export INSTANCE_DIR="path_to_training_images"
export CLASS_DIR="path_to_class_images"
export OUTPUT_DIR="path_to_saved_model"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --train_text_encoder \
  --instance_data_dir=$INSTANCE_DIR \
  --class_data_dir=$CLASS_DIR \
  --output_dir=$OUTPUT_DIR \
  --with_prior_preservation --prior_loss_weight=1.0 \
  --instance_prompt="a photo of sks dog" \
  --class_prompt="a photo of dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --use_8bit_adam
  --gradient_checkpointing \
  --learning_rate=2e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --num_class_images=200 \
  --max_train_steps=800

 

8 GB GPU で訓練する

DeepSpeed を使用すれば一部のテンソルを VRAM から CPU または NVME にオフロードすることさえ可能で、少ない GPU メモリで訓練を進めることを可能にします。

DeepSpeed は accelerate config で有効にする必要があります。configuration の間に、“Do you want to use DeepSpeed?“ に対して yes と答えます。DeepSpeed ステージ 2, fp16 混合精度, そしてモデルパラメータと optimizer 状態の両方を CPU にオフロードすることを組み合わせると、8 GB VRAM 以下で訓練することが可能です。欠点はこれはより多くのシステム RAM を必要とすることです (およそ 25 GB)。See the DeepSpeed documentation for more configuration options.

デフォルトの Adam optimizer を DeepSpeed の Adam の特別なバージョン deepspeed.ops.adam.DeepSpeedCPUAdam への変更はかなりのスピードアップを与えますが、それを有効にするためにはシステムの CUDA toolchain バージョンが PyTorch によってインストールされたものと同じである必要があります。8-bit optimizer は現時点で DeepSpeed と互換性がないようです。

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export INSTANCE_DIR="path_to_training_images"
export CLASS_DIR="path_to_class_images"
export OUTPUT_DIR="path_to_saved_model"

accelerate launch train_dreambooth.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --instance_data_dir=$INSTANCE_DIR \
  --class_data_dir=$CLASS_DIR \
  --output_dir=$OUTPUT_DIR \
  --with_prior_preservation --prior_loss_weight=1.0 \
  --instance_prompt="a photo of sks dog" \
  --class_prompt="a photo of dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --sample_batch_size=1 \
  --gradient_accumulation_steps=1 --gradient_checkpointing \
  --learning_rate=5e-6 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --num_class_images=200 \
  --max_train_steps=800 \
  --mixed_precision=fp16

 

推論

モデルを訓練したら、StableDiffusionPipeline を使用して、モデルがセーブされたパスを単純に指定することにより推論が実行できます。プロンプトが訓練中に使用された特殊な識別子を含むことを確実にしてください (前の例では sks)。

from diffusers import StableDiffusionPipeline
import torch

model_id = "path_to_saved_model"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to("cuda")

prompt = "A photo of sks dog in a bucket"
image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]

image.save("dog-bucket.png")

You may also run inference from any of the saved training checkpoints.

以上