Real-ESRGANの追加学習で「CUDA out of memory」が発生した話

AI・機械学習 コンピュータ
AI・機械学習

縮小などで劣化してしまった画像を綺麗な形に拡大処理してくれるAIの一つにReal-ESRGANがあります。学習結果であるモデルを差し替えることで画像の内容にに合わせた処理を行うことが出来るのですが、このモデルを追加学習することで自分好みにカスタマイズすることが出来るようなので試してみました。

GPUがRTX3070(VRAM8G)でOSがWindowsのローカル環境で試したところ、実行まではこぎつけるのですが全く進捗している感じがしませんでした。CPUが60%程の使用率でGPUは全く仕事をしていない感じでした。

Windows環境の問題かと思いLinux(Ubuntu)環境でも試してみました。

GPUはRTX3060(VRAM12G)でOSはUbuntuです。こちらも実行できる環境までは同様に構築出来ました。
学習を実行したところ「CUDA out of memory」が発生しました。

メモリが足りないということで、この場合のメモリはGPUが使うVRAMが足りないことになります。
メインメモリと異なり足りないからと言って簡単にVRAMは増設できるものではないので、Google Colabを借りることにします。

処理時間が長時間に及ぶと予想されるので、月1000円ちょっとの有料プランに加入して試してみます。
使えるVRAMは16GBでした。

実行したところこちらでも「CUDA out of memory」が発生しました。
VRAMが16GBでも足りないことになります。

仕方がないので、月5000円ちょっとの上のプランで試してみます。
実行することが出来ました。利用しているVRAMの容量は16G~17Gぐらいと予想通りでした。
処理時間の関係で完走する前に5000円分を使い切ってしまいました。

更に5000円を追加することも出来ますが、完走までどれだけお金が必要か不明で、完走しない可能性もあるのでこの方法はいったん中止します。なかなかお金がかかるお話です。

ローカル環境に戻って、Windowsのほうですが最近のNVIDAのドライバーはVRAMが足りなくなるとメインメモリー(共有メモリ?)を使うようになったと聞いています。ただメインメモリーを使うと極端に遅くなるようなので、実用に耐えない可能性があります。

何にしても処理能力に対して仕事が過大であると思われるので、処理する画像を小さく切ってみました。

学習用の画像を256×256ぐらいにトリミングした画像で学習したところ進捗するようになりました。メモリの使用状況はVRAMと共有メモリを合わせて17GぐらいでGoogle Colabで実行した状況と同じです。
進捗するようになりましたが、非常に遅く終わりが見えません。ちなみに実行しているWindows環境はメインマシンで、我が家で一番快適なスペックを誇るPCでこの処理(実験)に数日占有されるのは好ましくありません。

ということで、何とかしてUbuntuで実行させたいのですが要求されてるVRAMが17GBぐらいですので12GのRTX3060では足りないことになります。学習画像をさらに小さくすることも考えましたが、あまり小さくしすぎると学習結果にも影響が出ると思われるので別の方法を探しまします。

調べたところ、別のAI学習の記事でメモリが足りない場合batchサイズを小さくすることでVRAMの使用量を制御できるとありました。(batchの意味は理解していません)
Real-ESRGANのoptions/finetune_realesrgan_x4plus.ymlを開きbatchと検索したところ「batch_size_per_gpu: 12」がヒットしましたこちらの値を小さくして(半分の6)実行したところ、「CUDA out of memory」は解消され実行することが出来ました。
nvidia-smiを実行しVRAMの使用状況を見たところ8966MiB / 12288MiBでVRAM内に収まっています。

この記事を書いているタイミングでは、処理の最中で処理完了の予想時間は4日と12時間後とのことですので、気長に待ちたいと思います。

電気代がかかりますが、冬場ですのでPCの排熱で部屋を暖める暖房代だと思うことにします。

追加学習用の縮小画像の作成するコマンド

追加学習用の画像が保存されたディレクトリ:datasets\own_dataset_root\own_images
(ディレクトリ名は任意の名前で良さそう?)

python scripts/generate_multiscale_DF2K.py --input datasets/own_dataset_root/own_images --output datasets/own_dataset_root/own_images_multiscale
python scripts/generate_meta_info.py --input datasets/own_dataset_root/own_images datasets/own_dataset_root/own_images_multiscale --root datasets/own_dataset_root datasets/own_dataset_root --meta_info datasets/own_dataset_root/meta_info/meta_info_own_imagesmultiscale.txt

追加学習を実行するコマンド

options/finetune_realesrgan_x4plus.ymlは追加学習の設定ファイルですので事前に編集する。
・バッチサイズ
・学習画像のパスを実行環境に合わせて編集
・参照モデル名の変更(RealESRNet_x4plus.pth⇒RealESRGAN_x4plus.pth)

python realesrgan/train.py -opt options/finetune_realesrgan_x4plus.yml --auto_resume

追加学習したモデルを指定して拡大処理

python inference_realesrgan.py -n RealESRGAN_x4plus -i images -o results  --outscale 4 -t 512 --suffix ex --model_path experiments\finetune_RealESRGANx4plus_400k\models\net_g_30000.pth;[Console]::Beep(440,200)

追加学習したモデルファイルが保存されたディレクトリ:experiments\finetune_RealESRGANx4plus_400k\models
複数のファイルが作成されているのでファイル名の数字が大きいファイルを使っています。
;[Console]::Beep(440,200)は終了時ビープ音を鳴らすためのPowerShellコマンド。(無くともOK)

画像を分割して拡大処理するオプション

-t サイズ

感想

20240114
学習途中ですが追加学習用に使った画像を拡大処理すると思った通りの画像が出来上がりました。
ほかの画像を処理したところ結果はダメでした。学習完了後また試してみたいと思います。
20240130
数日かけて学習をしてみました。学習素材となる画像を縮小して学習させるのですが、解像度の高い綺麗な画像を多数用意することが難しく、期待した結果は得られませんでした。色々と思いついたこともあるのですが、試すには時間(又はお金)が掛かり過ぎるので、世の中のコンピュータの性能が向上する日を待ちたいと思います。

コメント