tensorflow python 多个GPU并行处理

神经网络对于电脑硬件非常依赖，尤其是显卡。显卡的优劣决定了训练网络的规模上限和训练速度。而对于一般的个人，甚至实验室来讲，配备多个性能较好的gpu要比配备单个顶级gpu划算很多。如果配备了多个gpu，在训练的时候如果没有指定device的话，你会发现你只有一个gpu的使用率上来了，其他的gpu都是被分配了一部分或全部的显存，但是使用率为0。这就有点浪费了嘛~  那么装备了一个以上的gpu之后，能否把这些gpu们充分调动起来，多个gpu同时进行训练呢？

瞎JB贴个图

---

答案当然是可以的，而且很简单。可惜tensorflow上的官网没写(至少我没看见。。) 最下边的for loop来loop过device会让每一个gpu挨个运行一次指令，这肯定不是我们想要的。

如果要调用多个gpu的话，你先要知道自己有几块gpu，并且它们的device_name是什么，才好手动分配。在上面tensorflow网页上有显示gpu device mapping的代码，运行一下就知道了：

# Creates a graph. 

a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a') 

b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b') 

c = tf.matmul(a, b) 

# Creates a session with log_device_placement set to True. 

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True)) 

# Runs the op. 

print(sess.run(c))

比如我的就是GPU:0 和 GPU:1。如果你打开NVIDIA X Server Setting的话也能看到这两个gpu。

下面重点来了，找到你的gpu的device_name以后怎么办？你需要指定哪一段代码用哪个gpu运行。用tf.device()命令。比如我再训练一个结构的时候要用到gpu:0，那就这么开头：

with tf.device('/device:GPU:0'):  #指定gpu

    with tf.Graph().as_default():  # 配置你的graph

        with tf.Session(config=tf.ConfigProto(log)device_placement=True)) as sess:  # 配置你的session

            xxxxxx  # 训练

同时我要训练另外一个网络，只要把上面的GPU:0 换成 GPU:1就行了(我的第二块gpu名字是GPU:1)。

但是光这样还不够。如果要用不同的进程调用不同的GPU，你还需要设置Cuda里的环境变量CUDA_VISIBLE_DEVICE。如果你用的命令行训练的话，只需要分别打开两个窗口分别进行训练就可以了。比如：

用命令行：

CUDA_VISIBLE_DEVICES=0 python script_one.py

CUDA_VISIBLE_DEVICES=1 python script_two.py

好像也可以两个gpu一起训练一个：

CUDA_VISIBLE_DEVICES=0,1 python script.py  # 这个没测试过 stackoverflow上有大神说可以

如果你用IDE，比如pycharm的话，可以分别打开两个文件页：

用Pycharm：

import os 

os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 第一个文件配置device GPU:0里开始输入这一行

os.environ["CUDA_VISIBLE_DEVICES"] = "1"  # 同理第二个文件配置GPU:1

然后分别运行两个文件，就可以做到同时训练多个模型了。再看你的NVIDIA X Server就会看到两个gpu使用率都起来了~

Ref:

https://www.tensorflow.org/guide/using_gpu

https://github.com/tensorflow/nmt/issues/60

https://stackoverflow.com/questions/44135538/tensorflow-using-2-gpu-at-the-same-time