PyTorch模型保存與加載實例詳解
- torch.save：保存序列化的對象到磁盤，使用了Python的pickle進行序列化，模型、張量、所有對象的字典。
- torch.load：使用了pickle的unpacking將pickled的對象反序列化到內(nèi)存中。
- torch.nn.Module.load_state_dict：使用反序列化的state_dict加載模型的參數(shù)字典。
state_dict 是一個Python字典，將每一層映射成它的參數(shù)張量。注意只有帶有可學習參數(shù)的層（卷積層、全連接層等），以及注冊的緩存（batchnorm的運行平均值）在state_dict 中才有記錄。state_dict同樣包含優(yōu)化器對象，存儲了優(yōu)化器的狀態(tài)，所使用到的超參數(shù)。
一個簡單的例子
# 定義模型
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()
# 初始化優(yōu)化器
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())
# 打印優(yōu)化器的 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]}]
保存/加載 state_dict（推薦）
保存：
torch.save(model.state_dict(), PATH)
加載：
model = TheModelClass(*args, **kwargs) model.load_state_dict(torch.load(PATH)) model.eval()
要注意這個細節(jié)，如果使用nn.DataParallel在一臺電腦上使用了多個GPU，那么加載模型的時候也必須先進行nn.DataParallel。
保存模型的推理過程的時候，只需要保存模型訓練好的參數(shù)，使用torch.save()保存state_dict，能夠方便模型的加載。因此推薦使用這種方式進行模型保存。
記住一定要使用model.eval()來固定dropout和歸一化層，否則每次推理會生成不同的結果。
注意，load_state_dict()需要傳入字典對象，因此需要先反序列化state_dict再傳入load_state_dict()
保存/加載整個模型
保存：
torch.save(model, PATH)
加載：
# 模型類必須在別的地方定義 model = torch.load(PATH) model.eval()
這種保存/加載模型的過程使用了最直觀的語法，所用代碼量少。這使用Python的pickle保存所有模塊。這種方法的缺點是，保存模型的時候，序列化的數(shù)據(jù)被綁定到了特定的類和確切的目錄。這是因為pickle不保存模型類本身，而是保存這個類的路徑，并且在加載的時候會使用。因此，當在其他項目里使用或者重構的時候，加載模型的時候會出錯。
一般來說，PyTorch的模型以.pt或者.pth文件格式保存。
一定要記住在評估模式的時候調(diào)用model.eval()來固定dropout和批次歸一化。否則會產(chǎn)生不一致的推理結果。
保存加載用于推理的常規(guī)Checkpoint/或繼續(xù)訓練
保存：
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()
在保存用于推理或者繼續(xù)訓練的常規(guī)檢查點的時候，除了模型的state_dict之外，還必須保存其他參數(shù)。保存優(yōu)化器的state_dict也非常重要，因為它包含了模型在訓練時候優(yōu)化器的緩存和參數(shù)。除此之外，還可以保存停止訓練時epoch數(shù)，最新的模型損失，額外的torch.nn.Embedding層等。
要保存多個組件，則將它們放到一個字典中，然后使用torch.save()序列化這個字典。一般來說，使用.tar文件格式來保存這些檢查點。
加載各個組件，首先初始化模型和優(yōu)化器，然后使用torch.load()加載保存的字典，然后可以直接查詢字典中的值來獲取保存的組件。
同樣，評估模型的時候一定不要忘了調(diào)用model.eval()。
保存多個模型到一個文件
保存：
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，一個序列-序列模型，或者組合模型，使用與保存常規(guī)檢查點的方式來保存模型。也就是說，保存每個模型的state_dict和對應的優(yōu)化器到一個字典中。我們可以保存任何能幫助我們繼續(xù)訓練的東西到這個字典中。
使用其他模型來預熱當前模型
保存：
torch.save(modelA.state_dict(), PATH)
加載：
modelB = TheModelBClass(*args, **kwargs) modelB.load_state_dict(torch.load(PATH), strict=False)
在遷移學習或者訓練新的復雜模型時，加載部分模型是很常見的。利用經(jīng)過訓練的參數(shù)，即使只有少數(shù)參數(shù)可用，也將有助于預熱訓練過程，并且使模型更快收斂。
在加載部分模型參數(shù)進行預訓練的時候，很可能會碰到鍵不匹配的情況（模型權重都是按鍵值對的形式保存并加載回來的）。因此，無論是缺少鍵還是多出鍵的情況，都可以通過在load_state_dict()函數(shù)中設定strict參數(shù)為False來忽略不匹配的鍵。
如果想將某一層的參數(shù)加載到其他層，但是有些鍵不匹配，那么修改state_dict中參數(shù)的key可以解決這個問題。
跨設備保存與加載模型
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動態(tài)地將tensors的底層存儲重新映射到CPU設備上。
上述代碼只有在模型是在一塊GPU上訓練時才有效，如果模型在多個GPU上訓練，那么在CPU上加載時，會得到類似如下錯誤：
KeyError: ‘unexpected key “module.conv1.weight” in state_dict'
原因是在使用多GPU訓練并保存模型時，模型的參數(shù)名都帶上了module前綴，因此可以在加載模型時，把key中的這個前綴去掉：
# 原始通過DataParallel保存的文件
state_dict = torch.load('myfile.pth.tar')
# 創(chuàng)建一個不包含`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
# 加載參數(shù)
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))
model.to(device)
# 往模型中輸入數(shù)據(jù)的時候不要忘記在任意tensor上調(diào)用input = input.to(device)
在把GPU上訓練的模型加載到GPU上時，只需要使用model.to(torch.devie('cuda'))將初始化的模型轉(zhuǎn)換為CUDA優(yōu)化模型。同時確保在模型所有的輸入上使用.to(torch.device('cuda'))。注意，調(diào)用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")) # 選擇希望使用的GPU
model.to(device)
保存torch.nn.DataParallel模型
保存：
torch.save(model.module.state_dict(), PATH)
總結
到此這篇關于PyTorch模型保存與加載的文章就介紹到這了,更多相關PyTorch模型保存加載內(nèi)容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關文章希望大家以后多多支持腳本之家！
相關文章
python通過設置WordCloud參數(shù)實現(xiàn)定制詞云
這篇文章主要為大家詳細介紹了python如何通過設置WordCloud參數(shù)實現(xiàn)定制詞云,文中的示例代碼講解詳細,感興趣的小伙伴可以跟隨小編一起學習一下2023-11-11
Python調(diào)用shell cmd方法代碼示例解析
這篇文章主要介紹了Python調(diào)用shell cmd方法代碼示例解析,文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值,需要的朋友可以參考下2020-06-06
Python自定義類的數(shù)組排序?qū)崿F(xiàn)代碼
這篇文章主要介紹了Python自定義類的數(shù)組排序?qū)崿F(xiàn)代碼,需要的朋友可以參考下2016-08-08
這篇文章主要介紹了M1芯片安裝python3.9.1的實現(xiàn)，文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值，需要的朋友們下面隨著小編來一起學習學習吧2021-02-02
Python+tkinter實現(xiàn)網(wǎng)站下載工具
這篇文章主要為大家詳細介紹了如何利用Python+tkinter實現(xiàn)網(wǎng)站下載工具，實現(xiàn)所有數(shù)據(jù)一鍵獲取，文中的示例代碼講解詳細，感興趣的可以了解一下2023-03-03
Starship定制shell提示符實現(xiàn)信息自由
這篇文章主要介紹了Starship定制shell提示符的實現(xiàn)，讓你需要的所有信息觸手可及，有需要的朋友可以借鑒參考下，希望能夠有所幫助，祝大家多多進步，早日升職加薪2022-03-03
這篇文章主要介紹了Python3之字符串比較_重寫cmp函數(shù)方式，具有很好的參考價值，希望對大家有所幫助。如有錯誤或未考慮完全的地方，望不吝賜教2022-02-02
這篇文章主要介紹了Pytorch十九種損失函數(shù)的使用詳解，文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值，需要的朋友們下面隨著小編來一起學習學習吧2020-04-04
這篇文章主要介紹了Python日志logging模塊功能與用法,結合實例形式詳細分析了Python日志logging模塊的基本功能、原理、用法及操作注意事項,需要的朋友可以參考下2020-04-04
如何利用python實現(xiàn)圖片轉(zhuǎn)化字符畫
這篇文章主要介紹了如何利用python實現(xiàn)圖片轉(zhuǎn)化字符畫，文章圍繞主題展開詳細的內(nèi)容介紹，具有一定的參考價值，需要的小伙伴可以參考一下2022-06-06