TensorFlow自定义掩码BCE损失函数训练失效问题排查

你的自定义损失实现存在3个核心错误，直接导致模型无法更新参数、完全学不到有效模式：

• 错误1：错误对预测值y_pred乘掩码。掩码的作用是把无效样本算出来的损失清0，不是直接改模型输出的预测值。你现在的写法虽然能让无效样本位置算出来的BCE值为0，但会破坏梯度计算的稳定性，一旦BCE的计算参数（比如是否用logits）调整，很容易出现异常梯度。
• 错误2：损失归一化逻辑错误。你直接返回逐样本的损失值后，Keras默认会把batch里所有样本（包括你要屏蔽的无效样本）拉通算平均。只要batch里有无效样本，有效样本产生的损失梯度就会被无效样本的0值拉低，等效学习率会变得极小，模型根本没法有效更新权重。
• 错误3：对哨兵值999的处理鲁棒性极差。你现在仅靠乘法把无效位置的y_true转成0，一旦后续调整损失计算逻辑，999这个异常值直接进入BCE计算会产生爆炸大的损失值，直接把模型训崩。

按照掩码损失的标准逻辑重写即可，核心流程是：先标记有效样本位置，清洗掉标签里的999异常值，逐样本算完BCE之后再用掩码清掉无效样本的损失，最后只对有效样本算平均损失，避免梯度被稀释。
修复后的代码如下：

def custom_loss_mask(y_true, y_pred):
    # 对齐标签和预测值的数据精度
    y_true = tf.cast(y_true, dtype=y_pred.dtype)
    # 生成掩码：有效样本位置为1，y_true=999的无效位置为0
    mask = tf.cast(tf.not_equal(y_true, 999), dtype=y_pred.dtype)
    # 把无效位置的999替换成0，避免异常值进入BCE计算
    y_true_clean = tf.where(mask == 1, y_true, tf.zeros_like(y_true))
    # 逐样本计算BCE，绝对不要修改原始y_pred
    # 注意：如果你的模型最后一层没有加sigmoid激活、直接输出原始logits，必须把from_logits改成True
    bce_per_sample = tf.keras.losses.binary_crossentropy(
        y_true_clean, 
        y_pred,
        from_logits=False
    )
    # 清掉无效样本的损失
    masked_loss = bce_per_sample * mask
    # 按有效样本总数做归一化，加1e-8防止出现全无效样本时除零报错
    loss = tf.reduce_sum(masked_loss) / (tf.reduce_sum(mask) + 1e-8)
    return loss

• 你现在用的内置accuracy、AUC指标默认不会跳过y_true=999的无效样本，算出来的指标本身就是错的。更简单的实现屏蔽逻辑的方式是不用写自定义损失：训练时把y_true!=999转成0/1权重作为sample_weight参数传给fit_generator，直接用Keras内置的二分类交叉熵损失就行，框架会自动用权重屏蔽无效样本的损失和指标计算，不容易出错。

传入sample_weight的方案兼容性最好，不需要单独修改指标计算逻辑，适合快速验证。

• 务必检查模型最后一层的激活函数：如果最后一层没有sigmoid，损失函数里一定要把from_logits设为True，否则BCE计算会出现数值不稳定，模型训不动。
• 从你贴的训练日志看，模型已经完全塌缩到固定输出了（AUC恒为0.5、损失和准确率全程不动），修复损失后必须重新初始化模型权重从头训练，在当前已经塌缩的权重上继续训没有意义。

内容的提问来源于stack exchange，提问作者Athreya H P