PyTorch 多GPU下模型的保存与加载(踩坑笔记)
2021-03-08 19:34:59
这几天在一机多卡的环境下,用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下模型的保存与加载内容请搜索IT技术网以前的文章或继续浏览下面的相关文章希望大家以后多多支持IT技术网!
您可能感兴趣的文章:
- 解决pytorch多GPU训练保存的模型,在单GPU环境下加载出错问题
- 将Pytorch模型从CPU转换成GPU的实现方法
相关阅读
- Pytorch1.5.1版本安装的方法步骤
- pytorch中index_select()的用法详解
- 解决pytorch下出现multi-target not supported at的一种可能原因
- pytorch 中forward 的用法与解释说明
- pytorch __init__、forward与__call__的用法小结
- Pytorch - TORCH.NN.INIT 参数初始化的操作
- Pytorch如何切换 cpu和gpu的使用详解
- 基于PyTorch中view的用法说明
- 对pytorch中x = x.view(x.size(0), -1) 的理解说明
- Pytorch 中的optimizer使用说明
- 解决pytorch 的state_dict()拷贝问题
- 对Pytorch 中的contiguous理解说明
- Pytorch实现WGAN用于动漫头像生成
- Pytorch之扩充tensor的操作
- pytorch 计算Parameter和FLOP的操作
- pytorch 把图片数据转化成tensor的操作
- pytorch 数据加载性能对比分析
- pytorch从csv加载自定义数据模板的操作