未知batch size时如何对tf.keras.Model输出的每个批样本应用函数

你不需要提前获取当前batch size的取值，直接使用TensorFlow内置的tf.map_fn接口即可自动适配任意batch size，对batch内的每个样本单独应用自定义处理逻辑，两种实现方案如下： #

方案1：将处理逻辑嵌入模型定义，调用predict_on_batch直接返回处理后结果 #

直接在构建tf.keras.Model时加入样本处理逻辑，不需要额外修改后续调用代码：

import tensorflow as tf

# 替换为你自己的单样本处理逻辑，入参为单个样本的输出张量，形状为(output_shape,)
def custom_process(single_sample_output):
    # 示例：对单样本输出做L2归一化，可替换为任意自定义运算
    return tf.math.l2_normalize(single_sample_output, axis=-1)

# 你的原模型构建逻辑
input_layer = tf.keras.Input(shape=(your_input_shape,))
# 替换为你自己的模型层堆叠
dense_layer = tf.keras.layers.Dense(units=your_output_shape)(input_layer)
# 对batch维度遍历，每个样本应用自定义处理函数
processed_output = tf.map_fn(
    fn=custom_process,
    elems=dense_layer,
    # 与custom_process返回的单样本张量形状、数据类型保持一致
    fn_output_signature=tf.TensorSpec(shape=(your_output_shape,), dtype=tf.float32)
)

model = tf.keras.Model(inputs=input_layer, outputs=processed_output)

完成模型构建后直接调用model.predict_on_batch(输入数据)，返回的结果就是每个样本处理后的batch输出。 #

方案2：在predict_on_batch调用后做后处理 #

如果不需要修改原有模型定义，可在拿到原始batch输出后再做样本级处理：

# 调用原有模型拿到原始batch输出，形状为(batch_size, output_shape)
raw_batch_output = original_model.predict_on_batch(your_input)
# 对每个样本应用自定义处理函数
processed_batch = tf.map_fn(custom_process, raw_batch_output).numpy()

注意事项 #

• 若你的自定义处理逻辑全部使用TensorFlow内置运算实现，tf.map_fn会自动纳入计算图优化，不会有明显性能损耗
• 若你的自定义逻辑包含Numpy等非TensorFlow运算，需要用tf.py_function包装后再传入tf.map_fn，示例如下：

# 非TensorFlow实现的自定义处理函数
def numpy_custom_process(single_sample):
    return single_sample * 3 + 2

def wrapped_process(single_sample):
    return tf.py_function(
        func=numpy_custom_process,
        inp=[single_sample],
        Tout=tf.float32
    )

processed_output = tf.map_fn(wrapped_process, raw_batch_output)

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