使用Generator生成训练数据并基于tf.data.from_generator训练NN的问题

一、Generator的正确写法与代码优化 #

Generator的输出规范 #

对于动态生成数据的场景，Generator每次生成一个批次的样本即可，不需要提前循环steps_per_epoch。model.fit的steps_per_epoch参数会控制每个epoch需要取多少个批次，Dataset会自动重复调用Generator直到取够指定步数。

现有代码的核心错误 #

1. 缩进逻辑错误：你的Generator中，X_train、Y_train的生成代码和yield语句都在for jj in range(steps_per_epoch)循环外部，导致Generator只生成一个批次就终止，这也是你被迫添加repeat(nb_epoch)的根本原因。
2. 验证集数据重复：用同一个datagen生成训练和验证数据，会导致验证集与训练集数据完全一致，无法有效评估模型泛化能力。
3. 冗余的repeat调用：model.fit的epochs参数会自动控制训练轮次，给Dataset加repeat(nb_epoch)会导致数据重复逻辑冲突，反而可能引发异常。

修正后的Generator与训练代码 #

import numpy as np
import tensorflow as tf

# 假设trans1、trans2是已定义的预处理变换
def datagen(batch_size, nw):
    while True:  # 无限生成批次，直到model.fit停止
        X_train = np.zeros(shape=(batch_size, nw))
        Y_train = np.zeros(shape=(batch_size, 3))
        for ii in range(batch_size):
            # 替换为你的实际数据生成逻辑
            X_train[ii] = np.random.uniform(low=-1.0, high=1.0, size=nw)
            Y_train[ii] = [np.mean(X_train[ii]), np.std(X_train[ii]), np.max(X_train[ii])]
        yield trans1.transform(X_train), trans2.transform(Y_train)

# 验证集单独生成，确保数据独立
def val_datagen(batch_size, nw):
    while True:
        X_val = np.zeros(shape=(batch_size, nw))
        Y_val = np.zeros(shape=(batch_size, 3))
        for ii in range(batch_size):
            X_val[ii] = np.random.uniform(low=-1.0, high=1.0, size=nw)
            Y_val[ii] = [np.mean(X_val[ii]), np.std(X_val[ii]), np.max(X_val[ii])]
        yield trans1.transform(X_val), trans2.transform(Y_val)

# 训练参数配置
batch_size = 1024
nb_epoch = 10
datatot_train = 1e6
datatot_val = 0.2 * datatot_train
steps_per_epoch = int(np.ceil(datatot_train / batch_size))
validation_steps = int(np.ceil(datatot_val / batch_size))
nw = 200  # 对应X_train的特征数200

# 构建训练集Dataset
dataset_train = tf.data.Dataset.from_generator(
    lambda: datagen(batch_size, nw),
    output_signature=(
        tf.TensorSpec(shape=(batch_size, nw), dtype=tf.float32),
        tf.TensorSpec(shape=(batch_size, 3), dtype=tf.float32)
    )
).prefetch(tf.data.AUTOTUNE)  # 自动优化预取批次数量

# 构建验证集Dataset
dataset_val = tf.data.Dataset.from_generator(
    lambda: val_datagen(batch_size, nw),
    output_signature=(
        tf.TensorSpec(shape=(batch_size, nw), dtype=tf.float32),
        tf.TensorSpec(shape=(batch_size, 3), dtype=tf.float32)
    )
).prefetch(tf.data.AUTOTUNE)

# 模型训练：无需给Dataset加repeat
history = model.fit(
    dataset_train,
    epochs=nb_epoch,
    steps_per_epoch=steps_per_epoch,
    verbose=1,
    validation_data=dataset_val,
    validation_steps=validation_steps
)

更高效的实现：用TensorFlow原生操作替代Python Generator #

Python Generator受GIL限制，速度不如TensorFlow原生操作，且无法利用图模式优化。建议用TF原生API生成数据：

def generate_single_sample(nw):
    # 用TensorFlow操作生成单个样本，避免NumPy开销
    X = tf.random.uniform(shape=(nw,), minval=-1.0, maxval=1.0, dtype=tf.float32)
    Y = tf.stack([tf.reduce_mean(X), tf.math.reduce_std(X), tf.reduce_max(X)], axis=0)
    # 若trans1、trans2是Scikit-learn预处理类，建议替换为TF原生层（如Normalization）
    X = trans1(X)
    Y = trans2(Y)
    return X, Y

# 构建数据集：先生成无限样本，再分批
dataset_train = tf.data.Dataset.from_generator(
    lambda: generate_single_sample(nw),
    output_signature=(
        tf.TensorSpec(shape=(nw,), dtype=tf.float32),
        tf.TensorSpec(shape=(3,), dtype=tf.float32)
    )
).batch(batch_size).prefetch(tf.data.AUTOTUNE)

dataset_val = tf.data.Dataset.from_generator(
    lambda: generate_single_sample(nw),
    output_signature=(
        tf.TensorSpec(shape=(nw,), dtype=tf.float32),
        tf.TensorSpec(shape=(3,), dtype=tf.float32)
    )
).batch(batch_size).prefetch(tf.data.AUTOTUNE)

注意：如果使用Scikit-learn的预处理类（如StandardScaler），建议替换为tf.keras.layers.Normalization，将预处理逻辑整合到模型中，进一步提升端到端效率。

二、训练重启后变慢的原因与解决方法 #

1. Generator状态残留：Python Generator是有状态的，中断训练后若未重新创建Generator实例，可能导致内部变量累积或生成逻辑异常。解决方法：每次训练前重新创建Dataset实例，或使用无状态的无限循环Generator（如上述代码中的while True）。
2. 内存泄漏：频繁创建NumPy数组未及时回收会导致内存占用攀升，拖慢训练。用TensorFlow原生操作生成数据可避免此问题，TF会自动管理内存。
3. TF会话/图残留：中断训练后，TF可能残留旧的计算图或缓存节点，导致后续训练效率下降。解决方法：每次训练前调用tf.keras.backend.clear_session()清理会话，重置所有状态。
4. 硬件资源未释放：中断训练后GPU/CPU资源可能未完全释放，导致后续训练资源不足。可重启Python环境，或用nvidia-smi（NVIDIA GPU）查看进程并手动释放。

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