解決pytorch多GPU訓(xùn)練保存的模型,在單GPU環(huán)境下加載出錯(cuò)問(wèn)題
背景
在公司用多卡訓(xùn)練模型，得到權(quán)值文件后保存，然后回到實(shí)驗(yàn)室，沒(méi)有多卡的環(huán)境，用單卡訓(xùn)練，加載模型時(shí)出錯(cuò)，因?yàn)閱慰C(jī)器上，沒(méi)有使用DataParallel來(lái)加載模型，所以會(huì)出現(xiàn)加載錯(cuò)誤。
原因
DataParallel包裝的模型在保存時(shí)，權(quán)值參數(shù)前面會(huì)帶有module字符，然而自己在單卡環(huán)境下，沒(méi)有用DataParallel包裝的模型權(quán)值參數(shù)不帶module。本質(zhì)上保存的權(quán)值文件是一個(gè)有序字典。
解決方法
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
補(bǔ)充知識(shí)：Pytorch：多GPU訓(xùn)練網(wǎng)絡(luò)與單GPU訓(xùn)練網(wǎng)絡(luò)保存模型的區(qū)別
測(cè)試環(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運(yùn)行
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運(yùn)行
net = LeNet().to(device)
由于多GPU訓(xùn)練使用了 nn.DataParallel(net, device_ids=gpu_ids) 對(duì)網(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) )
由于在測(cè)試模型時(shí)不需要用到多GPU測(cè)試，因此在保存模型時(shí)應(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)境下加載出錯(cuò)問(wèn)題就是小編分享給大家的全部?jī)?nèi)容了，希望能給大家一個(gè)參考，也希望大家多多支持腳本之家。
相關(guān)文章
這篇文章主要介紹了Python解析CDD文件的方法,使用Python 腳本解析CDD文件，統(tǒng)一定義，一鍵生成，十分快捷，本文給大家介紹的非常詳細(xì)，對(duì)大家的學(xué)習(xí)或工作具有一定的參考借鑒價(jià)值，需要的朋友可以參考下2022-05-05
Python3 pyecharts生成Html文件柱狀圖及折線圖代碼實(shí)例
這篇文章主要介紹了Python3 pyecharts生成Html文件柱狀圖及折線圖代碼實(shí)例,文中通過(guò)示例代碼介紹的非常詳細(xì)，對(duì)大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價(jià)值,需要的朋友可以參考下2020-09-09
id() 函數(shù)用于獲取對(duì)象的內(nèi)存地址。很多朋友不清楚python中的id函數(shù)到底是什么？接下來(lái)小編給大家分享本文幫助大家學(xué)習(xí)2017-10-10
Python 中檢查一個(gè)字符是否為數(shù)字的示例代碼
本文詳細(xì)介紹了在 Python 中檢查一個(gè)字符是否為數(shù)字的幾種常用方法。我們介紹了使用 isdigit() 方法、isnumeric() 方法和正則表達(dá)式等方法，并提供了示例代碼幫助你理解和應(yīng)用這些方法,需要的朋友可以參考下2023-06-06
這篇文章主要介紹了Python threading多線程編程實(shí)例,本文講解了使用函數(shù)和線程類實(shí)現(xiàn)多線程編程的例子,需要的朋友可以參考下2014-09-09
Python 模擬員工信息數(shù)據(jù)庫(kù)操作的實(shí)例
下面小編就為大家?guī)?lái)一篇Python 模擬員工信息數(shù)據(jù)庫(kù)操作的實(shí)例。小編覺(jué)得挺不錯(cuò)的，現(xiàn)在就分享給大家，也給大家做個(gè)參考。一起跟隨小編過(guò)來(lái)看看吧2017-10-10
Python實(shí)現(xiàn)基于POS算法的區(qū)塊鏈
這篇文章主要介紹了Python實(shí)現(xiàn)基于POS算法的區(qū)塊鏈,今天我們就來(lái)認(rèn)識(shí)POS(proof of stake)算法。需要的朋友可以參考下2018-08-08
Python中__str__()方法的實(shí)用技巧分享
在Python編程中,__str__()是一個(gè)特殊方法,它允許自定義對(duì)象的字符串表示形式,本文將深入探討__str__()的相關(guān)實(shí)用技巧,希望對(duì)大家有所幫助2023-11-11
這篇文章主要介紹了Python 寫入Excel格式和顏色的實(shí)現(xiàn)方式，具有很好的參考價(jià)值，希望對(duì)大家有所幫助。一起跟隨小編過(guò)來(lái)看看吧2021-03-03
Python中dumps與dump及l(fā)oads與load的區(qū)別
這篇文章主要介紹了Python中dumps與dump、loads與load的區(qū)別，json模塊提供了一種很簡(jiǎn)單的方式來(lái)編碼和解碼JSON數(shù)據(jù)。其中兩個(gè)主要的函數(shù)是json.dumps()和json.loads(),需要的朋友可以參考下2022-04-04