python - 联合 Tensorflow 内存重载单个 GPU

Question

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

我用 allow_growth ，所以我很确定显示的内存使用量是实际的内存使用量。我确实可以看到它随着时间的推移而增长。

TensorFlow 识别其他 GPU:在我运行程序之前，它们已占用 0MB。还有线路Adding visible gpu devices: 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 将尝试并行运行所有客户端。