tensorflow - 使用多个线程对数据进行多个分片是否可以最大限度地缩短训练时间？

Question

我的主要问题是:我有 204 GB 的训练 tfrecord 文件，其中包含 200 万张图像，还有 28GB 用于验证 tf.record 文件，其中包含 302900 张图像。训练一个 epoch 需要 8 小时，训练需要 33 天。我想通过使用多个线程和分片来加快速度，但我对一些事情有点困惑。

在 tf.data.Dataset API有分片功能，因此在文档中他们提到了以下关于分片功能的内容:

Creates a Dataset that includes only 1/num_shards of this dataset.

This dataset operator is very useful when running distributed training, as it allows each worker to read a unique subset.

When reading a single input file, you can skip elements as follows:

d = tf.data.TFRecordDataset(FLAGS.input_file)
d = d.shard(FLAGS.num_workers, FLAGS.worker_index)
d = d.repeat(FLAGS.num_epochs)
d = d.shuffle(FLAGS.shuffle_buffer_size)
d = d.map(parser_fn, num_parallel_calls=FLAGS.num_map_threads)

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(FLAGS.pattern)
d = d.shard(FLAGS.num_workers, FLAGS.worker_index)
d = d.repeat(FLAGS.num_epochs)
d = d.shuffle(FLAGS.shuffle_buffer_size)
d = d.repeat()
d = d.interleave(tf.data.TFRecordDataset,
             cycle_length=FLAGS.num_readers, block_length=1)

d = d.map(parser_fn, num_parallel_calls=FLAGS.num_map_threads)

这是我的问题:

1- tf.records 文件的数量和分片的数量之间有什么关系吗？分片(worker)的数量取决于您拥有的 CPU 数量，还是您拥有的 tf.records 文件数量？以及如何创建它，只需将分片数量设置为特定数量？ ，或者我们需要将文件拆分为多个文件，然后设置具体的分片数量。记下工作人员数量与分片数量的关系

2- 创建多个 tf.records 文件有什么好处？有人说here这与您需要以更好的方式对 tf.records 进行洗牌的情况有关，但由于 tf.Dataset API 中存在 Shuufle 方法，我们不需要这样做，其他人说 here它只是将您的数据分割成较小尺寸的部分。我的问题我是否需要第一步将 tf.records 文件拆分为多个文件

3-现在我们来到map函数中的num_threads(新版本的tensorflwo中的num_paralle_calls)应该与您拥有的分片数量相同。我搜索的时候发现有人说如果你有10个分片，2个线程，每个线程就需要5个分片。

4- d.interleave 函数怎么样，我知道它是如何工作的，正如 example 中提到的那样。但我再次错过了连接 num_threads，例如周期长度

5- 如果我想使用多个 GPU，我应该使用分片吗？如前所述in the accepted comment here

作为总结，我对( tf.records 文件数量、num_shards(workers)、循环长度、num_thread(num_parallel_calls) 之间的关系感到困惑。为了最大限度地减少训练时间，创建什么更好的情况两种情况(使用多个 GPU 和使用单个 GPU)

Answer 1

我是 tf.data 开发人员。让我们看看我是否可以帮助回答您的问题。

1) 听起来您有一个大文件。要使用多个工作进程处理它，将其拆分为多个较小的文件是有意义的，以便 tf.data 输入管道可以处理不同工作进程上的不同文件(通过在文件名列表)。

2) 如果不将单个文件拆分为多个记录，每个工作人员将不得不读取整个文件，从而消耗 n 倍的 IO 带宽(其中 n 是工作人员数量)。

3) map 转换的线程数量与分片数量无关。每个分片将由每个工作线程上的 num_parallel_calls 进行处理。一般来说，将 num_parallel_calls 设置为与工作线程上可用核心的数量成比例是合理的。

4) interleave 转换的目的是将多个数据集(例如从不同的 TFRecord 文件读取)合并为一个数据集。鉴于您的用例，我认为您不需要使用它。

5) 如果您想为多个 GPU 提供数据，我建议您使用您引用的评论中列出的第一个选项，因为它是最简单的。基于分片的解决方案需要在每个工作线程上创建多个管道(针对每个 GPU)。

为了最大限度地减少训练时间(对于单个或多个 GPU，您需要确保输入管道生成数据的速度快于 GPU 处理数据的速度。输入管道的性能调整已在 here 中讨论。