Tensorflow 数据 API - 预取

Question

我正在尝试使用 TF 的新功能，即 Data API，但我不确定 prefetch 是如何工作的。在下面的代码中

def dataset_input_fn(...)
    dataset = tf.data.TFRecordDataset(filenames, compression_type="ZLIB")
    dataset = dataset.map(lambda x:parser(...))
    dataset = dataset.map(lambda x,y: image_augmentation(...)
                      , num_parallel_calls=num_threads
                     )

    dataset = dataset.shuffle(buffer_size)
    dataset = dataset.batch(batch_size)    
    dataset = dataset.repeat(num_epochs)
    iterator = dataset.make_one_shot_iterator()

上面我输入dataset=dataset.prefetch(batch_size) 的每一行之间有关系吗？或者，如果数据集来自tf.contrib.data，则可能应该在使用output_buffer_size 的每个操作之后？

Answer 1

在github的讨论中，我发现了mrry的评论：

请注意，在 TF 1.4 中会有一个 Dataset.prefetch() 方法 可以更轻松地在管道中的任何位置添加预取，而不是 就在地图（）之后。 （您可以通过下载当前的 nightly 构建。）

和

例如，Dataset.prefetch() 将启动一个后台线程来 填充一个类似于 tf.FIFOQueue 的有序缓冲区，这样 下游管道阶段不需要阻塞。然而， prefetch() 实现要简单得多，因为它不需要支持 as 许多不同的并发操作作为 tf.FIFOQueue。

所以这意味着任何命令都可以进行预取，并且它适用于前一个命令。到目前为止，我注意到仅将其放在最后才能获得最大的性能提升。

Meaning of buffer_size in Dataset.map , Dataset.prefetch and Dataset.shuffle 上还有一个讨论，mrry 解释了更多关于预取和缓冲区的内容。

2018 年 10 月 1 日更新：

从 1.7.0 版开始，数据集 API（在 contrib 中）有一个选项 prefetch_to_device。请注意，此转换必须是管道中的最后一个转换，当 TF 2.0 到达时，contrib 将消失。要在多个 GPU 上进行预取工作，请使用 MultiDeviceIterator（示例参见 #13610）multi_device_iterator_ops.py。

https://www.tensorflow.org/versions/master/api_docs/python/tf/contrib/data/prefetch_to_device