TensorFlow 2.4.1中tf.keras.utils.Sequence动态调整batch_size触发数据耗尽警告问题

TF2.x（包括2.4.1）在训练启动阶段，会一次性把自定义Sequence的__len__()返回值缓存下来，作为整个训练周期的固定steps_per_epoch参数，不会在每个epoch结束后重新计算、读取__len__()的最新结果。
这种设计带来的直接问题：

• 你在on_epoch_end里修改self.batch_size时，Sequence自身的__len__()返回值其实已经变了，但训练循环完全不感知，始终用最开始缓存的步数调度取数
• 当调整后的batch_size比初始值大时，每步实际取的样本量比调度预设的多，按初始步数算下来总取数规模会超过全部样本量，很快就把数据集读空，触发数据耗尽警告直接终止训练

这也是为什么旧版Keras1没这个问题：Keras1的训练循环每个epoch都会重新调用len(generator)取最新的步数，TF2.x为了减少重复计算做了步数缓存优化，反而给动态调整的场景埋了坑。

核心改动：手动同步动态步数到训练循环 #

写一个极简的自定义Callback，在每个epoch开始前，把Sequence当前的最新长度手动赋值给训练循环的步数参数，覆盖框架之前缓存的旧值。代码实现：

import tensorflow as tf
class DynamicStepSyncCallback(tf.keras.callbacks.Callback):
    def __init__(self, data_sequence):
        super().__init__()
        self.data_sequence = data_sequence
    
    def on_epoch_begin(self, epoch, logs=None):
        # 用Sequence的实时长度替换训练循环缓存的旧步数
        self.params["steps"] = len(self.data_sequence)

优化Sequence类的边界健壮性 #

原示例代码直接对原始路径列表做切片取数，没有处理动态batch_size下的索引偏移、最后一个批次的越界问题，很容易在batch_size变化后出现样本错乱。优化后的完整Sequence实现：

from skimage.io import imread
from skimage.transform import resize
import numpy as np
import random
import math

class CIFAR10Sequence(tf.keras.utils.Sequence):
    def __init__(self, x_set, y_set, batch_size):
        self.x, self.y = x_set, y_set
        self.batch_size = batch_size
        # 维护内部样本索引数组，从根源上避免动态batch_size带来的索引错乱
        self.indices = np.arange(len(self.x))

    def __len__(self):
        return math.ceil(len(self.x) / self.batch_size)

    def on_epoch_end(self):
        # 每个epoch结束后打乱样本顺序、随机更新batch_size
        np.random.shuffle(self.indices)
        self.batch_size = random.randint(10, 100)
        print(f"当前epoch结束，更新batch_size为{self.batch_size}，下轮训练步数为{len(self)}")

    def __getitem__(self, idx):
        start_pos = idx * self.batch_size
        # 做边界截断，避免最后一个批次的索引超出样本总长度
        end_pos = min((idx + 1) * self.batch_size, len(self.indices))
        batch_indices = self.indices[start_pos:end_pos]
        batch_x = [self.x[i] for i in batch_indices]
        batch_y = [self.y[i] for i in batch_indices]

        return np.array([
            resize(imread(file_path), (200, 200)) for file_path in batch_x
        ]), np.array(batch_y)

训练时的调用方式 #

按常规流程实例化Sequence和自定义Callback，传入model.fit()即可，不需要手动指定steps_per_epoch参数：

# 实例化数据集序列
train_seq = CIFAR10Sequence(x_train_paths, y_train_labels, batch_size=16)
# 实例化步数同步回调
sync_step_callback = DynamicStepSyncCallback(train_seq)

model.fit(
    train_seq,
    epochs=20,
    callbacks=[sync_step_callback]
)

配置完成后：

• 每个epoch开始前，训练循环都会读取Sequence的实时长度，完全匹配当前动态batch_size对应的实际步数
• 每个epoch结束后自动更新batch_size、打乱样本顺序，不会再出现步数固定、数据读空触发中断的问题
• 即使总样本量不多，最后一个批次会自动做边界截断，不会出现索引越界的报错

内容的提问来源于stack exchange，提问作者Võ Hoàng Trọng