Colab Enterpriseを使ってStable DiffusionのLoRA学習を行なってみたいと思います。

Colab Enterpriseは、ひとことで言うと「GCP上のColabratory」に相当するサービスです。ColabratoryやJupyter Notebookと同じように、ブラウザ上でPythonのコードを実行できます。課金は、使用量に応じて課金される従量課金です。有料のColab環境という点ではColabratory Proに似ていますが、個人のクレジットカード登録などが必要ないため、企業内で使用しやすいサービスになっています。

今回はこちらをAI画像生成の実験環境として使用します。Colab Enterpriseを使うメリットは、開発者だけではなく、デザイナーやイラストレーターの方々にも気軽に試してもらいやすいことです。

ただし、今回は画像の生成ではなく、気軽にGPU環境を用意できるクラウドのメリットを活かし、LoRAの学習について紹介してみたいと思います。

LoRAは少ない計算コストでモデルの追加学習を行なう手法です。画像生成の場合、絵柄やキャラクターを学習するために使用されることが多いようです。LoRA学習用途でよく使用されるKohya's GUIをColab Enterprise上で動かしてみます。

1. 環境準備

Colab Enterpriseでは、ランタイムテンプレートを設定することで、ノートブックの実行環境をきめ細かく設定することができます。今回は、次のようなランタイムテンプレートを作成しました。

インスタンスタイプ: n1-standard-4
GPU: T4
ディスクサイズ: 100GB
アイドル状態でのシャットダウン: 20分

今回実行したノートブックはこちらに公開しました。
https://gist.github.com/takada-at/545214ed2ddb68b2c526e25bb0f22e0c

基本的にはこちらを順番に実行していくだけです。コードの中身はリポジトリをチェックアウトし、依存ライブラリをインストールしているだけのシンプルなものです。サービスはgradioを通じて公開されるので、パスワード認証を設定しています。

2. 実行

ノートブックを順番に実行していくと、Kohya's GUIが起動します。以下のようにURLが表示されるのでこちらをブラウザで開きます。

表示されたURLにアクセスし、パスワードを入力すると以下のような画面が表示され、LoRAの学習を実行することができます。

3. SDXLの学習

Stable Diffusion XL1.0(SDXL1.0)はStable Diffusionの最新版のモデルです。モデル自体が巨大なため、LoRAの学習にもかなりの性能が要求され、実行のハードルが高いです。今回は、Colab Enterprise上でもSDXLでの実行が確認できたので、実行時の設定を共有したいと思います。

ランタイムテンプレート

インスタンスタイプ: n1-highmem-4
GPU: T4
ディスクサイズ: 100GB
アイドル状態でのシャットダウン: 20分

メモリ16GBではSDXLの学習に不足するようだったので、メモリ32GBのn1-highmem-4インスタンスを使用します。

学習時の設定

以下のKohya's GUIの公式リポジトリに従って、SDXL用の設定を適用していきます。わかりにくいところのみ以下に解説します。

LR Scheduler / Optimizerの設定

LR Scheduler: constant_with_warmupが推奨されているようです。
Optimizer: Adafactorに設定し、Optimizer extra argumentに"scale_parameter=False", "relative_step=False", "warmup_init=False"を設定(下記スクリーンショットを参照)。

Cache text encoder outputs: 「Parameters」「Basic」の「Cache text encoder outputs」をオンにする。
UNETのみ学習: 「Parameters」「Advanced」「Additional parameters」に"--network_train_unet_only"を設定し、UNETのみ学習するように設定します。

Gradient checkpointing: 「Parameters」「Advanced」の「Gradient checkpointing」をオンにする。

以上ひととおりの設定をすれば、T4環境でもLoRA学習が実行できました。