使用tf.data.Dataset.from_generator训练模型性能下降问题排查

你的核心问题出在Generator的输出逻辑错误：自定义的Generator类一次性返回了整个训练集作为单个样本（单个超大batch），这直接导致model.fit中的batch_size=10_000参数完全失效。

在DataFrame训练方式中，每一轮会自动将80万条训练数据拆分为800000/10000=80个batch，每完成一个batch就更新一次模型参数；而Dataset方式下，每一轮仅处理1个超大batch，仅更新1次参数。10轮训练后，DataFrame方式累计完成800次参数更新，Dataset方式仅完成10次，自然训练效率极低，准确率无法快速提升。

方案1：修正Generator，按batch返回数据 #

修改Generator逻辑，每次yield一个指定大小的batch数据，同时添加循环和shuffle逻辑模拟epoch迭代：

class BatchGenerator:
    def __init__(self, X, y, batch_size):
        self.X = X.values
        self.y = y.values
        self.batch_size = batch_size
        self.n_samples = len(X)
        self.index = 0

    def __call__(self):
        while True:
            # 到达数据集末尾时重置索引并打乱数据
            if self.index + self.batch_size > self.n_samples:
                self.index = 0
                perm = np.random.permutation(self.n_samples)
                self.X = self.X[perm]
                self.y = self.y[perm]
            # 提取当前batch
            batch_X = self.X[self.index:self.index+self.batch_size]
            batch_y = self.y[self.index:self.index+self.batch_size]
            self.index += self.batch_size
            yield batch_X, batch_y

构建Dataset并训练：

ds_train = tf.data.Dataset.from_generator(
    lambda: BatchGenerator(X_train, y_train, batch_size=10_000),
    output_signature=(
        tf.TensorSpec(shape=(10_000, n_features), dtype=tf.float32),
        tf.TensorSpec(shape=(10_000, 1), dtype=tf.int32)
    )
)

set_seeds()
model = define_model()
compile(model)
# 这里steps_per_epoch需要指定为总样本数/ batch_size，避免无限迭代
history = model.fit(ds_train, epochs=10, steps_per_epoch=800000//10000, verbose=True)

方案2：更简洁的Dataset构建方式（无需自定义Generator） #

如果数据可部分载入内存，直接使用TensorFlow原生API构建Dataset，自动处理batch和shuffle：

ds_train = tf.data.Dataset.from_tensor_slices((X_train.values, y_train.values))
# 打乱缓冲区设为10万（根据内存调整），按batch划分，预取提升效率
ds_train = ds_train.shuffle(buffer_size=100_000).batch(10_000).prefetch(tf.data.AUTOTUNE)

set_seeds()
model = define_model()
compile(model)
history = model.fit(ds_train, epochs=10, verbose=True)

方案3：针对HDF多文件的专业优化 #

针对你实际的大体积HDF多文件场景，推荐用tf.data的文件并行读取方案，无需载入全部数据：

import h5py

def load_hdf_file(file_path):
    # 用tf.py_function包装HDF读取逻辑
    def _load(file_path):
        file_path_str = file_path.numpy().decode('utf-8')
        with h5py.File(file_path_str, 'r') as f:
            X = f['X'][:]
            y = f['y'][:]
        return X, y
    X, y = tf.py_function(_load, [file_path], [tf.float32, tf.int32])
    # 固定张量形状，保证Dataset兼容性
    X.set_shape((None, n_features))
    y.set_shape((None, 1))
    return tf.data.Dataset.from_tensor_slices((X, y))

# 获取所有HDF文件路径
hdf_files = tf.data.Dataset.list_files('/your/hdf/path/*.h5')
# 并行读取文件、打乱、batch、预取
ds_train = hdf_files.interleave(
    load_hdf_file,
    num_parallel_calls=tf.data.AUTOTUNE,
    cycle_length=4  # 同时处理的文件数，根据CPU核心调整
).shuffle(buffer_size=100_000).batch(10_000).prefetch(tf.data.AUTOTUNE)

• 确保Dataset的shuffle操作：DataFrame方式的model.fit默认会shuffle数据，Dataset方式需手动添加shuffle()避免模型过拟合到数据顺序。
• 控制迭代次数：使用无限循环的Generator时，需通过steps_per_epoch指定每轮的batch数量，避免训练无限进行。
• 预取优化：添加prefetch(tf.data.AUTOTUNE)让TensorFlow在训练当前batch时提前准备下一个batch，提升训练速度。

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