Pytorch：模型的保存與加載 torch.load()、torch.nn.Module.load_state_dict()

Pytorch：模型的保存與加載 torch.save()、torch.load()、torch.nn.Module.load_state_dict()

Pytorch 保存和加載模型后綴：.pt 和.pth

1 torch.save() [source]
保存一個序列化（serialized）的目標到磁盤。函數使用了Python的pickle程序用於序列化。模型（models），張量（tensors）和文件夾（dictionaries）都是可以用這個函數保存的目標類型。

torch.save(obj, f, pickle_module=<module '...'>, pickle_protocol=2)
參數 描述
obj 保存對象
f 類文件對象 (必須實現寫和刷新)或一個保存文件名的字符串
pickle_modul 用於 pickling 元數據和對象的模塊
 pickle_protocol 指定 pickle protocal 可以覆蓋默認參數
栗子
保存整個模型：

torch.save(model,'save.pt')
只保存訓練好的權重：

torch.save(model.state_dict(), 'save.pt')
2 torch.load() [source]
用來加載模型。torch.load() 使用 Python 的 解壓工具（unpickling）來反序列化 pickled object 到對應存儲設備上。首先在 CPU 上對壓縮對象進行反序列化並且移動到它們保存的存儲設備上，如果失敗了（如：由於系統中沒有相應的存儲設備），就會拋出一個異常。用戶可以通過 register_package 進行擴展，使用自己定義的標記和反序列化方法。

torch.load(f, map_location=None, pickle_module=<module 'pickle' from '...'>)
參數 描述
f 類文件對象 (返回文件描述符)或一個保存文件名的字符串
map_location 一個函數或字典規定如何映射存儲設備
pickle_module 用於 unpickling 元數據和對象的模塊 (必須匹配序列化文件時的 pickle_module )
栗子
torch.load('tensors.pt')

# Load all tensors onto the CPU
torch.load('tensors.pt', map_location=torch.device('cpu'))

# Load all tensors onto the CPU, using a function
torch.load('tensors.pt', map_location=lambda storage, loc: storage)

# Load all tensors onto GPU 1
torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))

# Map tensors from GPU 1 to GPU 0
torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})

# Load tensor from io.BytesIO object
with open('tensor.pt') as f:
buffer = io.BytesIO(f.read())
torch.load(buffer)
3 torch.nn.Module.load_state_dict(state_dict) [source]
使用 state_dict 反序列化模型參數字典。用來加載模型參數。將 state_dict 中的 parameters 和 buffers 復制到此 module 及其子節點中。

torch.nn.Module.load_state_dict(state_dict, strict=True)
參數 描述
state_dict 保存 parameters 和 persistent buffers 的字典
strict 可選，bool型。state_dict 中的 key 是否和 model.state_dict() 返回的 key 一致。
栗子
torch.save(model,'save.pt')
model.load_state_dict(torch.load("save.pt")) #model.load_state_dict()函數把加載的權重復制到模型的權重中去
3.1 什么是state_dict？
在PyTorch中，一個torch.nn.Module模型中的可學習參數（比如weights和biases），模型的參數通過model.parameters()獲取。而state_dict就是一個簡單的Python dictionary，其功能是將每層與層的參數張量之間一一映射。注意，只有包含了可學習參數（卷積層、線性層等）的層和已注冊的命令（registered buffers，比如batchnorm的running_mean）才有模型的state_dict入口。優化方法目標（torch.optim）也有state_dict，其中包含的是關於優化器狀態的信息和使用到的超參數。

因為state_dict目標是Python dictionaries，所以它們可以很輕松地實現保存、更新、變化和再存儲，從而給PyTorch模型和優化器增加了大量的模塊化（modularity）。

torch.nn.Module.state_dict

torch.nn.Module.state_dict(destination=None, prefix='', keep_vars=False)
返回一個包含模型狀態信息的字典。包含參數（weighs and biases）和持續的緩沖值（如：觀測值的平均值）。只有具有可更新參數的層才會被保存在模型的 state_dict 數據結構中。

栗子：

module.state_dict().keys()

# ['bias', 'weight']
torch.optim.Optimizer.state_dict

torch.optim.Optimizer.state_dict()
返回一個包含優化器狀態信息的字典。包含兩個 key：

state：字典，保存當前優化器的狀態信息。不同優化器內容不同。
param_groups：字典，包含所有參數組（eg：超參數）。
栗子：
from __future__ import print_function, division

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy


# 定義模型
class TheModelClass(nn.Module):
def __init__(self):
super(TheModelClass, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)

def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x

# 初始化模型
model = TheModelClass()

# 初始化優化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 打印模型的 state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
print(param_tensor, "\t", model.state_dict()[param_tensor].size())

# 打印優化器的 state_dict
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:
conv1.weight torch.Size([6, 3, 5, 5])
conv1.bias torch.Size([6])
conv2.weight torch.Size([16, 6, 5, 5])
conv2.bias torch.Size([16])
fc1.weight torch.Size([120, 400])
fc1.bias torch.Size([120])
fc2.weight torch.Size([84, 120])
fc2.bias torch.Size([84])
fc3.weight torch.Size([10, 84])
fc3.bias torch.Size([10])
Optimizer's state_dict:
state {}
param_groups [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]
4 保存/加載
4.1 state_dict（推薦）
保存：
torch.save(model.state_dict(), PATH)
加載：
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
保存模型的推理過程的時候，只需要保存模型訓練好的參數，使用torch.save()保存state_dict，能夠方便模型的加載。因此推薦使用這種方式進行模型保存。

記住一定要使用model.eval()來固定dropout和歸一化層，否則每次推理會生成不同的結果。

注意，load_state_dict()需要傳入字典對象，因此需要先反序列化state_dict再傳入load_state_dict()

4.2 整個模型
保存
torch.save(model, PATH)
加載
# 模型類必須在別的地方定義
model = torch.load(PATH)
model.eval()
這種保存/加載模型的過程使用了最直觀的語法，所用代碼量少。這使用Python的pickle保存所有模塊。這種方法的缺點是，保存模型的時候，序列化的數據被綁定到了特定的類和確切的目錄。這是因為pickle不保存模型類本身，而是保存這個類的路徑，並且在加載的時候會使用。因此，當在其他項目里使用或者重構的時候，加載模型的時候會出錯。

一般來說，PyTorch的模型以.pt或者.pth文件格式保存。

一定要記住在評估模式的時候調用model.eval()來固定dropout和批次歸一化。否則會產生不一致的推理結果。

4.3 保存加載用於推理的常規Checkpoint/或繼續訓練
保存
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
...
}, PATH)
加載
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']

model.eval()
# - 或者 -
model.train()
在保存用於推理或者繼續訓練的常規檢查點的時候，除了模型的state_dict之外，還必須保存其他參數。保存優化器的state_dict也非常重要，因為它包含了模型在訓練時候優化器的緩存和參數。除此之外，還可以保存停止訓練時epoch數，最新的模型損失，額外的torch.nn.Embedding層等。

要保存多個組件，則將它們放到一個字典中，然后使用torch.save()序列化這個字典。一般來說，使用.tar文件格式來保存這些檢查點。

加載各個組件，首先初始化模型和優化器，然后使用torch.load()加載保存的字典，然后可以直接查詢字典中的值來獲取保存的組件。

同樣，評估模型的時候一定不要忘了調用model.eval()。

4.4 保存多個模型到一個文件
保存
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)
加載
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'])

modelA.eval()
modelB.eval()
# - 或者 -
modelA.train()
modelB.train()
保存的模型包含多個torch.nn.Modules時，比如GAN，一個序列-序列模型，或者組合模型，使用與保存常規檢查點的方式來保存模型。也就是說，保存每個模型的state_dict和對應的優化器到一個字典中。我們可以保存任何能幫助我們繼續訓練的東西到這個字典中。

4.5 使用其他模型來預熱當前模型
保存
torch.save(modelA.state_dict(), PATH)
加載
modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)
在遷移學習或者訓練新的復雜模型時，加載部分模型是很常見的。利用經過訓練的參數，即使只有少數參數可用，也將有助於預熱訓練過程，並且使模型更快收斂。

在加載部分模型參數進行預訓練的時候，很可能會碰到鍵不匹配的情況（模型權重都是按鍵值對的形式保存並加載回來的）。因此，無論是缺少鍵還是多出鍵的情況，都可以通過在load_state_dict()函數中設定strict參數為False來忽略不匹配的鍵。

如果想將某一層的參數加載到其他層，但是有些鍵不匹配，那么修改state_dict中參數的key可以解決這個問題。

4.6 跨設備保存與加載模型
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))
當在CPU上加載一個GPU上訓練的模型時，在torch.load()中指定map_location=torch.device('cpu')，此時，map_location動態地將tensors的底層存儲重新映射到CPU設備上。

上述代碼只有在模型是在一塊GPU上訓練時才有效，如果模型在多個GPU上訓練，那么在CPU上加載時，會得到類似如下錯誤：

KeyError: ‘unexpected key “module.conv1.weight” in state_dict’
原因是在使用多GPU訓練並保存模型時，模型的參數名都帶上了module前綴，因此可以在加載模型時，把key中的這個前綴去掉：

# 原始通過DataParallel保存的文件
state_dict = torch.load('myfile.pth.tar')
# 創建一個不包含`module.`的新OrderedDict
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # 去掉 `module.`
new_state_dict[name] = v
# 加載參數
model.load_state_dict(new_state_dict)
GPU上保存，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")) # 選擇希望使用的GPU
model.to(device)
在把GPU上訓練的模型加載到GPU上時，只需要使用model.to(torch.devie('cuda'))將初始化的模型轉換為CUDA優化模型。同時確保在模型所有的輸入上使用.to(torch.device('cuda'))。注意，調用my_tensor.to(device)會返回一份在GPU上的my_tensor的拷貝。不會覆蓋原本的my_tensor，因此要記得手動將tensor重寫：my_tensor = my_tensor.to(torch.device('cuda'))。

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)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model
在 GPU 上加載 CPU 訓練保存的模型時，將 torch.load() 函數的 map_location 參數 設置為 cuda:device_id。這種方式將模型加載到指定設備。下一步，確保調用 model.to(torch.device(‘cuda’)) 將模型參數 tensor 轉換為 cuda tensor。最后，確保模型輸入使用 .to(torch.device(‘cuda’)) 為 cuda 優化模型准備數據。

注意：調用 my_tensor.to(device) 會在 GPU 上返回 my_tensor 的新副本，不會覆蓋 my_tensor。因此，使用 my_tensor = my_tensor.to(torch.device(‘cuda’)) 手動覆蓋。

4.7 保存torch.nn.DataParallel模型
保存
torch.save(model.state_dict(), PATH)
加載

# Load to whatever device you want
torch.nn.DataParallel 是支持模型使用 GPU 並行的封裝器。要保存一個一般的 DataParallel 模型， 請保存 model.module.state_dict()。這種方式，可以靈活地以任何方式加載模型到任何設備上。

5 參考
https://blog.csdn.net/LXYTSOS/article/details/90639524
https://blog.csdn.net/cpongo3/article/details/93624288
————————————————
版權聲明：本文為CSDN博主「寧靜致遠*」的原創文章，遵循CC 4.0 BY-SA版權協議，轉載請附上原文出處鏈接及本聲明。
原文鏈接：https://blog.csdn.net/weixin_40522801/article/details/106563354

 

-----------------------------------------------------------------

PyTorch模型保存深入理解

前面寫過一篇PyTorch保存模型的文章：Pytorch模型保存與加載，並在加載的模型基礎上繼續訓練 ，簡單介紹了PyTorch模型保存和加載函數的用法，足以快速上手，但對相關函數和參數的具體用法和代表的含義沒有進行展開介紹，這篇文章用於記錄之。

PyTorch保存模型的語句是這樣的：

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

加載是這樣的：

model.load_state_dict(torch.load(path)) 

下面我們將其拆開逐句介紹，深入理解。

1.torch.save()和torch.load()

顧名思義，save函數是PyTorch的存儲函數，load函數則是讀取函數。save函數可以將各種對象保存至磁盤，包括張量，列表，ndarray，字典，模型等；而相應地，load函數將保存在磁盤上的對象讀取出來。

用法：

torch.save(保存對象, 保存路徑)
torch.load(文件路徑)

應用舉例：

保存張量

In [3]: a = torch.ones(3)                                                       

In [4]: a                                                                       
Out[4]: tensor([1., 1., 1.])

In [5]: torch.save(a, './a.pth')          # 保存Tensor               

In [6]: a_load = torch.load('./a.pth')    # 讀取Tensor

In [7]: a_load                                                                  
Out[7]: tensor([1., 1., 1.])

保存字典

In [11]: b = {k:v for v,k in enumerate('abc',1)} In [12]: b Out[12]: {'a': 1, 'b': 2, 'c': 3} In [13]: torch.save(b, './b.rar') In [14]: torch.load('./b.rar') Out[14]: {'a': 1, 'b': 2, 'c': 3} 

可以看出，保存和讀取非常方便。這里需要注意的是文件的命名，命名必須要有擴展名，擴展名可以為‘xxx.pt’，‘xxx.pth’，‘xxx.pkl’，‘xxx.rar’等形式。

2.model.state_dict()

在PyTorch中，state_dict是一個從參數名稱隱射到參數Tesnor的字典對象。

In [15]: 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) ...: In [16]: net = MLP() In [17]: net.state_dict() Out[17]: OrderedDict([('hidden.weight', tensor([[ 0.4839, 0.0254, 0.5642], [-0.5596, 0.2602, -0.5235]])), ('hidden.bias', tensor([-0.4986, -0.5426])), ('output.weight', tensor([[0.0967, 0.4980]])), ('output.bias', tensor([-0.4520]))]) 

可以看出，state_dict()返回的是一個有序字典，該字典的鍵即為模型定義中有可學習參數的層的名稱+weight或+bias，值則對應相應的權重或偏差，無參數的層則不在其中。

state_dict()返回模型可學習參數的鍵值對，那么問題來了：model.parameters()不也是模型可學習參數的訪問方式嗎？它們有什么不同？

我們再來對比了一下常用的model.parameters()，在訓練時，我們常常將此語句放在優化器中，表示要優化學習的模型參數。關於model.parameters(), model.named_parameters(), model.children(), model.named_children(), model.modules(), model.named_modules()請看我的這篇博客

In [18]: net.parameters() Out[18]: <generator object Module.parameters at 0x7f009b78acf0> In [19]: list(net.parameters()) Out[19]: [Parameter containing: tensor([[ 0.4839, 0.0254, 0.5642], [-0.5596, 0.2602, -0.5235]], requires_grad=True), Parameter containing: tensor([-0.4986, -0.5426], requires_grad=True), Parameter containing: tensor([[0.0967, 0.4980]], requires_grad=True), Parameter containing: tensor([-0.4520], requires_grad=True)] ####################################################################### In [20]: net.named_parameters() Out[20]: <generator object Module.named_parameters at 0x7f00a009a390> In [21]: list(net.named_parameters()) Out[21]: [('hidden.weight', Parameter containing: tensor([[ 0.4839, 0.0254, 0.5642], [-0.5596, 0.2602, -0.5235]], requires_grad=True)), ('hidden.bias', Parameter containing: tensor([-0.4986, -0.5426], requires_grad=True)), ('output.weight', Parameter containing: tensor([[0.0967, 0.4980]], requires_grad=True)), ('output.bias', Parameter containing: tensor([-0.4520], requires_grad=True))] 

可以看出，model.parameters()是一個生成器，每個參數張量都是一個參數容器，它的對象是各個參數Tensor。這里簡單的提一下model.named_parameters()，它和model.parameters很像，都返回一個可迭代對象，但從上例可以看出它多返回一個參數名稱，這樣有利於訪問和初始化或修改參數。
我們在用優化器優化參數時，優化對象是純參數，所以用model.parameters():

In [22]: optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9) In [23]: optimizer.state_dict() Out[23]: {'state': {}, 'param_groups': [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [139638712861752, 139640593827520, 139640585217080, 139640585217296]}]} 

除了模型中可學習參數的層(卷積層、線性層等)有state_dict，優化器也有一個state_dict，其中包含關於優化器狀態以及所使用的超參數的信息。

3.model.load_state_dict()

這是模型加載state_dict的語句，也就是說，它的輸入是一個state_dict，也就是一個字典。模型定義好並且實例化后會自動進行初始化，上面的例子中我們定義的模型MLP在實例化以后顯示的模型參數都是自動初始化后的隨機數。
在訓練模型或者遷移學習中我們會使用已經訓練好的參數來加速訓練過程，這時候就用load_state_dict()語句加載訓練好的參數並將其覆蓋在初始化參數上，也就是說執行過此語句后，加載的參數將代替原有的模型參數。
既然加載的是一個字典，那么需要注意的就是字典的鍵一定要相同才能進行覆蓋，比如加載的字典中的'hidden.weight'只能覆蓋當前模型的'hidden.weight'，如果鍵不同，則不能實現有效覆蓋操作。鍵相同而值的shape不同，則會將新的鍵值對覆蓋原來的鍵值對，這樣在訓練時會報錯。所以我們在加載前一般會進行數據篩選，篩選是對字典的鍵進行對比來操作的：

pretrained_dict = torch.load(log_dir)  # 加載參數字典
model_state_dict = model.state_dict()  # 加載模型當前狀態字典
pretrained_dict_1 = {k:v for k,v in pretrained_dict.items() if k in model_state_dict}  # 過濾出模型當前狀態字典中沒有的鍵值對
model_state_dict.update(pretrained_dict_1)  # 用篩選出的參數鍵值對更新model_state_dict變量
model.load_state_dict(model_state_dict)  # 將篩選出的參數鍵值對加載到模型當前狀態字典中

以上代碼簡單的對預訓練參數進行了過濾和篩選，主要是通過第3條語句粗略的過濾了鍵值對信息，進行篩選后要用Python更新字典的方法update()來對模型當前字典進行更新，update()方法將pretrained_dict_1中的鍵值對添加到model_state_dict中，若pretrained_dict_1中的鍵名和model_state_dict中的鍵名相同，則覆蓋之；若不同，則作為新增鍵值對添加到model_state_dict中。顯然，這里需要的是將pretrained_dict_1中的鍵值對覆蓋model_state_dict的相應鍵值對，所以對應的鍵的名稱必須相同，所以第3條語句中按鍵名稱進行篩選，過濾出當前模型字典中沒有的鍵值對。否則會報錯。
如果想要細粒度過濾或更改某些參數的維度，如進行卷積核參數維度的調整，假如預訓練參數里conv1有256個卷積核，而當前模型只需要200個卷積核，那么可以采用類似以下語句直接對字典進行更改：

pretrained_dict['conv1.weight'] = pretrained_dict['conv1.weight'][:200,:,:,:]   # 假設保留前200個卷積核

以上。

 

原文鏈接：https://www.jianshu.com/p/6c558300130f