Pytorch保存和加載模型完全指南: 關於使用Pytorch讀寫模型的一切方法

本文是一篇關於如何用Pytorch保存和加載模型的指南。

本文主要涉及到3個函數：

1. torch.save: 使用Python的pickle實用程序將對象進行序列化，然後將序列化的對象保存到disk，可以保存各種對象,包括模型、張量和字典等。
2. torch.load: 使用pickle unpickle工具將pickle的對象文件反序列化爲內存。
3. torch.nn.Module.load_state_dict: 用反序列化的state_dict來加載模型參數。

1 讀寫tensor

1.1 單個張量

import torch

x = torch.tensor([3.,4.])
torch.save(x, 'x.pt')
x1 = torch.load('x.pt')
print(x1)

輸出：

tensor([3., 4.])

1.2 張量列表和張量詞典

y = torch.ones((4,2))
torch.save([x,y],'xy.pt')
torch.save({'x':x, 'y':y}, 'xy_dict.pt')
xy = torch.load('xy.pt')
xy_dict = torch.load('xy_dict.pt')
print(xy)
print(xy_dict)

輸出：

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

2 保存和加載模型

2.1 state_dict

state_dict是一個從每一個層的名稱映射到這個層的參數Tesnor的字典對象。

注意，只有具有可學習參數的層(卷積層、線性層等)和註冊緩存(batchnorm’s running_mean)纔有state_dict中的條目。優化器(torch.optim)也有一個state_dict，其中包含關於優化器狀態以及所使用的超參數的信息。

from torch import nn
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()
print(net.state_dict())
print('\n',net.state_dict()['output.weight'])

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

輸出：

OrderedDict([('hidden.weight', tensor([[ 0.1493, -0.1645,  0.5165],
        [-0.0773,  0.3487,  0.4961]])), ('hidden.bias', tensor([-0.4560, -0.5111])), ('output.weight', tensor([[0.3356, 0.4229]])), ('output.bias', tensor([-0.6374]))])

 tensor([[0.3356, 0.4229]])
{'state': {},
 'param_groups': [{'lr': 0.001,
   'momentum': 0.9,
   'dampening': 0,
   'weight_decay': 0,
   'nesterov': False,
   'params': [139788808398168,
    139788808397952,
    139788808397736,
    139788808398096]}]}

2.2 保存和加載

PyTorch中保存和加載訓練模型有兩種常見的方法:

1. 僅保存和加載模型參數(state_dict)；
2. 保存和加載整個模型。

2.2.1 保存和加載state_dict(推薦方式)

torch.save(net.state_dict(), 'net_state_dict.pt')## 後綴名一般寫爲: .pt或.pth
net1 = MLP()
net1.load_state_dict(torch.load('net_state_dict.pt'))

輸出：

<All keys matched successfully>

注意: load_state_dict() 接受一個詞典對象，而不是一個指向對象的路徑。因此你需要先使用torch.load()來反序列化。比如，你不能直接這麼用model.load_state_dict(PATH)。

2.2.2 保存和讀寫整個模型

這個就相對來說比較簡單了。

torch.save(net, 'net.pt')
net2 = torch.load('net.pt')

**注意：**以這種方式保存模型將使用Python的pickle模塊保存整個模塊。 這種方法的缺點是序列化的數據被綁定到特定的類，並且在保存模型時使用了確切的詞典結構。 這樣做的原因是因爲pickle不會保存模型類本身。 而是將其保存到包含這個類的文件的路徑，該路徑在加載時使用。 因此，在其他項目中使用或重構後，您的代碼可能會以各種方式中斷。

2.2.3 保存和加載checkpiont

## Save
torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

###########################
## Load
model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

保存用於檢查或繼續訓練的checkpiont時，你必須保存的不只是模型的state_dict。 保存優化器的state_dict也很重要，因爲它包含隨着模型訓練而更新的緩衝區和參數。 你可能要保存的其他項目包括你未設置的時間段，最新記錄的訓練損失，外部torch.nn.Embedding層等。

2.2.4 在一個文件中保存多個模型

#Save
torch.save({
            'modelA_state_dict': modelA.state_dict(),
            'modelB_state_dict': modelB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            ...
            }, PATH)


#Load
modelA = TheModelAClass(*args, **kwargs)
modelB = TheModelBClass(*args, **kwargs)
optimizerA = TheOptimizerAClass(*args, **kwargs)
optimizerB = TheOptimizerBClass(*args, **kwargs)

checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])

常見的PyTorch約定是使用.tar文件擴展名保存這些檢查點。

2.3 使用來自不同模型的參數進行模型熱啓動

這種方法一般用於遷移學習。利用經過訓練的參數，即使只有少數幾個可用的參數，也將有助於熱啓動訓練過程，並希望與從頭開始訓練相比，可以更快地收斂模型。

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

modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)

無論是從缺少某些鍵的部分state_dict加載，還是加載比要加載的模型更多的key的state_dict，都可以在load_state_dict()函數中將strict參數設置爲False，以忽略不匹配鍵。

如果你想要將一個層的參數加載到另一個層，但是一些keys不匹配，你只需改變你所加載的state_dict中的名稱即可。

3 跨設備保存和加載模型

3.1 在GPU中保存，在CPU中加載

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

device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))

3.2 在GPU中保存，在GPU中加載

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

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)

3.3 在CPU中保存，在GPU中加載

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

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
model.to(device)

4 保存torch.nn.DataParallel的模型

torch.save(model.module.state_dict(), PATH)

# Load to whatever device you want,加載方法使用常規方式即可。

---

參考鏈接:

1. 官方文檔
2. Dive-into-DL-PyTorch