PyTorchでモデルを保存および読み込む方法

目次

PyTorchモデルを保存およびロードstate_dict（推奨）

最初に、state_dictモデルの重要なパラメータ（ウェイトやバイアスなど）のみを保存し、ファイルサイズを小さくして簡単に操作できるようにします。第二に、柔軟性を提供します。state_dictはPythonの辞書であるため、モデルパラメータだけでなくオプティマイザの状態や他のメタデータも保存でき、トレーニングの再開やモデルの微調整が容易になります。さらに、state_dictは異なる環境間でのモデルのリロードを簡単にし、Pythonの組み込みモジュールであるpickleと互換性があるため、保存と読み込みが簡単です。このガイドでは、state_dictを使用してPyTorchでモデルを保存および読み込む手順を説明し、効果的なモデル管理のベストプラクティスを説明します。
第二に、state_dictはPythonの辞書なので、モデルのパラメータだけでなく、オプティマイザの状態やその他のメタデータも保存でき、トレーニングの再開が容易になります。モデルを微調整する。
さらに、state_dict はさまざまな環境でモデルを再ロードするプロセスを簡略化し、Python の組み込みモジュールである pickle と互換性があるため、簡単に保存およびロードできます。このガイドでは、state_dict を使用して PyTorch でモデルを保存およびロードする手順を説明し、効果的なモデル管理のためのベストプラクティスを説明します。

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辞書を簡単にピック、アンピック、更新、および復元できるようにします。 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() モデルを評価モードに設定する。

モデルのバージョン管理にはW&Bアーティファクトを使用する

PyTorchモデルをW&Bアーティファクトとして保存

# 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()

W&BアーティファクトとしてPyTorchモデルを読み込む

import wandb
run = wandb.init()

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

run.finish()

図1ライブアーティファクトダッシュボードをチェックしてくださいここ

Weights & Biasesを試す

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

おすすめの読書