PyTorch 多GPU下模型的保存与加载(踩坑笔记)
更新时间：2021年03月08日 14:06:36 作者：叶罅
这篇文章主要介绍了PyTorch 多GPU下模型的保存与加载(踩坑笔记)，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧
这几天在一机多卡的环境下，用pytorch训练模型，遇到很多问题。现总结一个实用的做实验方式：
多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下模型的保存与加载(踩坑笔记)的文章就介绍到这了,更多相关PyTorch 多GPU下模型的保存与加载内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！
相关文章
这篇文章主要为大家介绍了Flask学习笔记之日志操作配置实例讲解,有需要的朋友可以借鉴参考下,希望能够有所帮助,祝大家多多进步,早日升职加薪2023-11-11
抽了点时间体验了一把python分布式进程，有点像分布式计算的意思，这篇文章主要给大家介绍了关于python中如何使用分步式进程计算的相关资料，需要的朋友可以参考下2019-03-03
这篇文章主要介绍了python实现数组平移K位问题，具有很好的参考价值，希望对大家有所帮助。如有错误或未考虑完全的地方，望不吝赐教2023-02-02
今天小编就为大家分享一篇python 读取视频,处理后,实时计算帧数fps的方法，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧2018-07-07
这篇文章主要介绍了python定义函数如何判断年份是否为闰年，具有很好的参考价值，希望对大家有所帮助。如有错误或未考虑完全的地方，望不吝赐教2022-07-07
本文主要介绍了Pysvn 使用指南，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧2023-02-02
这篇文章主要和大家分享两个有意思的Python小工具,可以优雅地实现进度条和系统通知,文中的示例代码简洁易懂,有需要的小伙伴快也跟随小编一起学习一下2023-11-11
这篇文章主要介绍了使用pdb模块调试Python程序实例,本文着重讲解了pdb.run()函数、pdb.runeval()函数、pdb.runcall()函数、pdb.set_trace()函数的使用以及pdb调试命令等内容,需要的朋友可以参考下2015-06-06
这篇文章主要为大家详细介绍了如何使用Python实现获取某一进程的CPU利用率,文中的示例代码讲解详细,感兴趣的小伙伴可以跟随小编一起学习一下2024-02-02
这篇文章主要介绍了 Python 时间操作datetime,datetime 模块提供处理时间和日期的多种类，简单方便,下面文章将详细介绍其内容,需要的朋友可以参考一下2021-11-11
最新评论