Keras中如何实现批次内损失的自定义归约操作

自定义归约方法的实现方式 #

• 方式1：基于Keras损失基类的reduction参数配置
  Keras的tf.keras.losses.Loss基类原生支持归约策略配置，可选内置值包括tf.keras.losses.Reduction.NONE（不做归约，返回样本级损失向量）、tf.keras.losses.Reduction.SUM（批次损失求和）、tf.keras.losses.Reduction.AUTO（默认策略，对批次损失求平均）。如果需要自定义归约逻辑，可继承损失基类重写方法实现：

import tensorflow as tf
class CustomLoss(tf.keras.losses.Loss):
    def __init__(self, reduction=tf.keras.losses.Reduction.AUTO, name="custom_loss"):
        super().__init__(reduction=reduction, name=name)
        # 可在此处传入自定义损失超参数
    def call(self, y_true, y_pred):
        # 此处返回形状为(batch_size,)的样本级损失向量
        per_sample_loss = tf.reduce_mean(tf.square(y_true - y_pred), axis=-1)
        return per_sample_loss
# 使用时指定归约策略，比如按批次求和
loss_fn = CustomLoss(reduction=tf.keras.losses.Reduction.SUM)

如果内置归约策略不满足需求，也可以直接在call方法中实现自定义归约逻辑，比如按样本权重加权、忽略特定样本损失等。

• 方式2：自定义损失函数内直接完成归约
  该做法完全可行，适合归约逻辑固定、不需要动态调整的场景，写法更简单：

def custom_loss(y_true, y_pred):
    per_sample_loss = tf.reduce_mean(tf.square(y_true - y_pred), axis=-1)
    # 自定义归约示例：仅对损失大于0.1的样本求平均
    mask = per_sample_loss > 0.1
    reduced_loss = tf.reduce_mean(per_sample_loss[mask])
    return reduced_loss

注意：如果在损失函数内部完成归约返回单个标量，模型编译时指定的损失归约配置会自动失效，Keras会直接使用返回的标量作为批次损失值。

两种方案的适用场景 #

• 需要复用损失函数、动态切换归约策略时，推荐继承损失基类使用reduction参数配置
• 归约逻辑固定不需要调整时，直接在函数内部完成归约的写法更简洁高效

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