这篇教程PyTorch 多GPU下模型的保存与加载(踩坑笔记)写得很实用，希望能帮到您。
这几天在一机多卡的环境下，用pytorch训练模型，遇到很多问题。现总结一个实用的做实验方式： 多GPU下训练，创建模型代码通常如下： os.environ['CUDA_VISIBLE_DEVICES'] = args.cudamodel = 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下模型的保存与加载内容请搜索51zixue.net以前的文章或继续浏览下面的相关文章希望大家以后多多支持51zixue.net！
解决python3 整数数组转bytes的效率问题
python程序调用远程服务的步骤详解 |
