pytorch多GPU并行运算的实现

Pytorch多GPU运行

设置可用GPU环境变量。例如，使用0号和1号GPU'

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1'

设置模型参数放置到多个GPU上。在pytorch1.0之后的版本中，多GPU运行变得十分方便，先将模型的参数设置并行

if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)

将模型参数设置使用GPU运行

if torch.cuda.is_available():
model.cuda()

踩坑记录

在训练中，需要使用验证集/测试集对目前的准确率进行测试，验证集/测试集的加载也会占用部分显存，所以在训练开始时，不要将所有显存都几乎占满，稍微留一些显存给训练过程中的测试环节

pytorch并行后，假设batchsize设置为64，表示每张并行使用的GPU都使用batchsize=64来计算（单张卡使用时，使用batchsize=64比较合适时，多张卡并行时，batchsize仍为64比较合适，而不是64*并行卡数）。

参考

https://www.zhihu.com/question/67726969

以上就是本文的全部内容，希望对大家的学习有所帮助

您可能感兴趣的文章:

• 用Pytorch训练CNN(数据集MNIST,使用GPU的方法)
• 解决pytorch GPU 计算过程中出现内存耗尽的问题
• 将Pytorch模型从CPU转换成GPU的实现方法
• pytorch 使用单个GPU与多个GPU进行训练与测试的方法
• 在pytorch中为Module和Tensor指定GPU的例子
• pytorch使用指定GPU训练的实例
• 关于pytorch多GPU训练实例与性能对比分析