TensorFlow中基于神经网络部分输出定义损失函数的训练异常问题

问题背景 #

输入图像后，模型输出3维二值向量对应3个特征，需忽略第一个特征，仅基于后两个特征计算损失，让第一个特征不受约束。当前实现了自定义损失函数，但训练时损失和所有特征精度均平稳，模型完全无学习进展，且已确认参数处于可训练状态。

当前自定义损失函数代码：

def custom_loss(y_true, y_pred):
    n_feats = 3
    idx_to_ignore = 0
    loss = 0
    for i in range(n_feats):
        if i != idx_to_ignore:
            loss += tf.keras.losses.binary_crossentropy(y_true[:,i], y_pred[:,i], 
                    from_logits = False,                     
                    label_smoothing = 0.0, 
                    axis = -1)
    return loss

模型编译代码：

model.compile(
    loss = custom_loss,
    optimizer = optimizer,
    metrics = metrics
)

核心问题排查 #

1. from_logits 参数与模型输出激活不匹配（最可能原因） #

如果模型输出层未使用sigmoid激活函数（比如最后一层是Dense(3)而非Dense(3, activation='sigmoid')），from_logits=False的设置会导致二元交叉熵计算完全错误：

• 未经过sigmoid的输出是logit值，数值范围为(-∞, +∞)，但from_logits=False会默认将其当作[0,1]区间的概率值，此时计算出的损失会异常（要么极大要么极小），直接导致梯度消失或爆炸，模型无法更新参数。

2. 损失函数的循环操作存在冗余（非致命但影响效率） #

循环累加单个特征的损失属于TensorFlow中的低效操作，且容易因维度处理不当引入隐性问题（比如循环中张量拼接/累加的维度错误）。

3. 损失量级未归一化（可能影响优化效率） #

当前代码将两个特征的损失直接相加，导致单样本损失是单个特征损失的2倍，若优化器学习率是基于单特征损失量级设置的，会出现步长不匹配，减缓甚至停滞学习。

修复方案 #

步骤1：匹配from_logits参数与输出激活 #

• 若模型输出层用了sigmoid激活：保留from_logits=False
• 若模型输出层无激活（直接输出logit）：将from_logits改为True

步骤2：用向量化操作重构损失函数 #

替换循环逻辑，直接切片获取需要计算的特征，既高效又避免维度错误：

def custom_loss(y_true, y_pred):
    # 仅保留后两个特征（索引1和2）
    y_true_target = y_true[:, 1:]
    y_pred_target = y_pred[:, 1:]
    # 计算二元交叉熵，默认对每个样本的多个特征取平均（也可根据需求改为求和）
    return tf.keras.losses.binary_crossentropy(
        y_true_target, 
        y_pred_target,
        from_logits=False,  # 根据实际输出激活调整为True/False
        label_smoothing=0.0
    )

步骤3：可选——归一化损失量级 #

若需要让损失量级和单特征BCE一致，可对结果取平均（或除以特征数）：

return tf.reduce_mean(tf.keras.losses.binary_crossentropy(
    y_true_target, 
    y_pred_target,
    from_logits=False,
    label_smoothing=0.0
))

注：Keras训练时会自动对batch内的损失取平均，此处归一化针对的是单样本内的多个特征损失。

验证建议 #

1. 单独测试损失函数：输入几组已知的y_true和y_pred，查看损失值是否合理（比如y_pred接近y_true时损失应接近0，差异大时损失应较高）。
2. 检查模型输出范围：若为logit值，会有正负且绝对值可能较大；若为概率值，应在[0,1]区间内。
3. 监控梯度值：训练过程中打印关键层的梯度，确认梯度处于合理范围（既不是0也不是极大值）。

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