如何使用多GPU训练

1.Pytorch

首先设置 CUDA_VISIBLE_DEVICES，使得训练代码可以检测的显卡数量，具体设置如下，

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3" #此时显示4块显卡

之后就是通过 DataParallel 设置并行训练

device_ids = [0,1,2,3]
model = model.cuda(device_ids[0])
model = nn.DataParallel(model, device_ids=device_ids)

• 使用 model.gpu()，进入 gpu 模式
• model = nn.DataParallel(model) 启动多 GPU 训练
• inputs = Variable(inputs.cuda()),
  labels = Variable(labels.cuda()) 表示张量加载到GPU之中

ps: 笔者用了两张 Tesla K80 测试，发现用一个GPU，batch-size 可以为6，用三个GPU后batch-size只能支持到 4，如果四个全开就剩 3 了。并且发现第一个GPU卡占用率总是比之后的大，查阅相关资料，可能是第一个 GPU 还充当着正常的前传反传以及参数更新的作用，而其它的只负责前传和反传。所以随着开启的 GPU 数量增多，第一个的显存负担也会增加，RuntimeError: CUDA error : out of memory，多是因为第一个 GPU 达到了显存极限。如果这个问题解决，那么便可以提高效率，我觉得应该是有办法的。寻找中。

2. Tensorflow