TensorFlow训练报错：输入数据耗尽，如何解决fit_generator中断问题？

这个报错本质很明确：你设置的训练总批次（steps_per_epoch * epochs = 50*10=5000）超过了生成器能提供的批次量，导致训练到一半就没数据可用了。下面给你几个实用的解决思路：

方法1：给生成器添加repeat()循环生成数据 #

这是最直接的解决方案，让生成器在数据耗尽后自动循环从头输出，刚好能满足你需要的5000个批次需求。

如果你的生成器是用ImageDataGenerator的flow_from_directory创建的，只需要在末尾链式调用.repeat()即可：

train_generator = train_datagen.flow_from_directory(
    train_dir,
    target_size=(150, 150),
    batch_size=32,
    class_mode='binary'
).repeat()  # 加上这一行，让生成器循环输出数据

# 验证生成器也建议同步处理，避免验证阶段触发同样报错
validation_generator = validation_datagen.flow_from_directory(
    validation_dir,
    target_size=(150, 150),
    batch_size=32,
    class_mode='binary'
).repeat()

之后你的fit_generator调用可以保持不变，生成器会持续循环提供数据直到完成所有训练批次。

方法2：调整steps_per_epoch到实际批次数量 #

如果你不想重复使用训练数据（比如希望每个epoch只完整遍历一次数据集），那就要确保steps_per_epoch的数值等于训练集总样本数 ÷ 批次大小。

比如你的训练集有1600个样本、批次大小是32，那实际批次量是1600 // 32 = 50，刚好匹配你现在的设置——这时候要检查是不是生成器没正确加载所有样本（比如路径错误、文件夹内样本数统计失误）。如果样本数不是批次大小的整数倍，可以用math.ceil向上取整：

import math

train_samples = 1650  # 假设训练集实际有1650个样本
batch_size = 32
steps_per_epoch = math.ceil(train_samples / batch_size)  # 1650/32≈51.56，向上取整为52

# 验证集同理调整
val_samples = 400
validation_steps = math.ceil(val_samples / batch_size)

history = model.fit_generator(
    train_generator,
    steps_per_epoch=steps_per_epoch,
    epochs=10,
    verbose=1,
    validation_data=validation_generator,
    validation_steps=validation_steps
)

这样每个epoch刚好遍历一次数据集，不会出现数据耗尽的问题。

方法3：迁移到tf.data.Dataset（推荐TensorFlow 2.x用户） #

如果你用的是TensorFlow 2.x，其实fit_generator已经被官方弃用了，更推荐用tf.data.Dataset构建数据集，只需要调用.repeat()就能实现数据循环：

import tensorflow as tf

# 假设你已经实现了加载单张图片的函数load_image
train_dataset = tf.data.Dataset.list_files(train_dir + '/*/*')
train_dataset = train_dataset.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
train_dataset = train_dataset.batch(32).repeat()  # 添加repeat()实现循环

val_dataset = tf.data.Dataset.list_files(validation_dir + '/*/*')
val_dataset = val_dataset.map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
val_dataset = val_dataset.batch(32).repeat()

# 用model.fit替代fit_generator，兼容性更好
history = model.fit(
    train_dataset,
    steps_per_epoch=50,
    epochs=10,
    verbose=1,
    validation_data=val_dataset,
    validation_steps=50
)

额外注意点 #

• 如果选择重复数据的方式，要留意模型是否会过拟合——不过训练10个epoch的情况下，只要数据集足够大，一般问题不大；
• 验证阶段如果也触发同样报错，记得给validation_generator也加上.repeat()，或者调整validation_steps到验证集的实际批次量；
• TensorFlow 2.x中model.fit()完全支持生成器和tf.data.Dataset，建议替换掉fit_generator，避免后续版本兼容性问题。

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