18

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

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)吗?output_buffer_size或者,如果数据集来自,则可能应该在每个将要使用的操作之后tf.contrib.data

4

1 回答 1

18

github上的讨论中, 我发现了 mrry 的评论:

请注意,在 TF 1.4 中将有一个 Dataset.prefetch() 方法,可以更轻松地在管道中的任何位置添加预取,而不仅仅是在 map() 之后。(您可以通过下载当前的夜间版本来尝试。)

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

所以这意味着可以通过任何命令进行预取,并且它适用于上一个命令。到目前为止,我注意到最大的性能提升只是将它放在最后。

还有一个关于Dataset.map 、 Dataset.prefetch 和 Dataset.shuffle 中 buffer_size 的含义的讨论,其中 mrry 解释了更多关于预取和缓冲区的内容。

2018 年 10 月 1 日更新

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

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

于 2017-11-01T22:53:24.173 回答