DreamBoothスクリプトのfine-tuning機能について
※fine tuning専用のスクリプトを作成しましたので、通常はそちらをご利用いただければ幸いです。
はじめに
前回の記事の続きです。スクリプトのfine-tuning機能の解説です。
概要
DreamBoothはモデルのU-NetとText Encoderを学習しています。ですので、画像とプロンプト(キャプション)の取り扱いを変えるだけでそれらのfine-tuningが可能です。
教師データの準備
画像ファイルと、その画像のプロンプトが1行で記述されたテキストファイルをペアで用意します。「画像ファイル名の拡張子を.txtまたは.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の経験は私もほぼありませんが、学習率はかなり低くしたほうが良いようです。特に教師データの枚数が少ない場合は容易に過学習を起こすと思われます。
教師データの枚数が少なくエポックが短くなりすぎる場合は、お手数ですが画像ファイルとテキストファイルをコピーして水増ししてください。

雑なデータでもこのくらいの学習はできます。学習率1e-7、バッチサイズ10、4000step学習
おわりに
fine tuningは以下のページが参考になります。22,891枚で13epoch回しているので、バッチサイズ1なら30万step程度学習すればこのくらいの結果は得られるようです。
なお、DreamBoothを行うときに正則化画像(reg_data_dirオプション)を指定せず、フォルダ名をキャラクタ名などにすれば単なるモデルのfine-tuningになり、キャラクタを学ばせるならそれでも充分な場合もあります。学習率を1e-7のようにかなり低くし、ステップ数を長めにするのが良いようです。
DreamBoothとfine-tuningをうまく使い分けていただければ幸いです。