python - Tensorflow Dataset API shuffle 将性能降低 9 倍

Question

我正在使用 Tensorflow 数据集 API 来获取一堆文件名；打乱文件名；执行 python 函数来加载图像文件，对其进行预处理，并将其转换为张量；然后缓存、重复和批处理它们。到目前为止，一切顺利。

当我向张量添加 shuffle() 时，性能下降了 9 倍。同样，当我执行 self.dataset.apply(tf.data.experimental.shuffle_and_repeat(16384)) 时。

为什么shuffle会对性能造成如此严重的影响，我该如何解决它？

代码:

filenames = tf.data.Dataset.list_files(self.FILE_PATTERN).shuffle(buffer_size=16384)
dataset = filenames.map(lambda filename: self.pp(filename), 
num_parallel_calls=self.N_CPUS)
dataset = dataset.cache("./cachefile")
# The line below (shuffle_and_repeat) made performance very bad (1s/step without, 9s/step with)
# dataset = dataset.apply(tf.data.experimental.shuffle_and_repeat(16384))
# This too:
# dataset = dataset.repeat().shuffle(16384)
# This works fine, but doesn't shuffle:
dataset = dataset.repeat()
dataset = dataset.batch(self.BATCH_SIZE)
dataset = dataset.prefetch(4)

Answer 1

尝试更改预取参数buffer_size=2

数据集 = dataset.prefetch(2)

预取是一个性能标志，在后台读取下一个数据集以进行下一次迭代。如果预取的 buffer_size 很大，那么它会创建大量的数据集进行迭代，并且可能会因内存不足而减慢速度。