如何将TensorFlow数据集缓存指定epoch数后重新生成？

该需求可以实现，以下是两种高效实现方案，性能基本和当前4ms/epoch的训练速度一致。

方案1：数据集流水线原生实现（性能最优） #

核心逻辑是将「生成原始数据→缓存→重复使用X轮」作为一个周期，无限拼接周期即可，无需额外回调介入：

import tensorflow as tf

X = 10  # 自定义每X个epoch重建一次数据集

def build_single_cycle():
    # 每次调用都会重新生成一轮完整的原始数据集
    base_ds = tf.data.Dataset.from_generator(
        lambda: sample_generator(
            vocab_name, model_name, n_synonyms=5, use_closest_words=True
        ),
        output_signature=(
            (
                tf.TensorSpec(shape=(1,), dtype=tf.int32),
                tf.TensorSpec(shape=(None,), dtype=tf.int32),
                tf.TensorSpec(shape=(None,), dtype=tf.int32),
            ),
            tf.TensorSpec(shape=(1, None), dtype=tf.float64),
        )
    )
    # 单周期内复用缓存，仅第一次迭代生成数据
    return base_ds.batch(cfg.tf_parms.batch_size).prefetch(12).cache().repeat(X)

# 拼接多周期数据集，总训练轮数控制为600
train_dataset = tf.data.Dataset.from_generator(
    lambda: iter(build_single_cycle().repeat()),
    output_signature=build_single_cycle().element_spec
).take(600)

训练调用代码无需额外修改：

model.fit(
    x=train_dataset,
    callbacks=[cp_callback],
    epochs=600
)

方案2：自定义回调控制缓存重置（灵活度更高） #

如果需要根据训练状态动态调整重建间隔，可通过磁盘缓存+回调删除缓存文件触发重建：

1. 修改数据集定义，指定磁盘缓存路径：

cache_path = "./tf_dataset_cache"
base_ds = dataset.batch(cfg.tf_parms.batch_size).prefetch(12).cache(cache_path)

2. 自定义数据集重建回调：

import os
from tensorflow.keras.callbacks import Callback

class RebuildDatasetCallback(Callback):
    def __init__(self, rebuild_interval, cache_path):
        super().__init__()
        self.rebuild_interval = rebuild_interval
        self.cache_path = cache_path
    
    def on_epoch_end(self, epoch, logs=None):
        if (epoch + 1) % self.rebuild_interval == 0:
            # 删除旧缓存文件，下一轮迭代自动重新生成数据集并缓存
            for suffix in [".index", ".data-00000-of-00001"]:
                cache_file = self.cache_path + suffix
                if os.path.exists(cache_file):
                    os.remove(cache_file)

3. 训练时传入回调即可：

model.fit(
    x=base_ds,
    callbacks=[cp_callback, RebuildDatasetCallback(rebuild_interval=X, cache_path=cache_path)],
    epochs=600
)

注意：两种方案都要求sample_generator每次调用时都能生成完整的一轮训练数据，无全局状态残留。

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