【PyTorch】モデルの保存・読み込み時のエラー解決法

はじめに

PyTorchで発生するモデル保存/読込エラーの原因と解決法を解説します。

症状

PyTorchでモデルを保存・読み込みする際に、KeyError、RuntimeError、UnpicklingErrorなどが発生します。

原因

1. state_dictのキー不一致

モデル構造を変更した後、古いチェックポイントを読み込もうとしている場合です。

2. GPU/CPUデバイスの不一致

GPUで保存したモデルをCPU環境で読み込む際にエラーが出ます。

解決方法

手順1: 正しい保存方法

# 推奨: state_dictのみ保存
torch.save(model.state_dict(), 'model.pth')

# 読み込み
model.load_state_dict(torch.load('model.pth'))

手順2: デバイス変換付き読み込み

# GPU→CPU
model.load_state_dict(
    torch.load('model.pth', map_location='cpu')
)

# strict=Falseでキー不一致を許容
model.load_state_dict(
    torch.load('model.pth'), strict=False
)

まとめ

PyTorchのモデル保存/読込エラーについて解説しました。この記事が問題解決の参考になれば幸いです。