Tensorflow(Keras) U-Net语义分割训练损失突发跳变问题求助

核心诱因 #

• 标签预处理错误：从代码看输入图像做了/255归一化到0_{1区间，但mask读取后没有做对应归一化，若原始标注mask用255代表前景，会导致Y_train的前景值为255，和模型sigmoid输出的0}1范围不匹配。训练前期模型输出随机性强，损失表现平稳，当模型收敛到输出接近0/1时，损失计算逻辑突然崩溃，出现跳变。
• Jaccard损失实现缺陷：没有添加平滑项，当预测结果和真实标签交集接近0时，会出现除零错误，引发梯度爆炸。
• 梯度爆炸：训练过程中梯度累积到阈值，导致参数更新幅度过大，模型权重突然偏离最优区间。

排查步骤 #

• 打印Y_train的取值范围，确认最大值是否为255，若为255则是标签未归一化的问题。
• 检查jaccard损失/指标的实现代码，确认是否添加了平滑项，计算结果是否在0~1的合理区间（你日志中出现的63、82等数值明显超出合理范围，也验证了计算逻辑有问题）。
• 在训练过程中添加回调打印各层梯度范数，确认第22轮是否出现梯度突然飙升的情况。
• 抽检训练集和验证集的图像、mask样本，确认没有损坏、标注错误的异常数据。

解决方案 #

1. 修复标签预处理逻辑，mask读取后统一归一化到0~1区间：

mask_ = imread(path + '/masks/' + mask_file, as_gray=True) / 255.0

2. 优化Jaccard损失实现，添加平滑项避免除零错误：

def jaccard_coef_loss(y_true, y_pred, smooth=1e-6):
    intersection = K.sum(K.abs(y_true * y_pred), axis=[1,2,3])
    union = K.sum(y_true, axis=[1,2,3]) + K.sum(y_pred, axis=[1,2,3]) - intersection
    return 1 - K.mean( (intersection + smooth) / (union + smooth), axis=0)

3. 新增梯度裁剪限制参数更新幅度，避免梯度爆炸：

model.compile(optimizer = Adam(lr = 1e-4, clipnorm=1.0), loss = [jacard_coef_loss], metrics = [jaccard_distance])

4. 增加学习率衰减策略，训练后期自动降低学习率，减少参数波动：

callbacks.append(tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=3, min_lr=1e-6))

训练损失曲线如下：

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