PyTorch 多GPU下模型的保存與加載(踩坑筆記)
這幾天在一機(jī)多卡的環(huán)境下，用pytorch訓(xùn)練模型，遇到很多問(wèn)題?，F(xiàn)總結(jié)一個(gè)實(shí)用的做實(shí)驗(yàn)方式：
多GPU下訓(xùn)練，創(chuàng)建模型代碼通常如下：
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda model = MyModel(args) if torch.cuda.is_available() and args.use_gpu: model = torch.nn.DataParallel(model).cuda()
官方建議的模型保存方式，只保存參數(shù)：
torch.save(model.module.state_dict(), "model.pkl")
其實(shí)，這樣很麻煩，我建議直接保存模型（參數(shù)+圖）：
torch.save(model, "model.pkl")
這樣做很實(shí)用，特別是我們需要反復(fù)建模和調(diào)試的時(shí)候。這種情況下模型的加載很方便，因?yàn)槟Ｐ偷膱D已經(jīng)和參數(shù)保存在一起，我們不需要根據(jù)不同的模型設(shè)置相應(yīng)的超參，更換對(duì)應(yīng)的網(wǎng)絡(luò)結(jié)構(gòu)，如下：
if not (args.pretrained_model_path is None):
print('load model from %s ...' % args.pretrained_model_path)
model = torch.load(args.pretrained_model_path)
print('success!')
但是需要注意，這種方式加載的是多GPU下模型。如果服務(wù)器環(huán)境變化不大，或者和訓(xùn)練時(shí)候是同一個(gè)GPU環(huán)境，就不會(huì)出現(xiàn)問(wèn)題。
如果系統(tǒng)環(huán)境發(fā)生了變化，或者，我們只想加載模型參數(shù)，亦或是遇到下面的問(wèn)題：
AttributeError: 'model' object has no attribute 'copy'
或者
AttributeError: 'DataParallel' object has no attribute 'copy'
或者
RuntimeError: module must have its parameters and buffers on device cuda:0 (device_ids[0]) but found
這時(shí)候我們可以用下面的方式載入模型，先建立模型，然后加載參數(shù)。
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
# 建立模型
model = MyModel(args)
if torch.cuda.is_available() and args.use_gpu:
model = torch.nn.DataParallel(model).cuda()
if not (args.pretrained_model_path is None):
print('load model from %s ...' % args.pretrained_model_path)
# 獲得模型參數(shù)
model_dict = torch.load(args.pretrained_model_path).module.state_dict()
# 載入?yún)?shù)
model.module.load_state_dict(model_dict)
print('success!')
到此這篇關(guān)于PyTorch 多GPU下模型的保存與加載(踩坑筆記)的文章就介紹到這了,更多相關(guān)PyTorch 多GPU下模型的保存與加載內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！
相關(guān)文章
這篇文章主要介紹了python實(shí)現(xiàn)掃雷游戲的示例，幫助大家更好的理解和學(xué)習(xí)python，感興趣的朋友可以了解下2020-10-10
Python FastAPI 多參數(shù)傳遞的示例詳解
這篇文章主要介紹了Python FastAPI 多參數(shù)傳遞,FastAPI通過(guò)模板來(lái)匹配URL中的參數(shù)列表，大概分為三類方式傳遞參數(shù)，每種方式結(jié)合示例代碼給大家介紹的非常詳細(xì)，需要的朋友可以參考下2022-12-12
pandas創(chuàng)建DataFrame對(duì)象失敗的解決方法
本文主要介紹了pandas創(chuàng)建DataFrame對(duì)象失敗的解決方法，文中通過(guò)示例代碼介紹的非常詳細(xì)，對(duì)大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價(jià)值，需要的朋友們下面隨著小編來(lái)一起學(xué)習(xí)學(xué)習(xí)吧2023-01-01
基于OpenCV實(shí)現(xiàn)小型的圖像數(shù)據(jù)庫(kù)檢索功能
下面就使用VLAD表示圖像，實(shí)現(xiàn)一個(gè)小型的圖像數(shù)據(jù)庫(kù)的檢索程序。下面實(shí)現(xiàn)需要的功能模塊，分步驟給大家介紹的非常詳細(xì)，對(duì)OpenCV圖像數(shù)據(jù)庫(kù)檢索功能感興趣的朋友跟隨小編一起看看吧2021-12-12
這篇文章主要介紹了GraphQL在Django中的使用教程,本文結(jié)合示例代碼給大家介紹的非常詳細(xì)，對(duì)大家的學(xué)習(xí)或工作具有一定的參考借鑒價(jià)值，需要的朋友可以參考下2022-12-12
這篇文章介紹了Python的序列化模塊JSON與Pickle，文中通過(guò)示例代碼介紹的非常詳細(xì)。對(duì)大家的學(xué)習(xí)或工作具有一定的參考借鑒價(jià)值，需要的朋友可以參考下2022-06-06
解決windows下命令行執(zhí)行python3失效,會(huì)打開應(yīng)用商店問(wèn)題
這篇文章主要介紹了解決windows下命令行執(zhí)行python3失效,會(huì)打開應(yīng)用商店問(wèn)題,具有很好的參考價(jià)值,希望對(duì)大家有所幫助,如有錯(cuò)誤或未考慮完全的地方,望不吝賜教2024-02-02
Python enumerate函數(shù)遍歷數(shù)據(jù)對(duì)象組合過(guò)程解析
這篇文章主要介紹了Python enumerate函數(shù)遍歷數(shù)據(jù)對(duì)象組合過(guò)程解析,文中通過(guò)示例代碼介紹的非常詳細(xì)，對(duì)大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價(jià)值,需要的朋友可以參考下2019-12-12
Python開發(fā)的實(shí)用計(jì)算器完整實(shí)例
這篇文章主要介紹了Python開發(fā)的實(shí)用計(jì)算器,結(jié)合完整實(shí)例形式分析了Python實(shí)現(xiàn)計(jì)算器四則運(yùn)算、開方、取余等相關(guān)操作技巧,需要的朋友可以參考下2017-05-05
opencv+python實(shí)現(xiàn)鼠標(biāo)點(diǎn)擊圖像,輸出該點(diǎn)的RGB和HSV值
這篇文章主要介紹了opencv+python實(shí)現(xiàn)鼠標(biāo)點(diǎn)擊圖像,輸出該點(diǎn)的RGB和HSV值，具有很好的參考價(jià)值，希望對(duì)大家有所幫助。一起跟隨小編過(guò)來(lái)看看吧2020-06-06