python - 联合 Tensorflow 内存使单个 GPU 过载

Question

我正在尝试使用 tensorflow 进行联合学习。 我创建了一个基于 this tutorial for MNIST 的模型。现在我更改它:我为其提供大小为 112x112 的图像，每个图像有 3 个 channel (即输入层的大小为 112x112x3)。当我尝试使用多个模型(大约 50 个模型，每个模型大约 100 个图像)时，出现“内存不足”异常。查看 nvidia-smi 的输出，这是有道理的:出于某种原因，我的程序仅从单个 GPU 获取内存。我怎样才能避免它？

一些评论:

• 我使用allow_growth，因此我非常确定显示的内存使用情况是实际的内存使用情况。我确实可以看到它随着时间的推移而增长。
• TensorFlow 可以识别其他 GPU:在我运行程序之前，它们已占用 0MB。日志中还有行添加可见 GPU 设备:0, 1, 2, 3, 4, 5, 6, 7。

Answer 1

您在 TFF 中遇到了一个有趣且棘手的边缘。

TLDR:TFF 以非正统的方式使用 TF；这会遇到 TF 中的一个错误，其中设备放置无法通过 tf.data.Dataset.reduce 正确传播，并且 TFF 的操作只能放置在单个 GPU 上。

TFF 贡献者目前正在致力于向 TFF 的默认配置添加解决方法，这基本上涉及或多或少“欺骗”TensorFlow 的 AutoGraph 以生成 tf.while 循环，而不是数据集缩减。如果您从源代码构建，则可以交换 dataset.reduce对于特定的 for 循环:

num_examples_sum = 0
for batch in iter(dataset):
  num_examples_sum = reduce_fn(num_examples_sum, batch)

这应该会导致所有 GPU 都得到利用。

如果它们仍然出现，请尝试调整 clients_per_thread参数(与前一个 num_client_executors 参数相反)向上；默认情况下，TFF 将尝试并行运行所有客户端。