PyTorch でモデルを保存・読み込む方法

目次

PyTorch のモデルを保存・読み込む方法 state_dict （推奨）

まず、 state_dict モデルの本質的なパラメータ（重みやバイアスなど）のみを保存するため、ファイルサイズを小さく保てて扱いやすい。次に、柔軟性が高い点が挙げられる。state_dict は Python の辞書であるため、モデルのパラメータだけでなくオプティマイザの状態やその他のメタデータも一緒に保存でき、学習の再開や微調整が容易になる。さらに、state_dict は Python の標準 pickle モジュールと互換性があり、異なる環境間でもシンプルに保存・読み込みできるため、モデルの再現性確保が簡単になる。本ガイドでは、PyTorch で state_dict を用いてモデルを保存・読み込みする手順を解説し、効果的なモデル管理のためのベストプラクティスを紹介する。
次に、state_dict は Python の辞書であるため、モデルのパラメータだけでなくオプティマイザの状態やその他のメタデータも保存でき、学習の再開が容易になる。 モデルを微調整する
さらに、state_dict は Python の標準 pickle モジュールと互換性があるため、環境が異なっても保存・読み込みをシンプルに行え、モデルの再ロード手順を容易にする。本ガイドでは、PyTorch で state_dict を用いてモデルを保存・読み込みする手順を示し、効果的なモデル管理のためのベストプラクティスを解説する。

PyTorch でモデルを保存する方法 state_dict

torch.save(model.state_dict(), 'save/to/path/model.pth')

PyTorch でモデルを読み込む方法 state_dict

model = MyModelDefinition(args)
model.load_state_dict(torch.load('load/from/path/model.pth'))

長所

PyTorch は Python に依存しており pickle モジュールで、Python の辞書を簡単にシリアライズ（pickle）およびデシリアライズ（unpickle）し、更新や復元ができる。state_dict を用いると、モデルのパラメータ管理に柔軟性が生まれる。
モデルのパラメータに加えて、オプティマイザの状態やハイパーパラメータなども、キーと値のペアとして に保存できる。 state_dictこれらはモデルを再読み込みする際に容易に参照できる。

短所:

読み込むには state_dict、同一のモデル定義が必要である。これがないと、保存したパラメータを正しく読み込めない。

注意点：

必ず呼び出してください model.eval() 推論時は、ド���ップアウトやバッチ正規化などの層を評価モードに設定すること。
次の方法でモデルを保存します .pt または .pth 一貫性と互換性のための拡張子。

PyTorch のモデル全体を保存・読み込む

PyTorch のモデル全体を保存する

torch.save(model, 'save/to/path/model.pt')

PyTorch のモデル全体を読み込む

model = torch.load('load/from/path/model.pt')

長所:

モデル全体の保存は必要なコード量が最小で、保存・読み込みを素早く行う手軽な方法です。
モデル全体を保存・読み込むための API はシンプルで使いやすい。保存と読み込みの API はより直感的である。

短所:

Python の pickle このモジュールが内部で使用されるため、保存されたモデルは特定のクラス定義やディレクトリ構造に強く結び付けられます。つまり、リファクタリングやファイルパスの変更を行うと、読み込み時に問題が発生する可能性があります。
別プロジェクトでモデル全体の保存ファイルを使用するのは困難です。保存時と同一のディレクトリ構造とクラスパスでなければならないためです。

注意点：

必ず呼び出してください model.eval() ドロップアウトやバッチ正規化などの層が正しく動作するよう、推論の前に必ず適切な設定を行ってください。
次の方法でモデルを保存します .pt または .pth モデルを保存する際は、一貫性と可読性のために拡張子を統一してください。

チェックポイントから PyTorch モデルを保存・読み込みする方法

PyTorch のモデルチェックポイントを保存する

torch.save({'epoch': EPOCH,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': LOSS}, 
	    'save/to/path/model.pth')

PyTorch のモデルチェックポイントを読み込む

model = MyModelDefinition(args)
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

checkpoint = torch.load('load/from/path/model.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

注意点：

学習を再開するときは、必ず呼び出してください model.train() モデルを学習モードに設定するために必ず実行してください。
チェックポイントを読み込んだ後に推論を行う場合は、model.eval() を呼び出してください。 model.eval() モデルを評価モードに設定するために。

モデルのバージョン管理には W&B Artifacts を使用する

PyTorch モデルを W&B Artifacts として保存する

# Import
import wandb
# Save your model.
torch.save(model.state_dict(), 'save/to/path/model.pth')
# Save as artifact for version control.
run = wandb.init(project='your-project-name')
artifact = wandb.Artifact('model', type='model')
artifact.add_file('save/to/path/model.pth')
run.log_artifact(artifact)
run.finish()

PyTorch モデルを W&B Artifacts として読み込む

import wandb
run = wandb.init()

artifact = run.use_artifact('entity/your-project-name/model:v0', type='model')
artifact_dir = artifact.download()

run.finish()

図1ライブの Artifact ダッシュボードを確認する こちら

Weights & Biases を試す

下の緑色の「Run」ボタンをクリックしてください（初回に「Run」をクリックすると、Replit がマシンを割り当てるまで約30〜45秒かかります）
ターミナルウィンドウの指示に従ってください（右下のペイン）
ターミナルウィンドウのサイズを変更できます（右下）を拡大表示できます

参考資料