标签：count 模型 并行 ids Pytorch device GPU model

如果要用多GPU训练模型，那就要用到pytorch中的nn.DataParallel()函数对模型进行封装。

1 # code_for_use_cuda
2 device_count = torch.cuda.device_count()#获得本机的显卡的个数
3 print("cuda.device_count",device_count)
4 device_ids = list(range(device_count))
5 model = nn.DataParallel(model,device_ids=device_ids)#对模型进行封装，如果不填写device_ids那就默认的使用所有的GPU
6 criterion = nn.DataParallel(criterion,device_ids=device_ids)#

需要注意的几点：

（1）但是使用这种方法会造成多GPU的运行的负载不匹配的状况，在多GPU中选定一块主GPU进行Loss值的计算，这块GPU往往使用的memory更多。在有些模型中loss的计算是很耗费资源的，

这时候就建议将loss的计算放在model类的forward()方法中进行计算。

（2）在多GPU训练的情况下，如果模型的forward()返回的是单个数值，例如loss.item()。多GPU情况下最终会得到一个列表，是每个单独的GPU运算的结果，如果求loss需要再对他们进行平均。

（3）优化器optimization的处理不需要并行的计算。

（4）这一点笔者认为是最重要的一点，即并行模型的保存和加载的问题。使用了DataParallel()之后的model.state_dict().keys()相较于原来的，键名在前面多了“module.”。这样的模型文件

不能被不使用并行处理的模型所加载。解决方法：在保存模型的时候，使用torch.save(model.module.state_dict(),'filename.pt')，这样保存的模型的键名中就不含有module了。

（5）如果模型的训练和使用并不在同一台机器上，那保存和读取的时候最好都在CPU上进行，否则重新加载模型可能会发生错误。

标签：count,模型,并行,ids,Pytorch,device,GPU,model
来源： https://www.cnblogs.com/justkeen/p/16532515.html

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。