TensorFlow多分类自定义损失函数出现无梯度错误如何解决

无梯度的核心原因是你改造多分类召回率/特异度时用到了argmax、布尔判断、四舍五入这类不可导操作，哪怕转了float32，这些操作的输出节点本身没有梯度传递链路，反向传播时自然找不到可更新的梯度。

• 替换硬分类操作为可导近似：不要用tf.argmax取预测类别，改用tf.nn.softmax输出的概率值直接计算软混淆矩阵。比如对三分类的第i类，真阳性的软计算方式是tf.reduce_sum(y_true[:,i] * y_pred[:,i])，假阳性是tf.reduce_sum((1 - y_true[:,i]) * y_pred[:,i])，假阴性是tf.reduce_sum(y_true[:,i] * (1 - y_pred[:,i]))，全程用概率计算没有截断，梯度链路完整。
• 移除阈值判断类操作：如果你的原实现里加了大于0.5这类硬阈值判断，全部替换为连续的概率加权，不要做0/1硬截断。
• 确认损失输出维度适配：三分类场景下你可以逐类计算加权损失后求和，不要直接返回形状为(3,)的张量，要返回标量损失值，或者匹配y_true的形状避免广播错误导致梯度丢失。
• 检查标签格式匹配：如果用ImageDataGenerator的class_mode='categorical'输出的是独热编码标签，不要做多余的tf.cast到int类型的操作，全程保持float32参与计算。

可直接复用的损失函数示例 #

import tensorflow as tf

def weighted_recall_specificity_loss(y_true, y_pred, fn_penalty=[2.0, 3.0, 1.5], fp_penalty=[1.0, 1.0, 2.0]):
    # y_true: 独热编码标签 (batch_size, 3)
    # y_pred: 模型输出的logits，自动做softmax转概率
    y_pred = tf.nn.softmax(tf.cast(y_pred, tf.float32), axis=-1)
    y_true = tf.cast(y_true, tf.float32)
    
    total_loss = 0.0
    for class_idx in range(3):
        # 软混淆矩阵计算，全程可导
        tp = tf.reduce_sum(y_true[:, class_idx] * y_pred[:, class_idx])
        fn = tf.reduce_sum(y_true[:, class_idx] * (1 - y_pred[:, class_idx]))
        fp = tf.reduce_sum((1 - y_true[:, class_idx]) * y_pred[:, class_idx])
        tn = tf.reduce_sum((1 - y_true[:, class_idx]) * (1 - y_pred[:, class_idx]))
        
        # 平滑项避免除0
        recall = tp / (tp + fn + 1e-7)
        specificity = tn / (tn + fp + 1e-7)
        
        # 按自定义权重加惩罚项
        class_loss = fn_penalty[class_idx] * (1 - recall) + fp_penalty[class_idx] * (1 - specificity)
        total_loss += class_loss
    return total_loss

如果你要调整假阳性、假阴性的差异化惩罚力度，直接修改fn_penalty、fp_penalty数组对应类别的系数即可，不需要修改核心计算逻辑。

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