PyTorch 多GPU下模型的保存與加載(踩坑筆記)

這幾天在一機(jī)多卡的環(huán)境下，用pytorch訓(xùn)練模型，遇到很多問(wèn)題。現(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)文章希望大家以后多多支持腳本之家！