python - 将数据从 `tf.data.Dataset` 分发给多个工作人员(例如 Horovod)

Question

使用 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 问题。)

Answer 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 的低效率问题:

Important caveats:

• Be sure to shard before you use any randomizing operator (such as shuffle).

• Generally it is best if the shard operator is used early in the dataset pipeline. For example, when reading from a set of TFRecord files, shard before converting the dataset to input samples. This avoids reading every file on every worker. The following is an example of an efficient sharding strategy within a complete pipeline:

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)