Tensorflow 不使用多个 GPU - 获得 OOM答案

【问题标题】：Tensorflow not using multiple GPUs - getting OOMTensorflow 不使用多个 GPU - 获得 OOM
【发布时间】：2021-09-10 22:12:47
【问题描述】：

我在多 GPU 机器上遇到 OOM，因为 TF 2.3 似乎只使用一个 GPU 分配张量。

tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at conv_ops.cc:539 : 
Resource exhausted: OOM when allocating tensor with shape[20532,64,48,32] 
and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc.

但是当我运行我的代码时，tensorflow 确实可以识别多个 GPU：

Adding visible gpu devices: 0, 1, 2

我还需要做什么才能让 TF 使用所有 GPU？

【问题讨论】：

如果我没记错的话，具有这种形状的张量需要 8 GB 内存。你的机器上真的有足够的内存吗？
我机器上的每个 GPU 都有 16160MiB 应该够用了吧？
应该够了。我自己并没有真正使用过 tensorflow，只是偶然发现了你的问题。你能检查一下 TF 是否真的看到/有权访问完整的 GPU 内存吗？就像在stackoverflow.com/questions/36123740/… 中回答的一样@

标签： tensorflow gpu multi-gpu

【解决方案1】：

直接的答案是肯定的，您确实需要做更多工作才能让 TF 识别多个 GPU。您应该参考本指南，但 tldr 是

mirrored_strategy = tf.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  ...

https://www.tensorflow.org/guide/distributed_training#using_tfdistributestrategy_with_tfkerasmodelfit

但在你的情况下，正在发生其他事情。虽然这一个张量可能会触发 OOM，但很可能是因为之前分配了一些大张量。

第一个维度，你的批量大小，是 20532，真的很大。由于它的分解是 2**2 × 3 × 29 × 59，我猜你正在使用 CHW 格式，你的源图像是 3x64x128，经过几次卷积后被修剪。我怀疑是无意的广播。打印一个 model.summary() 然后查看每一层的张量的大小。您可能还需要查看批处理。

【讨论】：