Keras模型训练中如何获取批次数据以计算自定义指标？

方法1：使用自定义训练循环（最推荐） #

放弃依赖model.fit()，改用自定义训练循环可以完全掌控每个批次的数据，直接获取训练时使用的批次输入和标签，步骤清晰可控：

1. 准备训练数据集（如tf.data.Dataset）
2. 遍历数据集的每个批次
3. 对当前批次执行训练步骤
4. 用训练后的模型对该批次做预测
5. 计算自定义指标

示例代码：

import tensorflow as tf
import numpy as np

# 假设模型、损失函数、优化器已定义
model = ...
loss_fn = tf.keras.losses.CategoricalCrossentropy()
optimizer = tf.keras.optimizers.Adam()

# 准备训练数据集
train_dataset = tf.data.Dataset.from_tensor_slices((X_train, y_train)).batch(batch_size)

# 存储自定义指标历史
custom_metric_history = []

for epoch in range(num_epochs):
    print(f"Epoch {epoch+1}/{num_epochs}")
    for batch_idx, (x_batch, y_batch) in enumerate(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))
        
        # 训练后对当前批次预测
        post_train_predictions = model.predict_on_batch(x_batch)
        
        # 替换为你的自定义指标计算逻辑
        custom_metric = your_custom_metric(y_batch, post_train_predictions)
        custom_metric_history.append(custom_metric)
        
        print(f"Batch {batch_idx+1}, Custom Metric: {custom_metric:.4f}")

方法2：自定义数据生成器+回调 #

如果坚持使用model.fit()，可以通过自定义生成器跟踪最后输出的批次，再在回调中获取该批次数据：

步骤1：创建带批次跟踪的生成器 #

class TrackingGenerator(tf.keras.utils.Sequence):
    def __init__(self, X, y, batch_size):
        self.X = X
        self.y = y
        self.batch_size = batch_size
        self.indices = np.arange(len(X))
        # 存储最后一个批次的数据
        self.last_x_batch = None
        self.last_y_batch = None

    def __len__(self):
        return int(np.ceil(len(self.X) / self.batch_size))

    def __getitem__(self, idx):
        batch_indices = self.indices[idx*self.batch_size : (idx+1)*self.batch_size]
        self.last_x_batch = self.X[batch_indices]
        self.last_y_batch = self.y[batch_indices]
        return self.last_x_batch, self.last_y_batch

    def on_epoch_end(self):
        # 每个epoch结束后打乱数据（可选）
        np.random.shuffle(self.indices)

步骤2：修改回调获取批次数据 #

class CollectCallback(tf.keras.callbacks.Callback):
    def __init__(self, generator):
        super().__init__()
        self.generator = generator

    def on_batch_end(self, batch, logs=None):
        # 获取当前训练批次的原始数据
        x_batch = self.generator.last_x_batch
        y_batch = self.generator.last_y_batch
        
        # 训练后对该批次预测
        post_train_predictions = self.model.predict_on_batch(x_batch)
        
        # 计算自定义指标
        custom_metric = your_custom_metric(y_batch, post_train_predictions)
        print(f"Batch {batch+1}, Custom Metric: {custom_metric:.4f}")

步骤3：启动训练 #

train_generator = TrackingGenerator(X_train, y_train, batch_size=32)
callback = CollectCallback(train_generator)

model.fit(
    train_generator,
    epochs=num_epochs,
    callbacks=[callback]
)

方法3：基于索引计算批次（仅适用于非 shuffle 或固定 shuffle 种子场景） #

如果训练数据是numpy数组，且不使用shuffle（或固定shuffle种子并提前记录索引顺序），可以通过批次号计算当前批次的索引范围：

class CollectCallback(tf.keras.callbacks.Callback):
    def __init__(self, X, y, batch_size, shuffle_seed=None):
        super().__init__()
        self.X = X
        self.y = y
        self.batch_size = batch_size
        self.current_epoch_indices = None
        self.shuffle_seed = shuffle_seed

    def on_epoch_begin(self, epoch, logs=None):
        # 每个epoch生成索引（如需shuffle则固定种子保证一致性）
        self.current_epoch_indices = np.arange(len(self.X))
        if self.shuffle_seed is not None:
            np.random.seed(self.shuffle_seed + epoch)
            np.random.shuffle(self.current_epoch_indices)

    def on_batch_end(self, batch, logs=None):
        # 计算当前批次的索引范围
        start_idx = batch * self.batch_size
        end_idx = min((batch+1)*self.batch_size, len(self.X))
        batch_indices = self.current_epoch_indices[start_idx:end_idx]
        
        x_batch = self.X[batch_indices]
        y_batch = self.y[batch_indices]
        
        post_train_predictions = self.model.predict_on_batch(x_batch)
        custom_metric = your_custom_metric(y_batch, post_train_predictions)
        print(f"Batch {batch+1}, Custom Metric: {custom_metric:.4f}")

使用示例：

callback = CollectCallback(X_train, y_train, batch_size=32, shuffle_seed=42)
model.fit(X_train, y_train, batch_size=32, epochs=num_epochs, shuffle=True, callbacks=[callback])

注意：该方法依赖固定的shuffle顺序，若Keras内部shuffle逻辑变更（如版本更新）可能导致批次不匹配，可靠性不如前两种方法。

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