Python tensorflow多GPU并行使用

我想并行使用8GPU，而不是按顺序使用

例如，当我执行此代码时

import tensorflow as tf

with tf.device('/gpu:0'):
    for i in range(10):
        print(i)

with tf.device('/gpu:1'):
    for i in range(10, 20):
        print(i)

我尝试了cmd命令“CUDA\u VISIBLE\u DEVICE='0,1'，但结果相同

我想看到结果“01011112312…等等”

但实际结果是依次为“01 2 3 4 5…..10 11 12 13….”

---

我怎样才能得到想要的结果？

**我看到一个问题的编辑，所以将此添加到我的答案中**

您需要将您的操作传递给Tensorflow会话，否则，代码将被解释为顺序的（就像许多编程语言一样），然后操作将按顺序完成

对于之前对该问题的理解，以下讨论了如何创建具有多个GPU的神经网络训练：

坏消息是，没有神奇的功能可以帮你做到这一点

好消息是，有一些既定的方法

第一个是一些CUDA和其他GPU开发人员所熟悉的，可以将模型复制到多个GPU，通过CPU进行同步。一种方法是分批分割数据集，或者在本例中称为塔，然后为每个GPU提供一个塔。如果这是MNIST数据集，并且您有两个GPU，那么您可以显式地使用CPU作为设备来初始化此数据。现在，随着数据集变小，相对批处理大小可能会变大。一旦你完成了一个历元，你可以分享梯度和平均值来训练两个网络。当然，这很容易扩展到8 GPU的情况

在CPU上执行任务分配和收集结果的步骤如下所示：

# Creates a graph.
c = []
for d in ['/gpu:2', '/gpu:3']:
  with tf.device(d):
    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3])
    b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2])
    c.append(tf.matmul(a, b))
with tf.device('/cpu:0'):
  sum = tf.add_n(c)
# 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(sum))

但是，在多个设备之间传输数据会使您无法准确地获得您的gpu数倍加速度。因此，您需要优化每个GPU的工作负载，以最大限度地提高性能，并尽可能避免设备间通信

第二种方法是将你的神经网络分割成你拥有的设备，训练并合并它们

在多个GPU上显式运行模型需要您以这种方式设置算法。看看这些：