tensorflow - PREFETCH 样本后 TFRecord 读取管道速度减慢

Question

我已将训练数据分成多个 tf-record 文件，并使用这段代码读取它们:

SHUFFLE_BUFFER = 64
PREFETCH = 256
dataset = tf.data.TFRecordDataset(filenames)
dataset = dataset.shuffle(SHUFFLE_BUFFER) 
dataset = dataset.map(_parse_image_function, num_parallel_calls=tf.data.experimental.AUTOTUNE)
dataset = dataset.batch(BATCH_SIZE)
dataset = dataset.prefetch(PREFETCH)
dataset = dataset.repeat()

该数据集直接输入到 model.fit(dataset)。

第一个 PREFETCH 样本加载速度很快，GPU 利用率始终保持在 80% 以上。然而，在那之后，快速读取似乎停止了，GPU 利用率下降，训练时间大幅减慢。 有人知道可能出了什么问题吗？

Answer 1

如果不了解更多详细信息(存储后端、记录大小、每个文件的记录数、文件数、_parse_image_function 中的任何 io 操作？，...)，这有点难以诊断

我的第一个怀疑是 tf.data.TFRecordDataset(filenames) - 在下一个文件之后打开一个文件可能会导致延迟峰值，从而可能暂时导致数据集 CPU 管道不足。 (多个较小的文件从自动预读中获得的好处也可能较低)

我会尝试在 tf.data.TFRecordDataset(filenames) 之后添加额外的预取，以解耦 IO(并且可能交错来自不同文件的记录(num_parallel_reads 参数))。

如果预取没有帮助，我会尝试对 num_parallel_calls 进行硬编码(主要是因为我还没有阅读自动调整代码 - 如果您的管道需要的并行度超过默认并行度，则可能使用私有(private)线程池)。

根据您的存储后端 - 一旦训练速度减慢(以测试/优化数据集)，重复训练就会重新启动，可能只是从各种缓存中提取数据，并且一旦使用的数据集超出缓存，速度可能会减慢。