abinitio pytorch(十二):モデルの保存とロード

モデルの読み込みと保存
まとめると、いくつかの関数です

torch.load()/torch.save()

pythonのpickleによってシーケンス化と逆シーケンス化が完了する.メモリディスクの変換を完了します.

Module.state_dict()/Module.load_state_dict()

state_dict()はモデルパラメータを取得する.load_state_dict()ロードモデルパラメータ
読み書きTensorsave関数とload関数を直接使用して、それぞれTensorを格納し、読み取ることができます.save Pythonのpickleユーティリティを使用してオブジェクトをシーケンス化し、シーケンス化されたオブジェクトをdiskに保存します.saveを使用すると、モデル、テンソル、辞書など、さまざまなオブジェクトを保存できます.laodはpickle unpickleツールを使用してpickleのオブジェクトファイルをメモリに逆シーケンス化します.次の例では、Tensor変数xを作成し、ファイル名が同じx.ptのファイルに存在します.

import torch
from torch import nn

x = torch.ones(3)
torch.save(x, 'x.pt')

次に、格納されたファイルからメモリにデータを読み込みます.

x2 = torch.load('x.pt')
x2

出力:

tensor([1., 1., 1.])

また、Tensorのリストを格納し、メモリを読み込むこともできます.

y = torch.zeros(4)
torch.save([x, y], 'xy.pt')
xy_list = torch.load('xy.pt')
xy_list

出力:

[tensor([1., 1., 1.]), tensor([0., 0., 0., 0.])]

文字列からTensorにマッピングされた辞書を格納し、読み出します.

torch.save({'x': x, 'y': y}, 'xy_dict.pt')
xy = torch.load('xy_dict.pt')
xy

出力:

{'x': tensor([1., 1., 1.]), 'y': tensor([0., 0., 0., 0.])}

state_dict
PyTorchでは、Moduleの学習可能パラメータ(すなわち、重みおよび偏差)が、model.parameters()によってアクセスされるパラメータにモジュールモデルが含まれる.state_dictは、パラメータ名からパラメータTesnorに非表示にされる辞書オブジェクトである.

class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.hidden = nn.Linear(3, 2)
        self.act = nn.ReLU()
        self.output = nn.Linear(2, 1)

    def forward(self, x):
        a = self.act(self.hidden(x))
        return self.output(a)

net = MLP()
net.state_dict()

出力:

OrderedDict([('hidden.weight', tensor([[ 0.2448,  0.1856, -0.5678],
                      [ 0.2030, -0.2073, -0.0104]])),
             ('hidden.bias', tensor([-0.3117, -0.4232])),
             ('output.weight', tensor([[-0.4556,  0.4084]])),
             ('output.bias', tensor([-0.3573]))])

なお、state_dictのエントリ・オプティマイザ(optim)は、学習可能なパラメータを有する層(ボリューム層、線形層など)のみが、オプティマイザの状態および使用されるスーパーパラメータに関する情報を含むstate_dictも存在する.

optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
optimizer.state_dict()

出力:

{'state': {}, 'param_groups': [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [139952370292992, 139952370293784, 139952370294144, 139952370293496]}]}

モデルの保存とロード
PyTorchでは、トレーニングモデルの保存とロードには2つの一般的な方法があります.

モデルパラメータ(state_dict)

のみ保存およびロード

モデル全体の保存とロード

保存とロードstate_dict(推奨)
保存:

torch.save(model.state_dict(), PATH) #          pt pth

ロード:

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

モデル全体の保存とロード
保存:

torch.save(model, PATH)

ロード:

model = torch.load(PATH)

推奨する方法を用いて実験してみましょう.

X = torch.randn(2, 3)
Y = net(X)

PATH = "./net.pt"
torch.save(net.state_dict(), PATH)

net2 = MLP()
net2.load_state_dict(torch.load(PATH))
Y2 = net2(X)
Y2 == Y

出力:

tensor([[1],
        [1]], dtype=torch.uint8)

このnetとnet2には同じモデルパラメータがあるため、同じ入力Xに対する計算結果は同じになります.上の出力もこれを検証した.
また、GPUとCPUの間のモデルの保存と読み取り、複数のGPUを使用したモデルの記憶など、他にも使用シーンがあり、公式ドキュメントを参照することができます.