修复自定义BCE+Dice损失函数中的数值不稳定问题

问题根源 #

组合损失出现数值不稳定，主要有三个核心原因：

• 二元交叉熵（BCE）计算未适配logits场景：若模型输出是未经过Sigmoid的原始logits，直接计算BCE会因log(0)或log(1)触发极端值，导致梯度爆炸或NaN
• 重复的张量变形与类型转换：冗余操作可能引入维度隐式错误，放大数值波动
• 两个损失尺度未对齐：BCE默认返回逐元素损失，直接和标量Dice损失相加会导致训练梯度失衡

修复后的代码 #

import tensorflow as tf

def dice_coeff(y_true, y_pred):
    smooth = 1.
    # 用tf.flatten简化张量扁平化操作
    y_pred_f = tf.cast(tf.flatten(y_pred), tf.float32)
    y_true_f = tf.cast(tf.flatten(y_true), tf.float32)
    
    intersection = tf.reduce_sum(y_true_f * y_pred_f)
    denominator = tf.reduce_sum(y_true_f) + tf.reduce_sum(y_pred_f)
    score = (2. * intersection + smooth) / (denominator + smooth)
    return score

def dice_loss(y_true, y_pred):
    return 1. - dice_coeff(y_true, y_pred)

def bce_dice_loss(y_true, y_pred, from_logits=False, bce_weight=1.0, dice_weight=1.0):
    # 统一扁平化处理，避免重复操作
    y_true_flat = tf.cast(tf.flatten(y_true), tf.float32)
    y_pred_flat = tf.cast(tf.flatten(y_pred), tf.float32)
    
    # 处理BCE数值稳定性：根据模型输出类型设置from_logits
    bce_loss = tf.keras.losses.binary_crossentropy(
        y_true_flat, 
        y_pred_flat, 
        from_logits=from_logits
    )
    # 对BCE取均值，转为标量后和Dice损失匹配
    bce_loss = tf.reduce_mean(bce_loss)
    
    dice_loss_val = dice_loss(y_true, y_pred)
    
    # 加权组合，灵活平衡两个损失的贡献
    return bce_weight * bce_loss + dice_weight * dice_loss_val

关键修改说明 #

1. BCE数值稳定性处理：如果模型最后一层没有添加Sigmoid激活，调用损失函数时必须设置from_logits=True，TensorFlow会自动使用数值稳定的方式计算BCE，避免极端值问题；若已加Sigmoid，则设为False
2. 统一张量操作：用tf.flatten替代手动reshape，减少冗余操作，避免维度错误
3. 损失尺度对齐：对BCE损失取均值，确保和Dice损失同为标量，再通过权重参数调整两者贡献，稳定训练梯度

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