Diffusers版DreamboothをVRAM 8GBのWindowsで無理矢理動かす
概要
学習時のモデルの重みと勾配、テンソルをすべてfp16にすることで8GB VRAMで学習します。学習に失敗する可能性は高く、また学習できたとしても精度は大きく下がりますのであくまで実験的なものです。
PyTorchのソースコードを書き換えるなどとても無理矢理なので動いたらラッキーくらいの心持ちでお願いします。
※使用に当たっては自己責任でお願いいたします。
以下のTtl氏のバージョンでも8GB VRAMで動作するようですので、Linux/WSLでは使ったほうが良さそうです（Text Encoderは学習されないようです）。Ttl氏のバージョンはDeepSpeedを使っていますのである程度の精度は出るはずです。
この記事で使っている手法はTtl氏のものとは異なります（DeepSpeedは使っていません）。
手順
前回の記事に従い環境構築を行います。acclerate configの設定の最後の質問ではfp16を選択してください。
次にPyTorchのソースコードを書き換えて勾配にfp16を使えるようにします。torch/cuda/amp/grad_scaler.pyを書き換えます。
（ファイルの添付はありません。）
_unscale_grads_メソッドの最後の引数をFalseからTrueにするとfp16でも動くようになります。
次にfp16の学習に対応した以下のtrain_dreambooth_mod3.pyをmod2.pyと同じ場所に入れてください。mod2を元にTtl氏のソースを参考に重み等をfp16に設定したものです。
学習の実行
前回の記事と同様の設定で学習できます。--mixed_precision="fp16"を忘れずに指定してください。
生成画像が真っ黒になる場合は、精度が足りず重みが壊れています。ステップ数を短くしてみるか、学習率を変えるかなどすると学習できるかもしれません。
学習率5e-6、1000 steps学習した結果です。
