联合Tensorflow内存-重载单个GPU

问题描述

我正在尝试使用张量流进行联合学习。 我已经基于this tutorial for MNIST创建了一个模型。现在，我对其进行更改：我为其提供尺寸为112x112的图像，每个图像具有3个通道（即输入层的尺寸为112x112x3）。当我尝试使用多个模型（大约50个，每个大约100张图像）时，出现“内存不足”异常。查看--log-file的输出是有道理的：由于某种原因，我的程序仅从单个GPU占用内存。我该如何避免呢？

一些评论：

• 我使用nvidia-smi，所以我很确定显示的内存使用情况是实际的内存使用情况。我确实可以看到它随着时间的推移而增长。
• TensorFlow可以识别其他GPU：在我运行程序之前，它们已占用0MB内存。日志中也有第allow_growth行。

解决方法

您在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默认会尝试并行运行所有客户端。