分类：多GPU下模型保存

2021
03-09

PyTorch 多GPU下模型的保存与加载(踩坑笔记)

这几天在一机多卡的环境下，用pytorch训练模型，遇到很多问题。现总结一个实用的做实验方式：多GPU下训练，创建模型代码通常如下：os.environ['CUDA_VISIBLE_DEVICES']=args.cudamodel=MyModel(args)iftorch.cuda.is_available()andargs.use_gpu:model=torch.nn.DataParallel(model).cuda()官方建议的模型保存方式，只保存参数：torch.save(model.module.state_dict(),"model.pkl")其实，这样很麻烦，我建议直接保存模... 继续阅读 >