PyTorchでトレーニング済みモデルを保存する別の方法を探していました。これまでのところ、2つの代替案を見つけました。

1. torch.save（）はモデルを保存し、torch.load（）はモデルをロードします。
2. model.state_dict（）はトレーニング済みモデルを保存し、model.load_state_dict（）は保存されたモデルをロードします。

私はこの議論に出くわしました。アプローチ2はアプローチ1よりも推奨されます。

私の質問は、なぜ2番目のアプローチが好ましいのかということです。それはtorch.nnモジュールがこれら2つの機能を持っているからであり、それらを使用することをお勧めしますか？

このページをgithubリポジトリで見つけたので、ここにコンテンツを貼り付けます。

モデルを保存するための推奨されるアプローチ

モデルのシリアル化と復元には、主に2つの方法があります。

最初の（推奨）は、モデルパラメータのみを保存およびロードします。

torch.save(the_model.state_dict(), PATH)

じゃあ後で：

the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))

2番目は、モデル全体を保存してロードします。

torch.save(the_model, PATH)

じゃあ後で：

the_model = torch.load(PATH)

ただし、この場合、シリアル化されたデータは、特定のクラスと使用される正確なディレクトリ構造にバインドされるため、他のプロジェクトで使用したり、深刻なリファクタリングを行ったりすると、さまざまな方法で破損する可能性があります。

何をしたいかによります。

ケース＃1：モデルを保存して推論に使用する：モデルを保存して復元し、モデルを評価モードに変更します。これが行われるのは、通常、とがBatchNormありDropout、デフォルトでは構築時にトレーニングモードになっているためです。

torch.save(model.state_dict(), filepath)

#Later to restore:
model.load_state_dict(torch.load(filepath))
model.eval()

ケース＃2：モデルを保存して後でトレーニングを再開する：保存しようとしているモデルのトレーニングを継続する必要がある場合は、モデルだけでなく他のものも保存する必要があります。オプティマイザ、エポック、スコアなどの状態も保存する必要があります。次のように実行します。

state = {
    'epoch': epoch,
    'state_dict': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    ...
}
torch.save(state, filepath)

トレーニングを再開するにstate = torch.load(filepath)は、次のようにします。次に、個々のオブジェクトの状態を復元するには、次のようにします。

model.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])

トレーニングを再開しているので、ロード時に状態を復元した後は呼び出さないでmodel.eval()ください。

ケース＃3：コードにアクセスできない他のユーザーが使用するモデル：Tensorflow .pbでは、モデルのアーキテクチャと重みの両方を定義するファイルを作成できます。これは、特にを使用する場合に非常に便利Tensorflow serveです。Pytorchでこれを行う同等の方法は次のとおりです。

torch.save(model, filepath)

# Then later:
model = torch.load(filepath)

この方法はまだ完全な証拠ではありません。pytorchはまだ多くの変更が加えられているため、お勧めしません。

ピクルス Pythonライブラリを実装シリアライズとPythonオブジェクトをデシリアライズするためのバイナリプロトコル。

あなたimport torch（またはPyTorchを使用するとき）はそれをimport pickle実行し、オブジェクトを保存およびロードするメソッドであるpickle.dump()andをpickle.load()直接呼び出す必要はありません。

実際に、torch.save()そしてtorch.load()ラップされますpickle.dump()と、pickle.load()あなたのために。

state_dict他の答えはわずか数より多くのノートに値する言及しました。

state_dictPyTorchの内部には何がありますか？実際には2つstate_dictのがあります。

PyTorchモデルには、学習可能なパラメーター（wおよびb）を取得するための呼び出しがtorch.nn.Moduleありmodel.parameters()ます。これらの学習可能なパラメータは、ランダムに設定されると、時間の経過とともに更新されます。学習可能なパラメータが最初state_dictです。

2つ目state_dictは、オプティマイザの状態辞書です。オプティマイザーは学習可能なパラメーターを改善するために使用されることを思い出してください。しかし、オプティマイザstate_dictは修正されています。そこで学ぶことは何もありません。

state_dictオブジェクトはPython辞書であるため、簡単に保存、更新、変更、復元でき、PyTorchモデルとオプティマイザに大幅なモジュール性を追加します。

これを説明するための超シンプルなモデルを作成しましょう：

import torch
import torch.optim as optim

model = torch.nn.Linear(5, 2)

# Initialize optimizer
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

print("Model weight:")    
print(model.weight)

print("Model bias:")    
print(model.bias)

print("---")
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

このコードは以下を出力します：

Model's state_dict:
weight   torch.Size([2, 5])
bias     torch.Size([2])
Model weight:
Parameter containing:
tensor([[ 0.1328,  0.1360,  0.1553, -0.1838, -0.0316],
        [ 0.0479,  0.1760,  0.1712,  0.2244,  0.1408]], requires_grad=True)
Model bias:
Parameter containing:
tensor([ 0.4112, -0.0733], requires_grad=True)
---
Optimizer's state_dict:
state    {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [140695321443856, 140695321443928]}]

これは最小限のモデルであることに注意してください。シーケンシャルのスタックを追加してみてください

model = torch.nn.Sequential(
          torch.nn.Linear(D_in, H),
          torch.nn.Conv2d(A, B, C)
          torch.nn.Linear(H, D_out),
        )

モデルのにエントリがあるのは、学習可能なパラメーターを持つレイヤー（たたみ込みレイヤー、線形レイヤーなど）と登録されたバッファー（batchnormレイヤー）のみであることに注意してくださいstate_dict。

学習不能なものは、オプティマイザーオブジェクトに属しますstate_dict。オプティマイザーオブジェクトには、オプティマイザーの状態に関する情報と、使用されるハイパーパラメーターが含まれています。

物語の残りは同じです。予測のための推論フェーズ（これはトレーニング後にモデルを使用するフェーズです）。学習したパラメータに基づいて予測します。したがって、推論のために、パラメーターを保存する必要があるだけですmodel.state_dict()。

torch.save(model.state_dict(), filepath)

そして、後でmodel.load_state_dict（torch.load（filepath））model.eval（）を使用する

注：model.eval()モデルをロードした後、これは重要な最後の行を忘れないでください。

また、保存しようとしないでくださいtorch.save(model.parameters(), filepath)。これmodel.parameters()は単なるジェネレータオブジェクトです。

一方、torch.save(model, filepath)モデルオブジェクト自体は保存されますが、モデルにはオプティマイザのがないことに注意してstate_dictください。@Jadiel de Armasによる他の優れた回答をチェックして、オプティマイザの状態の辞書を保存してください。

一般的なPyTorchの規則は、.ptまたは.pthファイル拡張子を使用してモデルを保存することです。

モデル全体の保存/読み込み 保存：

path = "username/directory/lstmmodelgpu.pth"
torch.save(trainer, path)

負荷：

モデルクラスはどこかに定義する必要があります

model = torch.load(PATH)
model.eval()

モデルを保存し、後でトレーニングを再開する場合：

シングルGPU： 保存：

state = {
        'epoch': epoch,
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict(),
}
savepath='checkpoint.t7'
torch.save(state,savepath)

負荷：

checkpoint = torch.load('checkpoint.t7')
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint['epoch']

複数のGPU： 保存

state = {
        'epoch': epoch,
        'state_dict': model.module.state_dict(),
        'optimizer': optimizer.state_dict(),
}
savepath='checkpoint.t7'
torch.save(state,savepath)

負荷：

checkpoint = torch.load('checkpoint.t7')
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint['epoch']

#Don't call DataParallel before loading the model otherwise you will get an error

model = nn.DataParallel(model) #ignore the line if you want to load on Single GPU