這幾天在一機多卡的環(huán)境下，用pytorch訓練模型，遇到很多問題。現(xiàn)總結一個實用的做實驗方式：
多GPU下訓練，創(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()
官方建議的模型保存方式，只保存參數：
torch.save(model.module.state_dict(), "model.pkl")
其實，這樣很麻煩，我建議直接保存模型（參數+圖）：
torch.save(model, "model.pkl")
這樣做很實用，特別是我們需要反復建模和調試的時候。這種情況下模型的加載很方便，因為模型的圖已經和參數保存在一起，我們不需要根據不同的模型設置相應的超參，更換對應的網絡結構，如下：
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下模型。如果服務器環(huán)境變化不大，或者和訓練時候是同一個GPU環(huán)境，就不會出現(xiàn)問題。
如果系統(tǒng)環(huán)境發(fā)生了變化，或者，我們只想加載模型參數，亦或是遇到下面的問題：
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
這時候我們可以用下面的方式載入模型，先建立模型，然后加載參數。
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)
# 獲得模型參數
model_dict = torch.load(args.pretrained_model_path).module.state_dict()
# 載入參數
model.module.load_state_dict(model_dict)
print('success!')
到此這篇關于PyTorch 多GPU下模型的保存與加載(踩坑筆記)的文章就介紹到這了,更多相關PyTorch 多GPU下模型的保存與加載內容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關文章希望大家以后多多支持腳本之家！
您可能感興趣的文章:- 基于pytorch的保存和加載模型參數的方法
- Pytorch之保存讀取模型實例
- pytorch模型的保存和加載、checkpoint操作
