将数据从 `tf.data.Dataset` 分发给多个工作人员（例如，对于 Horovod）答案

【问题标题】：Distribute data from `tf.data.Dataset` to multiple workers (e.g. for Horovod)将数据从 `tf.data.Dataset` 分发给多个工作人员（例如，对于 Horovod）
【发布时间】：2020-09-10 12:33:17
【问题描述】：

使用 Horovod，您基本上可以运行 N 个独立实例（因此它是一种图间复制），它们通过特殊的 Horovod 操作（基本上是广播 + 减少）进行通信。

现在假设实例 0 或其他一些外部实例加载您的数据（通过 tf.data.Dataset）。您将如何将iterator.get_next() 分配给每个实例？使用 Horovod 广播效率低下，因为您会将所有数据复制到所有实例。

在每个实例中都有数据集，并在那里进行所有加载，然后在数据集上使用shard 也会效率低下，因为您将在任何地方加载数据，然后丢弃 (N-1)/ N个。所以这就是为什么也不想要分片，而是只在单个（生产者/数据集工作人员）实例中加载数据集，然后将批次分配给所有火车工作人员。

我猜 TF MultiDeviceIterator 提供了一些类似的功能（或基本上完全一样），但我不确定这是否与 Horovod 一起使用，以及您将如何设置它？

或者也许您可以通过 TF 工作人员以某种方式进行分发（guide？（也许这也是您配置 MultiDeviceIterator 的方式？）

如果可能，这应该通过 TensorFlow 操作/函数（有许多相关函数可能已经给了我这个，但我可能不知道它们，或者误解了它的工作原理）。或者答案可能是 TensorFlow 还没有提供任何此类功能？（知道这点还是有用的。然后我会用 C++ 实现我自己的解决方案，包装成 TensorFlow。但在这样做之前，最好知道这是否真的有必要。）

（相关也是this Horovod issue。）

（这个问题实际上比 Horovod 更通用，虽然 Horovod 可能是一个很好的例子。分布式 TensorFlow 可能总是有这个问题？）

（我收集了所有分布式 TensorFlow 术语和方面的概述here，主要是为了澄清。）

（相关的还有（可能？）this、this、this、this 或 this 问题。）

【问题讨论】：

我认为 Horovod 不可能做到这一点，而且绝对不能使用普通的 TensorFlow。你需要的是一个MPI scatter operation，但Horovod 只实现了all-reduce、all-gather 和broadcast 语义。如果您将数据存储为 TFRecord 文件，您最好的选择可能是 split them into many files 并从每个实例中读取一个分片，如答案中所建议...
@jdehesa 正如你提到的，我添加了一个feature request about adding MPI scatter support for Horovod。
TensorFlow 2.3.0 正在引入 tf.data service 用于分布式训练。我不确定它是否能特别解决您的问题，但可能值得您关注。
@jdehesa 哦，谢谢。当您使用共享作业（设置job_name）时，这看起来与我正在寻找的完全一样。或者至少差不多，因为它不会均匀分布，而是按照先到先得的原则（这可能没问题）。有趣的是，当作业不共享时，这是我目前已经实施的另一种解决方案，让每个工作人员只使用不同的随机种子进行数据集洗牌。
阿尔伯特，看看粉丝商店。如果您有兴趣，那么我可以发布一个答案，这是如何工作的。 github.com/TACC/FanStore

标签： python tensorflow tensorflow-datasets horovod

【解决方案1】：

正如你所说，复制每个实例中的数据并为每个实例分片数据是不切实际的。

然后一种解决方案是在数据流程中分离数据，并让每个实例从数据流程中提取数据，如下图所示。例如，可以使用队列来建立这种通信。

在这样的系统中，数据处理会加载数据集，将其预处理成批次，然后将批次推送到队列中。然后，每个训练实例将从该队列中提取批次。例如，您可以将队列作为生成器传递到数据集 API（请参阅tf.data.Dataset.from_generator）。此外，如果批次生产速度不够快，则可以创建更多数据处理来提高批次吞吐量。

根据您的用例，实现细节会有所不同。更多信息，您可以查找Networking and Interprocess communication和Multiprocessing pipes and queues。

                                                             Training        
                                                         +--------------+  ++
                                                         |              |   |
                                                    +----+  Instance 1  |   |
                                                    |    |              |   |
                                                    |    +--------------+   |
                                                    |                       |
                      Preprocessing                 |                       |
                  +--------------------+            +---->      X           |
                  |                    |            |                       |
             Load |                    | Batches    +           X           |
    Dataset+------>    Data Process    +--------->Queue                     |  N instances
                  |                    |            +           X           |  Distributed training
                  |                    |            |                       |  For example, using
                  +--------------------+            +---->      X           |  Horovod broadcast + reduce
                                                    |                       |
                                                    |        Training       |
                                                    |    +--------------+   |
                                                    |    |              |   |
                                                    +----+  Instance N  |   |
                                                         |              |   |
                                                         +--------------+  ++

对于 TensorFlow 实现，您可以使用 tf.data.Dataset.shard 和 tf.data.TFRecordDataset。

该文档解决了您使用 TFRecords 效率低下的问题：

重要提示：

确保在使用任何随机化操作符（如随机播放）之前进行分片。

通常，最好在数据集管道的早期使用分片运算符。例如，从一组 TFRecord 文件中读取时，在将数据集转换为输入样本之前进行分片。这避免了读取每个工作人员的每个文件。以下是完整管道中高效分片策略的示例：
d = Dataset.list_files(pattern)
d = d.shard(num_workers, worker_index)
d = d.repeat(num_epochs)
d = d.shuffle(shuffle_buffer_size)
d = d.interleave(tf.data.TFRecordDataset,
                 cycle_length=num_readers, block_length=1)
d = d.map(parser_fn, num_parallel_calls=num_map_threads)

【讨论】：

是的，这基本上就是我想要的。但我的问题是：如何实现它（使用 TensorFlow）。我已经将最初的数据集作为tf.data.Dataset （在您的图中最左边的数据集）。 TF 中的预处理部分非常简单。但是你会如何组织这个？我想这将在专用的 TF 工作人员（通过分布式 TF）上运行。但是如何实现队列（在 TF 中）？这是我的主要问题。我认为 Python 函数在这里用处不大。如果你说TF中不存在这部分，我宁愿用C++实现这部分。
@Albert 感谢您的评论。我已经更新了答案，提供了有关如何使用 TensorFlow 解决此问题的更多详细信息。
但是这段代码正是我不想做的（使用分片）。我想要一个单一的数据集管道并将其分发给所有火车工作人员（就像我描述的那样，就像你在图片中一样）。
@Albert 您提到不想进行分片，因为它“也会效率低下，因为您会在任何地方加载数据，然后丢弃 (N-1)/N 个数据”。在建议的解决方案中，只为每个实例加载必要的数据，并且不会抛出任何内容。 shard 解决方案在实现上确实不同，但在均匀采样的结果上应该是等效的。
你确定吗？那么shard 上的TFRecordDataset 是专门实现的吗？我在哪里可以看到这个？但无论如何，我的问题是关于不要那样做。我想知道我怎样才能做到我所要求的（以及您在图片中概述的内容）。

【解决方案2】：

我建议查看YogaDL。它允许您缓存数据集，这样在训练（或重新训练）期间，您将只访问该分片上所需的数据，而不是丢弃 (N-1)/N 次读取的数据。

【讨论】：