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下模型的保存与加载内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！
相关文章
这篇文章主要介绍了python实现扫雷游戏的示例，帮助大家更好的理解和学习python，感兴趣的朋友可以了解下2020-10-10
这篇文章主要介绍了Python FastAPI 多参数传递,FastAPI通过模板来匹配URL中的参数列表，大概分为三类方式传递参数，每种方式结合示例代码给大家介绍的非常详细，需要的朋友可以参考下2022-12-12
本文主要介绍了pandas创建DataFrame对象失败的解决方法，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧2023-01-01
下面就使用VLAD表示图像，实现一个小型的图像数据库的检索程序。下面实现需要的功能模块，分步骤给大家介绍的非常详细，对OpenCV图像数据库检索功能感兴趣的朋友跟随小编一起看看吧2021-12-12
这篇文章主要介绍了GraphQL在Django中的使用教程,本文结合示例代码给大家介绍的非常详细，对大家的学习或工作具有一定的参考借鉴价值，需要的朋友可以参考下2022-12-12
这篇文章介绍了Python的序列化模块JSON与Pickle，文中通过示例代码介绍的非常详细。对大家的学习或工作具有一定的参考借鉴价值，需要的朋友可以参考下2022-06-06
解决windows下命令行执行python3失效,会打开应用商店问题
这篇文章主要介绍了解决windows下命令行执行python3失效,会打开应用商店问题,具有很好的参考价值,希望对大家有所帮助,如有错误或未考虑完全的地方,望不吝赐教2024-02-02
Python enumerate函数遍历数据对象组合过程解析
这篇文章主要介绍了Python enumerate函数遍历数据对象组合过程解析,文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下2019-12-12
这篇文章主要介绍了Python开发的实用计算器,结合完整实例形式分析了Python实现计算器四则运算、开方、取余等相关操作技巧,需要的朋友可以参考下2017-05-05
opencv+python实现鼠标点击图像,输出该点的RGB和HSV值
这篇文章主要介绍了opencv+python实现鼠标点击图像,输出该点的RGB和HSV值，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧2020-06-06
最新评论