Keras训练第二轮及后续epoch出现训练数据耗尽问题求助

问题根源 #

手动设置steps_per_epoch和validation_steps为样本数整除批次大小的结果，导致生成器在第一个epoch结束后，迭代指针停留在数据末尾；且新版Keras中flow_from_directory生成器默认不会自动重置指针，后续epoch无法获取新数据，触发数据耗尽警告。

修复方案 #

方案1：移除手动步数设置（推荐） #

直接删除steps_per_epoch和validation_steps参数，Keras会自动根据生成器的样本量与批次大小计算每个epoch的步数，并且自动在每个epoch结束后重置生成器的迭代状态。

修改后的训练代码：

history = model.fit(
    train_generator,
    validation_data=validation_generator,
    epochs=epochs
)

方案2：手动重置生成器（保留步数控制时使用） #

如果需要手动指定步数，添加回调函数在每个epoch结束后重置生成器：

1. 定义回调类：

class ResetGenerator(tf.keras.callbacks.Callback):
    def on_epoch_end(self, epoch, logs=None):
        train_generator.reset()
        validation_generator.reset()

2. 在训练时引入回调：

history = model.fit(
    train_generator,
    steps_per_epoch=steps,
    validation_data=validation_generator,
    validation_steps=steps_v,
    epochs=epochs,
    callbacks=[ResetGenerator()]
)

额外优化建议 #

• 若想确保所有样本都被训练到，将步数计算改为向上取整，避免丢弃末尾不足一个批次的样本：

steps = (train_generator.samples + batch_size - 1) // batch_size
steps_v = (validation_generator.samples + batch_size - 1) // batch_size

• 确认生成器shuffle参数为True（默认开启），保证每个epoch的数据顺序随机，降低过拟合风险。

修复后重新训练，即可消除数据耗尽警告，每个epoch都能完整执行训练与验证流程，便于观察loss和accuracy变化判断模型是否过拟合。

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