tensorflow 多 GPU 并行使用

Question

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

例如，当我执行这段代码时，

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_VISIBLE_DEVICE='0,1' 但结果是一样的。

我想看到结果“0 10 1 11 2 3 12 ....等”

但实际结果依次为“0 1 2 3 4 5 ..... 10 11 12 13..”

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

Answer 1

** 我看到问题的编辑，因此将其添加到我的答案中**

您需要将操作传递给 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))

但是，在许多设备之间传输数据将阻止您获得恰好 your_gpu_number 倍的加速。因此，您需要优化每个 GPU 的工作负载，以最大限度地提高性能，并尽可能避免设备间通信。

第二个是将你的神经网络分成你拥有的多个设备，训练并合并它们。

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

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

https://gist.github.com/j-min/69aae99be6f6acfadf2073817c2f61b0