TensorFlow缓存数据集后迭代为空，如何统计缓存数据集样本数？

问题场景 #

使用tf.data.Dataset.from_generator创建数据集时遇到两个矛盾问题：

1. 先添加repeat()、prefetch()再cache()，调用reduce统计样本数时得到空数据集（结果为0）：

output_signature = (tensorflow.TensorSpec(shape=(64, 64, 2), dtype=tensorflow.float32),
                        tensorflow.TensorSpec(shape=(), dtype=tensorflow.int16))
dataset = tensorflow.data.Dataset.from_generator(MSUMFSDFrameGenerator(pathlib.Path(dataset_locations["mfsd"]), True), output_signature=output_signature)
dataset = dataset.repeat().prefetch(buffer_size=tensorflow.data.AUTOTUNE).cache(filename)
size = dataset.reduce(0, lambda x,_: x+1).numpy()

2. 先统计样本数再添加cache()，虽然能得到正常样本数，但缓存不生效，训练时仍需重新调用生成器生成数据：

output_signature = (tensorflow.TensorSpec(shape=(64, 64, 2), dtype=tensorflow.float32),
                        tensorflow.TensorSpec(shape=(), dtype=tensorflow.int16))
dataset = tensorflow.data.Dataset.from_generator(MSUMFSDFrameGenerator(pathlib.Path(dataset_locations["mfsd"]), True), output_signature=output_signature)
dataset = dataset.repeat().prefetch(buffer_size=tensorflow.data.AUTOTUNE)
size = dataset.reduce(0, lambda x,_: x+1).numpy()
dataset = dataset.cache(filename)

解决方法 #

核心问题是**repeat()和cache()的顺序错误**：将无限重复的数据集传给cache()会导致缓存逻辑异常，同时reduce无法处理无限序列得到正确结果。正确步骤是先缓存原始有限数据集，再添加重复和预取操作，同时在缓存前统计原始数据集的样本数：

import tensorflow as tf
import pathlib

output_signature = (
    tf.TensorSpec(shape=(64, 64, 2), dtype=tf.float32),
    tf.TensorSpec(shape=(), dtype=tf.int16)
)

# 1. 创建原始生成器数据集（不添加repeat）
dataset = tf.data.Dataset.from_generator(
    MSUMFSDFrameGenerator(pathlib.Path(dataset_locations["mfsd"]), True),
    output_signature=output_signature
)

# 2. 统计原始有限数据集的样本数
size = dataset.reduce(0, lambda x, _: x + 1).numpy()

# 3. 先缓存原始数据，再添加repeat和prefetch
dataset = dataset.cache(filename).repeat().prefetch(buffer_size=tf.data.AUTOTUNE)

原理说明 #

• 先缓存原始有限数据集，cache()会将生成器产出的所有数据写入缓存文件/内存，后续repeat()直接基于缓存数据重复，训练时无需重新调用生成器，缓存完全生效。
• 统计样本数针对原始有限数据集操作，reduce可以正常遍历完成并得到真实的单轮样本数量。
• 错误写法中，repeat()先将数据集转为无限序列，cache()无法缓存无限数据，导致统计时出现异常；先统计再缓存的写法，会提前耗尽生成器的迭代资源，后续缓存无法捕获有效数据。

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