使用PyTorch怎么多GPU中對模型進行保存

這篇文章將為大家詳細講解有關使用PyTorch怎么多GPU中對模型進行保存，文章內容質量較高，因此小編分享給大家做個參考，希望大家閱讀完這篇文章后對相關知識有一定的了解。

多GPU下訓練，創建模型代碼通常如下：

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下模型。如果服務器環境變化不大，或者和訓練時候是同一個GPU環境，就不會出現問題。

如果系統環境發生了變化，或者，我們只想加載模型參數，亦或是遇到下面的問題：

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中對模型進行保存就分享到這里了，希望以上內容可以對大家有一定的幫助，可以學到更多知識。如果覺得文章不錯，可以把它分享出去讓更多的人看到。