从零实现单类别单框YOLOv1训练损失发散输出NaN问题咨询

核心问题定位 #

你遇到的NaN问题和跳过预训练无关，核心出在损失函数的数值逻辑和输出约束上：

• 致命数值错误：直接对模型预测的w、h做平方根计算。训练初期网络参数随机初始化，未加约束的w、h预测值大概率为负，对负数开平方会直接生成NaN，经过反向传播后整个模型参数会在1~2个epoch内被完全污染。
• 梯度不稳定问题：所有损失项直接用K.sum做全局求和，损失值和梯度的量级会随batch大小、单图目标数量剧烈波动，极容易触发梯度爆炸冲垮参数。
• 输出值域缺失约束：网络最后一层没有对应激活限制输出范围，置信度P、坐标x/y、宽高w/h的输出无界，完全不符合YOLOv1的输出值域要求，会进一步加剧训练震荡。
• 额外缺失项：原YOLOv1损失包含无目标网格的置信度惩罚项，你当前实现完全跳过了这部分，会导致模型置信度预测失控，不过这不是NaN的直接诱因。

修复方案 #

1. 给输出层加对应激活，从源头限制值域
   • 置信度P、网格内相对偏移x/y接sigmoid激活，把输出约束在0~1区间
   • 宽高w/h接softplus激活，或者对原始输出求平方，保证预测值恒为非负
2. 重写损失函数，增加数值稳定性保护，归一化损失量级
   • 宽高计算前做截断，加极小epsilon值避免零值、负值进入平方根计算
   • 把全局求和替换为按batch维度求均值，稳定梯度量级
   • 补上无目标网格的置信度损失项，按原论文设置0.5的权重
3. 训练初期将Adam优化器学习率从默认1e-3调低到1e-4，等损失稳定下降后再按需调回。

修复后的损失函数代码 #

import keras.backend as K

def custom_loss(y_true, y_pred):
    # 标记包含目标/无目标的网格
    obj_mask = y_true[..., 0]
    no_obj_mask = 1 - obj_mask
    # 极小值做数值保护，避免非法运算
    eps = 1e-7

    # 中心坐标损失
    loss_xy = K.sum(
        (K.square(y_pred[..., 1] - y_true[..., 1]) + K.square(y_pred[..., 2] - y_true[..., 2])) * obj_mask,
        axis=-1
    )

    # 宽高损失：先截断保证非负，再开平方计算
    pred_w = K.maximum(y_pred[..., 3], eps)
    pred_h = K.maximum(y_pred[..., 4], eps)
    true_w = K.maximum(y_true[..., 3], eps)
    true_h = K.maximum(y_true[..., 4], eps)
    loss_wh = K.sum(
        (K.square(K.sqrt(pred_w) - K.sqrt(true_w)) + K.square(K.sqrt(pred_h) - K.sqrt(true_h))) * obj_mask,
        axis=-1
    )

    # 置信度损失：分有目标/无目标两项
    loss_conf_obj = K.sum(K.square(y_pred[..., 0] - y_true[..., 0]) * obj_mask, axis=-1)
    loss_conf_noobj = K.sum(K.square(y_pred[..., 0] - y_true[..., 0]) * no_obj_mask, axis=-1) * 0.5

    # 按batch维度求平均，稳定梯度量级
    total_loss = K.mean(loss_xy + loss_wh + loss_conf_obj + loss_conf_noobj)
    return total_loss

验证步骤 #

• 正式全量训练前，先取2~5张样本做小样本过拟合测试：如果小样本上损失能稳定下降到接近0，说明数值逻辑没有问题，再加载全量数据训练
• 训练前检查标签数据集：确认所有标签的w、h均为非负值，置信度标签只有0/1两种取值，不存在NaN、inf类脏数据

内容的提问来源于stack exchange，提问作者Francesco Pasti