DreamBoothスクリプトのfine-tuning機能について

※fine tuning専用のスクリプトを作成しましたので、通常はそちらをご利用いただければ幸いです。

はじめに

前回の記事の続きです。スクリプトのfine-tuning機能の解説です。

概要

DreamBoothはモデルのU-NetとText Encoderを学習しています。ですので、画像とプロンプト（キャプション）の取り扱いを変えるだけでそれらのfine-tuningが可能です。

教師データの準備

画像ファイルと、その画像のプロンプトが1行で記述されたテキストファイルをペアで用意します。「画像ファイル名の拡張子を.txtまたは.captionに変えたもの」がテキストファイルのファイル名になります。

ざっくりこんな感じです

caption

これはあくまでサンプルのため15枚ですが、目的にもよりますが、教師データの画像はもっとたくさん用意したほうが良いでしょう。

学習の実行

fine_tuningオプションを指定してスクリプトを実行します。最大限にメモリを節約した例は以下になります（実際には1行で入力します）。

accelerate launch --num_cpu_threads_per_process 8 train_db_fixed_v6.py 
    --pretrained_model_name_or_path=<Diffusers版モデルのディレクトリ> 
    --train_data_dir=<学習用データのディレクトリ> 
    --output_dir=<学習したモデルの出力先ディレクトリ> 
    --fine_tuning 
    --resolution=512 
    --train_batch_size=1 
    --learning_rate=1e-7 
    --max_train_steps=1600 
    --use_8bit_adam 
    --xformers 
    --mixed_precision="bf16" 
    --cache_latents
    --gradient_checkpointing

オプションは、reg_data_dirとprior_loss_weightを指定しないことを除き、DreamBoothとほぼ同じです。

shuffle_captionオプションを追加すると、学習時にaugmentationのため、プロンプトをカンマで区切られた要素ごとにシャッフルします。

雑ですがこんな感じです

メモリが許す場合はmixed_precisionを削除するか"no"を指定してfloat32で学習すると良いでしょう（これだけならVRAM 12GBに収まると思います）。さらに余裕がある場合は、gradient_checkpointingを外すと速くなります。またバッチサイズも大きくすると速度および精度が上がります。

DreamBoothより長い学習が必要になりますので、不慮の事態に備えてsave_every_n_epochsオプションをお使いください。

dataset_repeatsオプションを指定するとデータを指定数だけ繰り返したものを1 epochとします。教師データの枚数が少なく1 epochが極端に短くなってしまう場合にお使いください。

※モデル変換して保存している都合上、optimizerの状態を保存しないため、途中で中断した場合には厳密な再開はできません（中断せずに学習した場合と結果が変わります）。

学習時の注意

モデルのfine tuningの経験は私もほぼありませんが、学習率はかなり低くしたほうが良いようです。特に教師データの枚数が少ない場合は容易に過学習を起こすと思われます。

教師データの枚数が少なくエポックが短くなりすぎる場合は、お手数ですが画像ファイルとテキストファイルをコピーして水増ししてください。

green frog, riding bicycle, gray background
雑なデータでもこのくらいの学習はできます。学習率1e-7、バッチサイズ10、4000step学習

おわりに

fine tuningは以下のページが参考になります。22,891枚で13epoch回しているので、バッチサイズ1なら30万step程度学習すればこのくらいの結果は得られるようです。

Stable Diffusionを「いらすとや」で追加学習する - TadaoYamaokaの開発日記

なお、DreamBoothを行うときに正則化画像（reg_data_dirオプション）を指定せず、フォルダ名をキャラクタ名などにすれば単なるモデルのfine-tuningになり、キャラクタを学ばせるならそれでも充分な場合もあります。学習率を1e-7のようにかなり低くし、ステップ数を長めにするのが良いようです。

DreamBoothとfine-tuningをうまく使い分けていただければ幸いです。