TensorFlow中如何遍历训练误差并实现自定义损失函数？

实现思路 #

无需逐个遍历批次样本，利用TensorFlow的向量化操作即可高效完成需求：先计算预测值与真实值的误差，通过掩码分别筛选出误差大于0.5和小于0.5的样本子集，再分别计算两个子集的误差均值，最终返回两者之和作为损失。

代码实现 #

import tensorflow as tf

def custom_loss(y_true, y_pred):
    # 计算每个样本的误差
    error = y_pred - y_true
    
    # 处理误差大于0.5的部分：生成掩码并计算均值
    mask_above = tf.cast(error > 0.5, tf.float32)
    sum_above = tf.reduce_sum(error * mask_above)
    count_above = tf.reduce_sum(mask_above)
    mean_above = sum_above / tf.maximum(count_above, 1e-8)  # 避免分母为0
    
    # 处理误差小于0.5的部分：生成掩码并计算均值
    mask_below = tf.cast(error < 0.5, tf.float32)
    sum_below = tf.reduce_sum(error * mask_below)
    count_below = tf.reduce_sum(mask_below)
    mean_below = sum_below / tf.maximum(count_below, 1e-8)
    
    # 返回两个均值之和作为最终损失
    return mean_above + mean_below

关键细节说明 #

• 掩码处理：用tf.cast将布尔型筛选结果转为浮点型，方便后续计算加权误差和。
• 避免NaN：用tf.maximum将分母最小值限制为1e-8，防止批次中无符合条件样本时出现除以0的错误。
• 向量化优势：相比循环遍历，向量化操作能充分利用TensorFlow的GPU加速能力，提升训练效率。

使用方式 #

在模型编译时直接指定该损失函数即可：

# 假设已定义好模型结构
model.compile(optimizer='adam', loss=custom_loss)

内容的提问来源于stack exchange，提问作者张铭舟