如何在TensorFlow中正确在首个epoch期间缓存Dataset数据？

这个问题我之前也碰到过，本质是TensorFlow的cache()要求必须完整遍历一遍目标数据集才能完成缓存，你的代码顺序刚好踩了坑。

错误原因分析 #

你当前把cache()放在了prefetch之后，而且shuffle_and_repeat在cache之前。这种情况下，cache()要缓存的是已经被repeat处理过的数据集，而当首个epoch结束时，可能因为prefetch预取的残留数据，导致迭代器没有完全读取完整个要缓存的数据集，触发了TensorFlow的保护机制——直接丢弃部分缓存，所以后续还是会从硬盘读数据。

正确的代码调整方案 #

你需要把cache()移动到所有预处理操作（map、padded_batch）之后，且在shuffle_and_repeat和prefetch之前。这样既可以缓存预处理后的结果（避免重复执行耗时的解析和batch操作），又能确保迭代器完整读取整个数据集，顺利完成缓存。

调整后的代码如下：

dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=1)
# 先执行耗时的预处理操作：解析样本、padding成统一batch
dataset = dataset.map(_parser_a, num_parallel_calls=12)
dataset = dataset.padded_batch(20, padded_shapes=padded_shapes, padding_values=padding_values)
# 缓存预处理后的最终样本格式
dataset = dataset.cache()
# 再执行shuffle和repeat，保证训练随机性
dataset = dataset.apply(tf.contrib.data.shuffle_and_repeat(buffer_size=5000, count=1))
# 最后prefetch提升训练效率
dataset = dataset.prefetch(buffer_size=1)

额外优化建议 #

如果你希望shuffle的是单个样本而非batch（通常训练时更推荐这种方式），可以把shuffle提前到预处理之前，拆分shuffle_and_repeat成独立操作：

dataset = tf.data.TFRecordDataset(filenames, num_parallel_reads=1)
# 先shuffle原始样本，保证单样本级别的随机性
dataset = dataset.shuffle(buffer_size=5000)
# 预处理操作
dataset = dataset.map(_parser_a, num_parallel_calls=12)
dataset = dataset.padded_batch(20, padded_shapes=padded_shapes, padding_values=padding_values)
# 缓存预处理结果
dataset = dataset.cache()
# 重复数据集（count值可根据训练需求调整）
dataset = dataset.repeat(count=1)
dataset = dataset.prefetch(buffer_size=1)

为什么这样调整有效 #

1. cache()作用在预处理后的数据集上，直接缓存最终要输入模型的样本格式，彻底省去重复解析TFRecord和padding的开销。
2. 将shuffle/repeat放在cache()之后，每次迭代都会从缓存读取数据并重新shuffle，既保证训练随机性，又不会破坏缓存的完整性。
3. prefetch放在最后，利用预取机制提升训练速度，同时不会干扰缓存的生成过程。

这样调整后，首个epoch结束时，迭代器会完整遍历并缓存所有预处理后的样本，后续epoch就会直接从内存缓存读取数据了。

内容的提问来源于stack exchange，提问作者Maosi Chen