PyTorchモデルをGPUメモリからアンロードする
はじめに
機械学習モデルはそのサイズが増大する一方です。また、1つのプロジェクトで複数の機械学習モデルを使うことも珍しくありません。
しかし、潤沢なGPUリソースを持っている方は限られているかと思います。私もその1人で、自宅のGPUはメモリが16GBしかありません。
そこで今回は、PyTorchモデルをGPUメモリからアンロード（GPUメモリを解放）する方法を調べてみました。
Pythonプロセスを実行したままモデルをアンロードすることができれば、モデルを入れ替えながら推論するのも容易になります。
なお、Pythonプロセスを実行したままにするのは、メインメモリは比較的潤沢にあり、キャッシュを有効に使いたいためです。
Pythonプロセスを終了して良いのであれば、それが一番確実にGPUメモリを解放できます。
なお、TorchServe、TensorFlow Serving、NVIDIA Triton Inference Serverなど、機械学習モデルのライフサイクルを管理できるツールは存在しますが、モデルの変換、前処理/後処理の移植など、小さいプロジェクトでは採用しづらい導入コスト、運用コストがあるため、ここでは触れません。
結論
PyTorchのモデルをGPUメモリからアンロードするのは簡単でした。方法は以下の通りです。
import gc
import torch
model = load_model().to("cuda")
del model
gc.collect()
torch.cuda.empty_cache()
モデルへの参照をすべて削除したあと、gc.collect()でガベージコレクションを実行し、torch.cuda.empty_cache()を実行することで、完璧ではないにせいよGPUメモリを解放することができました。
調査した中には「to("cpu")してCPUに転送する」という情報もありましたが、これを実行せずともGPUメモリは解放されたので、少なくとも現時点のPyTorch（バージョン2.5.1）では不要そうです。
具体例: Jina CLIP v2をロード→アンロード
具体例として、Jina CLIP v2をロードし、アンロードしてみます。
Jina CLIP v2については別に記事を書いたので、そちらをご参照ください。
以下、PythonのREPL環境で1行1行実行しながら、nvidia-smiコマンドでGPUメモリの使用量を調べた結果です。
import gc
import torch
from transformers import AutoModel
# 0 MiB
model = AutoModel.from_pretrained("jinaai/jina-clip-v2", trust_remote_code=True).to("cuda")
# 0 MiB → 1914 MiB
del model
# 1914 MiB
gc.collect()
# 1914 MiB
torch.cuda.empty_cache()
# 1914 MiB → 236 MiB
model = AutoModel.from_pretrained("jinaai/jina-clip-v2", trust_remote_code=True).to("cuda")
# 236 MiB → 1914 MiB
del model
# 1914 MiB
gc.collect()
# 1914 MiB
torch.cuda.empty_cache()
# 1914 MiB → 236 MiB
モデルがアンロードされ、GPUメモリの使用量が1914 MiBから236 MiBに低下しているのが分かります。
この236 MiBについてはよく分かっていませんが、再度モデルをロードし、アンロードしてもGPUメモリの最大値、最小値は変わらないので「最低限必要なGPUメモリ」として捉え、深追いはしていません。
参考
本記事の作成にあたり、以下のページを参考にしました。
おわりに
割と簡単にPyTorchモデルをGPUメモリからアンロードすることができました。
今後は実際のサービスに組み込み、GPUメモリをやりくりできるかどうか試してみたいと思います。
本記事が何らかの参考になれば幸いです。
Discussion