如何强制TensorFlow使用所有可用的GPU？

Question

如何强制TensorFlow使用所有可用的GPU？

10

我有一个8 GPU的集群，当我运行来自Kaggle的一段Tensorflow代码（如下所示），它只使用了一个GPU而不是所有8个。我使用nvidia-smi进行了确认。

# Build model
inputs = Input((IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS))
...
outputs = Conv2D(1, (1, 1), activation='sigmoid') (c9)

model = Model(inputs=[inputs], outputs=[outputs])

sgd = optimizers.SGD(lr=0.03, decay=1e-6, momentum=0.9, nesterov=True)
model.compile(optimizer=sgd, loss='binary_crossentropy', metrics=[mean_iou])
model.summary()
    
# Fit model
results = model.fit(X_train, Y_train, validation_split=0.05, batch_size = 32, verbose=1, epochs=100)

我想使用MXNet或其他方法在所有可用的GPU上运行此代码。但是，我不确定该怎么做。所有资源都只显示如何在mnist数据集上执行此操作。我有自己的数据集，读取方式不同。因此，不太确定如何修改代码。

- Jonathan

1个回答

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Peter Szoldan · Accepted Answer

TL;DR: 使用 tf.distribute.MirroredStrategy() 作为作用域，例如

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
    [...create model as you would otherwise...]

如果您不指定任何参数，tf.distribute.MirroredStrategy() 将使用所有可用的GPU。如果您希望指定使用哪些GPU，可以像这样进行设置: mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])。

有关实现细节和其他策略，请参考 Distributed training with TensorFlow 指南。

早期的回答（现已过时：deprecated, removed as of April 1, 2020.）：使用 Keras 的 multi_gpu_model()。

TensorFlow 2.0现在有tf.distribute模块，它是“跨多个设备运行计算的库”。它建立在“分布策略”的概念上。您可以指定分发策略，然后将其用作作用域。 TensorFlow会透明地为您拆分输入，并行化计算，并将输出合并。反向传播也受此影响。由于所有处理现在都在后台完成，因此您可能希望熟悉可用策略及其参数，因为它们可能会极大地影响训练速度。例如，您想让变量驻留在CPU上吗？那么使用tf.distribute.experimental.CentralStorageStrategy()。有关更多信息，请参阅Distributed training with TensorFlow指南。

从Tensorflow Guide中得到的早期答案（现已过时，仅供参考）：

如果您的系统中有多个GPU，则默认情况下会选择ID最低的GPU。

如果要使用多个GPU，不幸的是，您必须手动指定要放置在每个GPU上的张量。

with tf.device('/device:GPU:2'):

更多信息请参见Tensorflow指南：使用多个GPU。

在如何将网络分布到多个GPU上方面，有两种主要方法。

你可以将网络按层分布在各个GPU上。这种实现较为简单，但性能提升不大，因为GPU需要等待其他GPU完成操作。
你可以在每个GPU上创建一个名为“tower”的独立网络副本。当你输入八元组网络时，你需要将输入批次分成8个部分并进行分发。让网络前向传播，然后对梯度进行求和和反向传播。这样可以实现与GPU数量几乎线性的加速。然而，这种实现较为困难，因为你还需要处理与批量归一化相关的复杂性，并且非常建议确保正确随机批处理。这里有一个很好的教程。你还应该查看Inception V3代码的相关内容，以获取如何构建这种东西的想法。特别是_tower_loss()、_average_gradients()和从for i in range(FLAGS.num_gpus):开始的train()部分。

如果你想尝试Keras，现在它已经通过multi_gpu_model()大大简化了多GPU训练。它可以为你完成所有繁重的工作。