这篇文章将为大家详细讲解有关使用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中对模型进行保存就分享到这里了，希望以上内容可以对大家有一定的帮助，可以学到更多知识。如果觉得文章不错，可以把它分享出去让更多的人看到。
亿速云「云服务器」，即开即用、新一代英特尔至强铂金CPU、三副本存储NVMe SSD云盘，价格低至29元/月。点击查看>>
免责声明：本站发布的内容（图片、视频和文字）以原创、转载和分享为主，文章观点不代表本网站立场，如果涉及侵权请联系站长邮箱：is@yisu.com进行举报，并提供相关证据，一经查实，将立刻删除涉嫌侵权内容。
