如何在批量结束触发的回调中获取指定批次的输入张量？

这个问题我之前在项目里也碰到过——默认的Keras/TensorFlow回调确实没法直接拿到批次输入，得动点小手脚才行。毕竟每轮epoch都会打乱数据，批次编号本身没法直接映射到原始数据，必须把epoch+批次编号和对应的输入绑定起来存储。下面给你几个实用的方案：

方案1：自定义记录器+回调追踪（适配Keras fit()） #

核心思路是先搞一个“记录器”类来存储每个epoch下各批次的输入，再用回调同步当前的epoch和批次编号，最后在批量结束的回调里从记录器中取数据。

步骤1：写一个批次输入记录器 #

这个类负责把每个批次的输入按epoch和批次号存起来：

class BatchInputRecorder:
    def __init__(self):
        self.epoch_batch_inputs = {}  # 结构: {epoch_num: {batch_num: 输入数据}}

    def save_batch(self, epoch, batch_num, input_tensor):
        # 先检查当前epoch的存储键是否存在，不存在就创建
        if epoch not in self.epoch_batch_inputs:
            self.epoch_batch_inputs[epoch] = {}
        # 把Tensor转成numpy数组存储（也可以直接存Tensor，看你需求）
        self.epoch_batch_inputs[epoch][batch_num] = input_tensor.numpy()

步骤2：给模型加一个记录输入的层 #

因为Keras的回调拿不到输入数据，我们可以在模型的输入层后面加一个自定义层，每次前向传播时自动记录输入：

import tensorflow as tf

class InputCaptureLayer(tf.keras.layers.Layer):
    def __init__(self, recorder, **kwargs):
        super().__init__(**kwargs)
        self.recorder = recorder
        self.current_epoch = 0
        self.current_batch = 0

    def call(self, inputs):
        # 每次模型处理输入时，把当前批次的输入存到记录器
        self.recorder.save_batch(self.current_epoch, self.current_batch, inputs)
        return inputs  # 不改变原始输入，只做记录

步骤3：用回调同步epoch和批次编号 #

写一个回调来更新上面自定义层里的当前epoch和批次号：

class EpochBatchTracker(tf.keras.callbacks.Callback):
    def __init__(self, capture_layer):
        super().__init__()
        self.capture_layer = capture_layer

    def on_epoch_begin(self, epoch, logs=None):
        # 每轮开始时重置批次编号，更新当前epoch
        self.capture_layer.current_epoch = epoch
        self.capture_layer.current_batch = 0

    def on_batch_end(self, batch, logs=None):
        # 每批结束后，批次编号+1
        self.capture_layer.current_batch += 1

步骤4：在你的批量结束回调中取数据 #

现在你可以写自己的批量结束回调，从记录器里拿到对应批次的输入：

class MyBatchEndCallback(tf.keras.callbacks.Callback):
    def __init__(self, recorder):
        super().__init__()
        self.recorder = recorder

    def on_batch_end(self, batch, logs=None):
        # 拿到当前epoch（可以从回调的self.model里取，或者从追踪器里拿）
        current_epoch = self.model.layers[1].current_epoch  # 假设InputCaptureLayer是第二层
        # 从记录器中取出当前批次的输入数据
        batch_input = self.recorder.epoch_batch_inputs[current_epoch][batch]
        # 这里就可以用batch_input做你想做的事了
        print(f"当前批次输入形状: {batch_input.shape}")

组装起来训练 #

# 初始化记录器和捕获层
recorder = BatchInputRecorder()
capture_layer = InputCaptureLayer(recorder)

# 构建你的模型（示例）
input_layer = tf.keras.layers.Input(shape=(28,28,1))
x = capture_layer(input_layer)
x = tf.keras.layers.Flatten()(x)
x = tf.keras.layers.Dense(128, activation='relu')(x)
output_layer = tf.keras.layers.Dense(10, activation='softmax')(x)
model = tf.keras.Model(inputs=input_layer, outputs=output_layer)

# 编译模型
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 准备回调列表
callbacks = [
    EpochBatchTracker(capture_layer),
    MyBatchEndCallback(recorder)
]

# 开始训练
model.fit(train_dataset, epochs=5, callbacks=callbacks)

方案2：自定义训练循环（更灵活） #

如果你不想折腾Keras的回调和自定义层，直接用TensorFlow的自定义训练循环会更直观——你可以完全掌控每个epoch和批次的处理流程，批量结束时直接拿到输入数据。

示例代码：

import tensorflow as tf

# 初始化记录器
recorder = BatchInputRecorder()

# 假设你已经有了model、train_dataset、loss_fn、optimizer
epochs = 5

for epoch in range(epochs):
    print(f"Epoch {epoch+1}/{epochs}")
    batch_num = 0
    # 遍历当前epoch的所有批次
    for x_batch, y_batch in train_dataset:
        # 标准的训练步骤
        with tf.GradientTape() as tape:
            predictions = model(x_batch, training=True)
            loss = loss_fn(y_batch, predictions)
        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))

        # 记录当前批次的输入
        recorder.save_batch(epoch, batch_num, x_batch)

        # 直接在这里执行你的批量结束逻辑
        print(f"第{epoch}轮第{batch_num}批输入形状: {x_batch.shape}")
        
        batch_num += 1

注意事项 #

• 如果你处理的是超大数据集，别长期存储所有批次的输入——用完对应epoch的数据后，可以手动删除recorder.epoch_batch_inputs[epoch]，避免内存溢出。
• 如果用的是PyTorch，思路是一样的：要么用钩子（Hook）记录输入，要么在训练循环里直接处理，核心都是把epoch+批次号和输入绑定存储。

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