服务器之家

服务器之家 > 正文

PyTorch 多GPU下模型的保存与加载(踩坑笔记)

时间:2021-09-17 00:02     来源/作者:叶罅

这几天在一机多卡的环境下,用pytorch训练模型,遇到很多问题。现总结一个实用的做实验方式:

多GPU下训练,创建模型代码通常如下:

?
1
2
3
4
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()

官方建议的模型保存方式,只保存参数:

?
1
torch.save(model.module.state_dict(), "model.pkl")

其实,这样很麻烦,我建议直接保存模型(参数+图):

?
1
torch.save(model, "model.pkl")

这样做很实用,特别是我们需要反复建模和调试的时候。这种情况下模型的加载很方便,因为模型的图已经和参数保存在一起,我们不需要根据不同的模型设置相应的超参,更换对应的网络结构,如下:

?
1
2
3
4
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

这时候我们可以用下面的方式载入模型,先建立模型,然后加载参数。

?
1
2
3
4
5
6
7
8
9
10
11
12
13
14
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下模型的保存与加载内容请搜索服务器之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持服务器之家!

原文链接:https://www.cnblogs.com/blog4ljy/p/11711173.html

标签:

相关文章

热门资讯

yue是什么意思 网络流行语yue了是什么梗
yue是什么意思 网络流行语yue了是什么梗 2020-10-11
2020微信伤感网名听哭了 让对方看到心疼的伤感网名大全
2020微信伤感网名听哭了 让对方看到心疼的伤感网名大全 2019-12-26
背刺什么意思 网络词语背刺是什么梗
背刺什么意思 网络词语背刺是什么梗 2020-05-22
苹果12mini价格表官网报价 iPhone12mini全版本价格汇总
苹果12mini价格表官网报价 iPhone12mini全版本价格汇总 2020-11-13
2021年耽改剧名单 2021要播出的59部耽改剧列表
2021年耽改剧名单 2021要播出的59部耽改剧列表 2021-03-05
返回顶部