PyTorch模型保存與加載實例詳解
- torch.save：保存序列化的對象到磁盤，使用了Python的pickle進行序列化，模型、張量、所有對象的字典。
- torch.load：使用了pickle的unpacking將pickled的對象反序列化到內存中。
- 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文件格式保存。
一定要記住在評估模式的時候調用model.eval()來固定dropout和批次歸一化。否則會產生不一致的推理結果。
保存加載用于推理的常規(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()加載保存的字典，然后可以直接查詢字典中的值來獲取保存的組件。
同樣，評估模型的時候一定不要忘了調用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)
在遷移學習或者訓練新的復雜模型時，加載部分模型是很常見的。利用經過訓練的參數(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上調用input = input.to(device)
在把GPU上訓練的模型加載到GPU上時，只需要使用model.to(torch.devie('cuda'))將初始化的模型轉換為CUDA優(yōu)化模型。同時確保在模型所有的輸入上使用.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")) # 選擇希望使用的GPU
model.to(device)
保存torch.nn.DataParallel模型
保存：
torch.save(model.module.state_dict(), PATH)
總結
到此這篇關于PyTorch模型保存與加載的文章就介紹到這了,更多相關PyTorch模型保存加載內容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關文章希望大家以后多多支持腳本之家！
相關文章
不論什么語言我們都需要注意性能優(yōu)化問題，提高執(zhí)行效率，這里就為大家分享下Python的性能優(yōu)化技巧,需要的朋友可以參考下2014-10-10
本文給大家分享的是Python隨機數(shù)random模塊的幾個常用的方法，非常的簡單，小伙伴們喜歡的話，后續(xù)繼續(xù)深入探討2016-09-09
pip?install?jupyterlab失敗的原因問題及探索
在學習Yolo模型時,嘗試安裝JupyterLab但遇到錯誤,錯誤提示缺少Rust和Cargo編譯環(huán)境,因為pywinpty包需要它們來編譯,由于在conda環(huán)境下操作,Rust和Cargo已經安裝,問題是pywinpty包丟失,安裝pywinpty包后,再次執(zhí)行pip?install?jupyterlab即可正常下載2025-02-02
這篇文章主要給大家介紹了關于如何將.ipynb文件轉換成.py文件的詳細步驟,文中通過圖文介紹的非常詳細,大家基本一看就會,需要的朋友可以參考下2023-07-07
pycharm創(chuàng)建一個python包方法圖解
在本篇文章中小編給大家分享了關于pycharm怎么創(chuàng)建一個python包的相關知識點，需要的朋友們學習下。2019-04-04
Anaconda 查看、創(chuàng)建、管理和使用python環(huán)境的方法
這篇文章主要介紹了Anaconda 查看、創(chuàng)建、管理和使用python環(huán)境的方法，文中通過示例代碼介紹的非常詳細，對大家的學習或者工作具有一定的參考學習價值，需要的朋友們下面隨著小編來一起學習學習吧2019-12-12
本篇給大家分享一下通過Python的PyPDF2庫提取pdf中的圖片方法，文中有詳細的代碼示例和流程步驟，感興趣的同學可以閱讀一下2023-05-05
Python實現(xiàn)微信自動回復信息的功能(根據(jù)不同信息回復對應的信息)
這篇文章主要介紹了Python實現(xiàn)微信自動回復信息的功能(根據(jù)不同信息回復對應的信息),我們使用的第三方包是UIAutomation,結合示例代碼給大家介紹的非常詳細,需要的朋友可以參考下2023-09-09
基于Python實現(xiàn)中秋佳節(jié)月餅搶購腳本
這篇文章主要介紹了Python版中秋佳節(jié)月餅搶購腳本,今天要用的是一個測試工具的庫Selenium，今天我們就是用它去實現(xiàn)自動化搶購月餅，其實就是用這個工具"模擬"人為操作瀏覽器相應的操作，比如登陸，勾選購物車商品，下單購買等等操作，需要的朋友可以參考下2022-09-09
這篇文章主要介紹了編寫Python小程序來統(tǒng)計測試腳本的關鍵字的方法,文中的實例不僅可以統(tǒng)計關鍵字數(shù)量,還可以按主關鍵字來歸類,需要的朋友可以參考下2016-03-03