模型讀取和儲存

總結下來,就是幾個函式

1. torch.load()/torch.save()

通過python的pickle完成序列化與反序列化.完成記憶體<-->磁碟轉換.

1. Module.state_dict()/Module.load_state_dict()

state_dict()獲取模型引數.load_state_dict()載入模型引數

讀寫Tensor

我們可以直接使用save函式和load函式分別儲存和讀取Tensor。save使用Python的pickle實用程式將物件進行序列化，然後將序列化的物件儲存到disk，使用save可以儲存各種物件,包括模型、張量和字典等。而laod

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

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中儲存和載入訓練模型有兩種常見的方法:

1. 僅儲存和載入模型引數(state_dict)
2. 儲存和載入整個模型

儲存和載入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的模型的儲存等等，使用的時候可以參考官方文件