解決pytorch多GPU訓(xùn)練保存的模型,在單GPU環(huán)境下加載出錯問題
背景
在公司用多卡訓(xùn)練模型，得到權(quán)值文件后保存，然后回到實驗室，沒有多卡的環(huán)境，用單卡訓(xùn)練，加載模型時出錯，因為單卡機器上，沒有使用DataParallel來加載模型，所以會出現(xiàn)加載錯誤。
原因
DataParallel包裝的模型在保存時，權(quán)值參數(shù)前面會帶有module字符，然而自己在單卡環(huán)境下，沒有用DataParallel包裝的模型權(quán)值參數(shù)不帶module。本質(zhì)上保存的權(quán)值文件是一個有序字典。
解決方法
1.在單卡環(huán)境下，用DataParallel包裝模型。
2.自己重寫Load函數(shù)，靈活。
from collections import OrderedDict
def myOwnLoad(model, check):
modelState = model.state_dict()
tempState = OrderedDict()
for i in range(len(check.keys())-2):
print modelState.keys()[i], check.keys()[i]
tempState[modelState.keys()[i]] = check[check.keys()[i]]
temp = [[0.02]*1024 for i in range(200)] # mean=0, std=0.02
tempState['myFc.weight'] = torch.normal(mean=0, std=torch.FloatTensor(temp)).cuda()
tempState['myFc.bias'] = torch.normal(mean=0, std=torch.FloatTensor([0]*200)).cuda()
model.load_state_dict(tempState)
return model
補充知識：Pytorch：多GPU訓(xùn)練網(wǎng)絡(luò)與單GPU訓(xùn)練網(wǎng)絡(luò)保存模型的區(qū)別
測試環(huán)境：Python3.6 + Pytorch0.4
在pytorch中，使用多GPU訓(xùn)練網(wǎng)絡(luò)需要用到 【nn.DataParallel】：
gpu_ids = [0, 1, 2, 3]
device = t.device("cuda:0" if t.cuda.is_available() else "cpu") # 只能單GPU運行
net = LeNet()
if len(gpu_ids) > 1:
net = nn.DataParallel(net, device_ids=gpu_ids)
net = net.to(device)
而使用單GPU訓(xùn)練網(wǎng)絡(luò)：
device = t.device("cuda:0" if t.cuda.is_available() else "cpu") # 只能單GPU運行
net = LeNet().to(device)
由于多GPU訓(xùn)練使用了 nn.DataParallel(net, device_ids=gpu_ids) 對網(wǎng)絡(luò)進(jìn)行封裝，因此在原始網(wǎng)絡(luò)結(jié)構(gòu)中添加了一層module。網(wǎng)絡(luò)結(jié)構(gòu)如下：
DataParallel( (module): LeNet( (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1)) (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1)) (fc1): Linear(in_features=400, out_features=120, bias=True) (fc2): Linear(in_features=120, out_features=84, bias=True) (fc3): Linear(in_features=84, out_features=10, bias=True) ) )
而不使用多GPU訓(xùn)練的網(wǎng)絡(luò)結(jié)構(gòu)如下：
LeNet( (conv1): Conv2d(3, 6, kernel_size=(5, 5), stride=(1, 1)) (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1)) (fc1): Linear(in_features=400, out_features=120, bias=True) (fc2): Linear(in_features=120, out_features=84, bias=True) (fc3): Linear(in_features=84, out_features=10, bias=True) )
由于在測試模型時不需要用到多GPU測試，因此在保存模型時應(yīng)該把module層去掉。如下：
if len(gpu_ids) > 1: t.save(net.module.state_dict(), "model.pth") else: t.save(net.state_dict(), "model.pth")
以上這篇解決pytorch多GPU訓(xùn)練保存的模型,在單GPU環(huán)境下加載出錯問題就是小編分享給大家的全部內(nèi)容了，希望能給大家一個參考，也希望大家多多支持腳本之家。
相關(guān)文章
這篇文章主要介紹了Python解析CDD文件的方法,使用Python 腳本解析CDD文件，統(tǒng)一定義，一鍵生成，十分快捷，本文給大家介紹的非常詳細(xì)，對大家的學(xué)習(xí)或工作具有一定的參考借鑒價值，需要的朋友可以參考下2022-05-05
Python3 pyecharts生成Html文件柱狀圖及折線圖代碼實例
這篇文章主要介紹了Python3 pyecharts生成Html文件柱狀圖及折線圖代碼實例,文中通過示例代碼介紹的非常詳細(xì)，對大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價值,需要的朋友可以參考下2020-09-09
id() 函數(shù)用于獲取對象的內(nèi)存地址。很多朋友不清楚python中的id函數(shù)到底是什么？接下來小編給大家分享本文幫助大家學(xué)習(xí)2017-10-10
本文詳細(xì)介紹了在 Python 中檢查一個字符是否為數(shù)字的幾種常用方法。我們介紹了使用 isdigit() 方法、isnumeric() 方法和正則表達(dá)式等方法，并提供了示例代碼幫助你理解和應(yīng)用這些方法,需要的朋友可以參考下2023-06-06
這篇文章主要介紹了Python threading多線程編程實例,本文講解了使用函數(shù)和線程類實現(xiàn)多線程編程的例子,需要的朋友可以參考下2014-09-09
Python 模擬員工信息數(shù)據(jù)庫操作的實例
下面小編就為大家?guī)硪黄狿ython 模擬員工信息數(shù)據(jù)庫操作的實例。小編覺得挺不錯的，現(xiàn)在就分享給大家，也給大家做個參考。一起跟隨小編過來看看吧2017-10-10
這篇文章主要介紹了Python實現(xiàn)基于POS算法的區(qū)塊鏈,今天我們就來認(rèn)識POS(proof of stake)算法。需要的朋友可以參考下2018-08-08
在Python編程中,__str__()是一個特殊方法,它允許自定義對象的字符串表示形式,本文將深入探討__str__()的相關(guān)實用技巧,希望對大家有所幫助2023-11-11
這篇文章主要介紹了Python 寫入Excel格式和顏色的實現(xiàn)方式，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧2021-03-03
Python中dumps與dump及l(fā)oads與load的區(qū)別
這篇文章主要介紹了Python中dumps與dump、loads與load的區(qū)別，json模塊提供了一種很簡單的方式來編碼和解碼JSON數(shù)據(jù)。其中兩個主要的函數(shù)是json.dumps()和json.loads(),需要的朋友可以參考下2022-04-04