西川和久の不定期コラム
自前でStable Diffusion用美女モデル作成の後日談。そしてVRAM 80GBのA100も使ってみた！
2024年4月26日 10:12
3月25日に「自前のグラビア写真でStable Diffusion用美女モデルを作る」記事を掲載してちょうど約1カ月。最後の最後にSDXLの--learning_rate=の最適値が「4e-7」と分かったのだが、やや半端な終わり方だったので、今回は後日談としてまとめてみたい。
--learning_rate=4e-7でも……。
前回の記事では大前提として、accelerate configやsd-scriptsのパラメータにbf16を使っていた。これは理由があり、より精度の高いfp16はVRAM 24GBだと不足で動かなかったからだ。
この状態で--learning_rate=4e-7とし、ほかのパラメータは全く同じでsd-scriptsを動かすと、学習するにはするものの、何epochs回しても以下のように、顔立ちなどは若干良くなるが、sd_xl_base_1.0とほとんど変わらない絵となってしまう。
もう一点、オプティマイザーに--use_8bit_adamに指定しており、これも精度が落ちる(同じくVRAM 24GBが理由)。
この2点を何とかしないと、どれだけ学習しても、できたモデルは、学習元のsd_xl_base_1.0と大差ない絵が出てしまう。
latents cacheを使ってVRAM 24GBでfp16指定
Xで前回の記事をシェアし、この辺りの話を呟いていたところ、ありがたいことに、sd-scriptsの作者であるKohya Tech氏から、VRAM 24GBでfp16ができる方法を教えて頂いた。
無事に結果が出ているようで幸いです。full_bf16は通常のmixed precision学習に比べると精度面では若干不利ですので、latentsの事前キャッシュ（ディスクにlatentsを保存します）とAdaFactorオプティマイザを使用していただくと、24GB VRAMならfull_bf16なしで学習できるかと思います。
— Kohya Tech (@kohya_tech)March 25, 2024
追加オプションは以下のような感じになると思います。
— Kohya Tech (@kohya_tech)March 25, 2024
--cache_latents_to_disk --highvram --vae_batch_size=4 --optimizer_type=adafactor --optimizer_args "scale_parameter=False" "relative_step=False" "warmup_init=False" --lr_scheduler=constant_with_warmup --lr_warmup_steps=100
仕組みとしてはlatentsをcache、使用するVRAM容量を減らし、fp16を有効にすること。確かにこのオプションを使うと24GBでもaccelerate configをfp16、--mixed_precision='fp16' --full_fp16としてもVRAM不足にならず、学習を開始する。
以下、オプティマイザーをLionに変え、ほかのオプションも見直したsd-scriptsのパラメータとなる(config.tomlは触っていない)。
accelerate launch --num_cpu_threads_pre_process 8 sdxl_train.py
--pretrained_model_name_or_path='/sd_checkpoint/SDXL/sd_xl_base_1.0.safetensors'
--dataset_config='/checkpoint/config.toml'
--output_dir='/checkpoint/output'
--output_name='MyCkpt3-4.5k-4e7-fp16-e100'
--save_model_as=safetensors --save_precision='fp16' --save_every_n_epochs=10
--learning_rate=4e-7
--mixed_precision='fp16' --full_fp16
--gradient_checkpointing
--enable_bucket
--no_half_vae
--optimizer_type='Lion'
--lr_scheduler="constant_with_warmup" --lr_warmup_steps=“100"
--train_text_encoder
--max_token_length=225
--max_train_epochs=100
--xformers --sdpa
--cache_latents --cache_latents_to_disk
※ 実際は改行せず1行にする
output_nameは何でもいいが、どんな内容のオプションだったのか、肝の部分を含めるようにしないと、いろいろなパターンで作って、後から「これ何だっけ!?」となる。MyCkpt3-4.5k-4e7-fp16-e100は、datasetの改良版(tagやcaptionの見直し)という意味でMyCkpt→MyCkpt2→MyCkpt3、4.5k枚、fp16、100 epochsという意味だ。
datasetの枚数がもっと増えれば違うかもしれないが、4.5k枚だと20 epochs辺りから効果が出て、40 epochsでもういいかも!? 状態になる、が、倍の80 epochsにすると、やっぱりこっち! と(笑)。
目安となるロスは各epochsで眺めていると0.131から0.124程度。あまり変化はない。これは顔LoRAと違い、広く浅く学習するので、さほど下がらないのだろうと勝手に思っている。従って「何epochsでOK」かは、実際生成して確認するしか手がない。
なお、datasetの見直しとして、
- 4.5k枚から厳選した3.6k枚に絞ったもう1セットを用意
- BLIP-2で生成したcaptionを眺めていると asian woman が大量に出ているので、 japanese woman に置き換え
- tagに自分のdatasetという意味で knp (Kazuhisa Nishikawa Photoの略)を追加、そして1girlがある時は japanese を追加(japaneseがなかった場合)
と、少しだけ改良している。sd-scriptsの環境設定、画像の長辺を1,024pxにしたり、BLIP-2でcaption、JoyTagでtagを作る方法は、全く同じなので、興味のある方は前回の記事を参考にして欲しい。
さて、実際生成してみるとおもしろいことに気が付く。一般的にちょっと下からのアングルを出す場合は、from belowとPromptに入れるのだが、それなしでも気持ち下からのアングルが多く出てくる。言われてみれば「ここならこんな感じで撮るかな!?」的なものも。筆者が撮った写真を学習しているので、その辺りの好みや癖も合わせて学習しているのだろう(笑)。
次なる手法は!?
fp16で学習させた結果、望ましい絵が出るようになったのだが、何かもう少し手を加えるには!? と、ネットを検索すると、まず目についたのが「ANIMAGINE XL 2.0」。
資料を眺めると、1.0のLearning Rateは同じく4e-7となっているが、2.0は1e-6と大きな値に変わっている。以前、bf16の時もLearning Rateを大きくしたり小さくしたり、いろいろ試したところ、大きくすると例外なく体などが破綻。大丈夫なのか? と試した見たのが以下のxyzチャートとなる。さすがに1e-6はダメだったので8e-7で学習させた。
ご覧のように、sd_xl_base_1.0とは大きく雰囲気が異なりおもしろい。ただxyzチャートでは大丈夫なのだが結構体が壊れやすい。やはりリアル系では無理がありそうな感じだ。これがうまくできる方法があれば、sd_xl_base_1.0と比較してテイストを変えたモデルができるのだが……。
次にソースはちょっと忘れてしまったが、できたモデルに低いLearning Rateで再学習する方法。試しに5e-8で40 epochsだけ再学習させると以下のような感じとなった。
細部が変わるだけで、あまり変化はなくまさにFineTune。こちらはありと言えばあり。オプションといった感じだろうか。
A100(VRAM 80GB)搭載GPUサーバーを使ってみた
さて、上記にxx epochsとさらっと書いているが、4.5k枚、RTX 4090で1 epochs約1時間かかる(笑)。この1月間、実験で少しだけ学習したのも含めると、何百時間使ったのか!? というほど、もはや分からないほどGPU時間を浪費した。幸い5 epochsでも結果が出るのは5時間後。sd-scripts実行後は、普通に仕事し、気がつけばできているので、仕事にはほとんど影響は出ていない(はず)。
とは言え、ちょっと時間がかかり過ぎるのと、ファンの音がうるさく、熱で室温も気持ち上がるため、「A100欲しい!(笑)」と、Xで呟いていたところ、DMが入り、B2C向けのGPUサーバーをクローズドβテスト中だけど参加しますか? 的な内容が。テスト中は費用0! 渡に船と言うこともあり参加した(記事掲載の許諾も得ている)。
βテストは既に公開しているB2B版「Cloud Bouquet」のB2C版となる。執筆時、サイトが更新されており、使えるGPUはNVIDIA A6000、A100、そして何とH100! 各最小プランで順に297円/時(VRAM 48GB)、660円/時(VRAM 40GB)、990円/時(VRAM 80GB)。
話によると、B2B版と内容的には同じだが、B2C用は個人向けということもあり、コンパネなどの内容が変わるということらしい。今回はA100(VRAM 80GB)を使用した。このタイプだと770円/時間。以下、ログイン後のコンパネだが、βテスト中なので、最終的には変わる可能性がある。
サーバーの起動自体は簡単。サーバーを追加、サーバー名、ストレージ容量、OS(Ubuntu 20.04)を選びLaunchでOK。起動まで若干時間がかかる。ステータスがrunningになった後、事前に連絡のあった踏み台にSSH鍵1を使って接続する。
GPUサーバーへのアクセスは、踏み台にssh後、GPUサーバー用のSSH鍵2を転送、さらにそのSSH鍵2を使いGPUサーバーへ接続となる。言葉で書くと分かりにくいが、実際のコマンドはこんな感じだ。
ssh -i SSH鍵1 <user1>@<踏み台サーバー> -p <ポート番号>
これで踏み台にログイン。ログイン後、GPUサーバーのSSH鍵2を転送しておく
↓
ssh -i SSH鍵2 <user2>@<GPUサーバー> -p <ポート番号>
つまり、GPUサーバーはネットに接続されていないというわけだ。2回操作するのが面倒な場合、以下で一発接続できる。
ssh -o ProxyCommand='ssh -W %h:%p -i <SSH鍵1> -p <踏み台ポート番号> <user1>@<踏み台サーバー>' -i <SH鍵2> -p <GPUサーバーポート番号> <user2>@<GPUサーバー>
これでCLIならGPUサーバーの操作ができるのだが、ファイル転送はどうするの? という話になる。方法はローカルマシンでProxyサーバーを起動する。
ssh -L 10022:<GPUサーバー>:<GPUサーバーポート番号> <user1>@<踏み台サーバー> -p <踏み台サーバーポート番号> -i <SSH鍵1>
※ 転送など作業が終わるまでは実行したまま
冒頭の10022はプロキシのポート番号(何でも良い)。これでFileZillaなどを起動し、SFTP、localhost:10022、<user2>、<SSH鍵2>をセットすればファイル転送が可能となる。
さて接続までの話で長くなってしまったが、GPUサーバーにログインできれば、そこはまっさらなUbuntu 20.04。NVIDIAのドライバなども含め、一からの環境構築となる。チャージは、サーバーのステータスがrunningになった以降なので、このGPUを使わない環境設定時間も加算される。ちょっと痛いが仕方ないところか。sd-scripts学習用の環境なら(手慣れていれば)1時間以内に終わる。
なお、NVIDIAドライバをsudo apt installすると再起動が必要になるが、これはOKとのこと。ただし、シャットダウンはNG。落とす時にはコンパネ側を操作する。
NVIDIA A100はVRAM 40GBもしくは80GB。アーキテクチャ的には一世代前のAmpere。つまりGeForce RTX 3090などと同じ世代となる。メモリ帯域幅などが違うため、単純比較はできないものの、ざっくり同じsd-scriptsの設定だとGeForce RTX 4090の方が処理は速い。
しかし、VRAMが多いので、config.tomlのbatch_size(並列処理的なもの)を1より大きくできるため、処理速度向上が期待できる。実際batch_size=12で試したところ、1 epochs 約20分! これはGeForce RTX 4090の3倍速に相当する(後に24でも試したが大差なかった)。従って確認の目安となる10 epochsは3時間ちょっと。これならさほどストレスもたまらない。
現在GeForce RTX 4090は円安と品薄で30万円前後。この手のGPUサーバーに個人が1回分で払える金額を考えると1.5万円程度だろうか? RTX A100 80GBだとdataset 4.5k枚で19.48時間分だ。ただし環境設定が入るため、実質学習に使える時間は18時間=54 epochs回せることになる。
キリがいいところで考えると、60 epochs=20時間+1時間環境設定として21 * 770 = 1万6,170円。これなら行ける範囲ではないだろうか!? H100だと時間単価は上がるものの、その分早く処理が終われば、結果安上がりの可能性も。参考までにNVIDIA A100 80GBを購入しようとすると300万円超だ(笑)。
ただこの手のサーバーは、まだオプション指定などが確定していない試行錯誤中だと、無駄に時間だけを浪費する。確定後、あとは任せた! 的に使うのがベストだろう。
最後に、できあがったモデルを使って、3枚をまとめてみた。一般的なリアル系モデルとは一味違う肌の質感になっているように思う。やはり学習する元画像は、無修正でしっかり撮影したものを使った方が効果的……を確認できたのではないだろうか。
以上、自分で撮影したグラビアを使い、Stable Diffusion用美女モデルを作成してみた……から1カ月間の後日談と、NVIDIA A100 VRAM 80GB搭載のGPUサーバーを使ったレポートをご紹介した。何かの参考になれば幸いだ。なるのか!?(笑)。完