Tensorflow 数据集预取和缓存选项的正确用法是什么？

Question

我已经阅读了 TF 页面和一些关于使用 prefetch() 和 cache() 来加速模型输入管道的帖子，并尝试在我的数据上实现它。Cache() 按预期为我工作，即在第一个纪元中从 dist 读取数据，在所有后续纪元中它只是从内存中读取数据。但我在使用 prefetch() 时遇到很多困难，我真的不明白何时以及如何使用它。有人可以帮我吗？我真的需要一些帮助。我的应用程序是这样的：我有一组大型 TFRecord 文件，每个文件都包含一些在馈送到我的网络之前要处理的原始记录。它们将被混合（不同的样本流），所以我要做的是：

def read_datasets(pattern, numFiles, numEpochs=125, batchSize=1024, take=dataLength):

    files = tf.data.Dataset.list_files(pattern)

    def _parse(x):
        x = tf.data.TFRecordDataset(x, compression_type='GZIP')
    return x

    np = 4 # half of the number of CPU cores
    dataset = files.interleave(_parse, cycle_length=numFiles, block_length=1, num_parallel_calls=np)\
    .map(lambda x: parse_tfrecord(x), num_parallel_calls=np)
    dataset = dataset.take(take)
    dataset = dataset.batch(batchSize)
    dataset = dataset.cache()
    dataset = dataset.prefetch(buffer_size=10)
    dataset = dataset.repeat(numEpochs)
    return dataset

interleave 函数中的 parse_tfrecord(x) 函数是数据应用于模型之前所需的预处理，我的猜测是预处理时间与网络的批处理时间相当。我的整个数据集（包括所有输入文件）包含大约 500 批 1024 个样本。我的问题是：

1-如果我进行缓存，我真的需要预取吗？

2- 映射、批处理、缓存、预取和重复的顺序是否正确？

3-Tensorflow文档说预取的缓冲区大小指的是数据集元素，如果是批处理的，则指的是批次的数量。那么在这种情况下我将读取 10 批 1024 个示例，对吧？我的问题是，通过更改预取缓冲区大小，我没有看到运行时间有任何差异，并且即使将缓冲区大小设置为 1000 或更大，内存消耗也没有太大变化。

Answer 1

我为斯坦福大学的 Andrew Ng 找到了这个很好的解释：https ://cs230.stanford.edu/blog/datapipeline/#best-practices

“当 GPU 对当前批次进行前向/后向传播时，我们希望 CPU 能够处理下一批数据，以便立即准备就绪。作为计算机中最昂贵的部分，我们希望 GPU 能够充分发挥作用。”在训练过程中一直使用，我们称之为消费者/生产者重叠，其中消费者是GPU，生产者是CPU。

使用，您可以通过在管道末尾（批处理之后）tf.data进行简单的调用来完成此操作。dataset.prefetch(1)这将始终预取一批数据并确保始终有一批数据准备就绪。

在某些情况下，预取多个批次可能很有用。例如，如果预处理的持续时间变化很大，则预取 10 个批次将平均 10 个批次的处理时间，而不是有时等待更长的批次。

举一个具体的例子，假设 10% 的批次需要 10 秒来计算，90% 的批次需要 1 秒。如果 GPU 需要 2 秒来训练一个批次，那么通过预取多个批次，您可以确保我们永远不会等待这些罕见的较长批次。”

我不太确定如何确定每个批次的处理时间，但这是下一步。如果您的批次大致花费相同的时间来处理，那么我相信 prefetch(batch_size=1) 应该足够了，因为您的 GPU 不会等待 CPU 完成处理计算量大的批次。

Answer 2

您可以查看此Stackoverflow 答案，快速了解 TensorFlow 数据集的函数cache()和prefetch().

另外，我发现这个Tensorflow 文档对于优化 Api 的性能非常有帮助tf.Data。他们为各种执行方式指定了基准和执行时间。您还可以分别找到有关数据的序列化和并行加载和转换及其执行时间的信息。