学習済みモデルのpthファイルを軽量化する

ディープラーニングモデルは年々サイズが大きくなり、GPUメモリも大きな容量が必要となっています。
Pretrainedモデルのpthをロードする際に、CUDA out of memory となることもしばしばですね。
ここでは学習時に保存されたオプティマイザの状態を削除したpthファイルを作成することでそれを回避できる場合があるので方法を紹介します。
(もちろんGPUリソースを湯水のごとく使えるならそれが一番ですね!)

解決方法

以下のコードで、オプティマイザの状態を削除した軽量なチェックポイントを作成しますのが：

import torch

# 保存されたチェックポイントファイルのパス
checkpoint_path = "path/to/your_checkpoint.pth"
new_checkpoint_path = "path/to/new_checkpoint.pth"

# チェックポイントをCPUでロード
checkpoint = torch.load(checkpoint_path, map_location="cpu")

# オプティマイザの状態を削除
if "optimizer" in checkpoint:
    del checkpoint["optimizer"]

# 新しいチェックポイントとして保存
torch.save(checkpoint, new_checkpoint_path)

効果

軽量化することで、ややVRAMが少ないGPUでもモデルをロードできる場合があります。
学習は強力なGPUで行い、推論はやや非力なGPUで作業するといったかたちで、マシンリソースの使い分けができるメリットもあります。
(私のケースでは24GB強のチェックポイントが、10GB弱になりました)

注意

学習時にもこの方法が使うこともできますが、学習の連続性については採用するオプティマイザによって影響が異なります。特に学習の中断・再開時にはご注意ください。
こちらの記事が参考になりました。