如何用Keras.backend.function获取训练批次预测及原数据行索引？

1. 构建包含原始行索引的数据集 #

首先要确保训练数据集加载时，把原始行索引和特征、标签一起打包，让每个批次都能输出索引信息：

import numpy as np
import tensorflow as tf

# 假设已有训练数据x_train, y_train
x_train = ...  # 你的特征数据
y_train = ...  # 你的标签数据
idx_train = np.arange(len(x_train))  # 生成原始行索引数组

# 构建包含索引的数据集：每个元素格式为(特征, (标签, 索引))
dataset = tf.data.Dataset.from_tensor_slices((x_train, (y_train, idx_train)))
dataset = dataset.shuffle(1000).batch(32)  # 按需设置shuffle和batch大小

2. 自定义回调类获取实际数据 #

因为你已经设置run_eagerly=True，回调中可以直接拿到numpy格式的真实数据，无需使用Keras.backend.function()。在on_train_batch_end方法中，通过模型接口获取批次输入、索引和预测概率：

class BatchDataCallback(tf.keras.callbacks.Callback):
    def on_train_batch_end(self, batch, logs=None):
        # 获取当前批次的特征输入
        current_batch_inputs = self.model.train_data_adapter.get_data()[0]
        # 从目标数据中拆分出标签和原始行索引
        _, (_, current_batch_indices) = self.model.train_data_adapter.get_data()
        
        # 获取当前批次的概率预测值（numpy数组格式）
        pred_probs = self.model.predict_on_batch(current_batch_inputs)
        
        # 这里可直接用pred_probs和current_batch_indices计算自定义指标
        # 示例：打印部分数据验证
        print(f"批次 {batch} 部分索引: {current_batch_indices[:5].tolist()}")
        print(f"批次 {batch} 第一个样本预测概率: {pred_probs[0].tolist()}")
        
        # 插入你的自定义指标计算逻辑
        # ...

3. 训练时加载回调 #

在模型fit阶段传入自定义回调，同时保持run_eagerly=True：

model = tf.keras.Sequential([
    # 你的模型层定义
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
model.fit(dataset, epochs=5, callbacks=[BatchDataCallback()], run_eagerly=True)

关键说明 #

• 放弃Keras.backend.function()的原因：run_eagerly=True模式下，模型所有操作即时执行，直接通过predict_on_batch和train_data_adapter.get_data()就能获取真实numpy数组，无需构建图模式函数。
• 数据集结构可灵活调整：如果标签是多维度格式，只需保证索引被包含在批次输出中即可，不用局限于示例中的打包方式。

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