从交叉熵到mIoU：图像分割损失函数避坑指南

做图像分割这几年CrossEntropyLoss是我用得最多的损失函数也是让我花式踩坑最多的函数。就在上周A同学拿了一个训练好的分割模型找我训练loss已经从0.68稳降到0.12验证集loss也像模像样地降着但mIoU就死死卡在0.42不动。一问用的正是默认的CrossEntropyLoss()而前景像素只占整个数据集的5%。这个现象太典型了几乎每个从图像分类转到分割的开发者都会撞上。这篇文章就把我在实际项目里踩过的、以及帮别人排查过的CrossEntropyLoss相关坑一次性整理完从把softmax结果又喂给CELoss这种最隐蔽的低级失误到类别不平衡、ignore_index、weight、NaN最后给出一套完整的排查链路和替代方案。不管你是刚接触分割还是已经跑过好几个模型想精调下面的经验都值得你对着自己的代码逐条检查一遍。1. 先把CrossEntropyLoss的真身看清楚它替你做了三件事1.1 分割任务本质上是像素级分类图像分割和图像分类最大的区别在于输出粒度分类是对整张图给出一个标签而分割是对每个像素分别给出一个标签。所以从损失函数的角度看一张512×512的图、19个类别就等于一次性做了262144个独立的分类问题每个像素都要计算一次交叉熵最后把所有像素的loss平均起来。于是CrossEntropyLoss的公式就变得很直白对每个像素取预测分布中真实类别的概率取负对数。L -1/N * Σ_i log(q_i[t_i])其中q_i[t_i]是第i个像素在真实类别t_i上的预测概率N是总像素数。预测越准这个概率越接近1loss越小预测错了概率接近0loss就急剧变大。这个错了就重罚的特性正是分割任务喜欢它的原因。1.2 内置的log_softmax、NLL和reduction一个函数三层逻辑很多新手第一次用PyTorch的CrossEntropyLoss()时会以为它只是套公式算交叉熵。实际上它内部做了三件事# CrossEntropyLoss(x, target) 等价于: loss F.nll_loss(F.log_softmax(x, dim1), target, reductionmean)第一步对输入的logits在类别维度做log_softmax得到log概率第二步用NLLLoss把真实类别对应的log概率挑出来取负第三步按reduction规则汇总默认是mean即对所有有效像素取平均。这里最关键的是它期望你喂进去的是logits也就是模型最后一层卷积的原始输出而不是已经过softmax的概率。这个点后面会专门展开这里先记住结论。同时log_softmax NLL的拆法不是闲得没事而是为了数值稳定性。直接算softmax再取log在概率接近0时会出现log(0)的inf问题合并成log_softmax后数值就稳多了。1.3 一个百试百灵的初始校验首迭代loss应约等于ln(C)这是我自己最常用的一个自查手段。模型刚初始化、还没训练时logits基本是随机噪声softmax之后预测分布近似均匀分布每个类别的预测概率都接近1/C。此时交叉熵的理论值就是ln(C)。比如二分类任务首迭代loss应该在ln2≈0.693附近19个类别应该在ln19≈2.94附近。import torch import torch.nn as nn model nn.Conv2d(3, 19, kernel_size3, padding1) x torch.randn(2, 3, 64, 64) logits model(x) # [2, 19, 64, 64] target torch.randint(0, 19, (2, 64, 64)) loss nn.CrossEntropyLoss()(logits, target) print(loss.item()) # 随机初始化时应该在 2.8~3.1 之间如果你跑出来的首迭代loss远小于这个范围比如0.01那基本可以断定模型输出有问题最常见的就是模型已经被固定成某个类别、或者你传进去的是概率而不是logits。这个检查30秒就能做完能省下后面几小时的调试时间。2. 第一个大坑把softmax的结果又喂给CELoss训练慢到让人怀疑人生2.1 我是怎么踩进去的有次做一个实时分割项目为了方便可视化我在模型forward里顺手加了softmax网络直接返回概率图。训练时图省事直接把这个输出丢给了CrossEntropyLoss()。loss表现得很正常一直在降但就是降得特别慢——别人同样的模型训练50个epoch能到mIoU 0.65我的跑了80个epoch还在0.55附近晃。当时根本没往损失函数上想一直怀疑是数据增强、学习率、模型结构的问题。最后排查到forward函数才发现我传进去的x已经是概率p而CrossEntropyLoss内部又对它做了一次log_softmax等于对softmax的结果再做了一次softmax。2.2 为什么不能这样喂前面说过CrossEntropyLoss(x, target)本质是nll_loss(log_softmax(x), target)。如果你的x已经是p softmax(z)那么它实际计算的是nll_loss(log_softmax(p), target)。p本身就是一个分布再对它做log_softmax等于又过了一层softmax变换把本来已经比较尖锐的概率分布重新压平。这样做的直接后果是梯度被扭曲。原本模型输出z经过一次softmax就能把梯度传到卷积层现在要先softmax到p再从p经过log_softmax梯度在中间被稀释了。模型仍然能学只是学得极其缓慢最终精度也上不去。这种坑是最阴间的它不报错、不NaN、loss也在降只是效果不达标特别容易误导排查方向。2.3 自查手段和修正先看forward返回值。正确做法是模型永远返回logits只有在推理可视化时才在外面套softmax# 训练时 def forward(self, x): return self.backbone(x) # 返回 logits # 推理时才转概率 with torch.no_grad(): prob torch.softmax(model(x), dim1)如果已经踩坑了可以用一个简单的办法自检看训练后模型输出概率的分布。正常训练到后期模型的max概率会普遍接近0.95甚至更高如果是二次softmax的模型max概率往往被压在0.5~0.7之间很难达到高置信度因为第二层softmax会把原本的高置信度拉平。这个特征比看loss曲线好使多了。3. 第二个大坑类别不平衡loss在降、mIoU原地躺平3.1 症状和背后的数学回到开头A同学那个案例二分类分割背景占95%前景建筑占5%。交叉熵是逐像素平均意味着loss里有95%的贡献来自背景像素。模型只要学会把所有像素都预测为背景就能拿到非常低的loss和97%的像素准确率但前景类的IoU是0mIoU自然也就上不去。这就是交叉熵在分割任务里最大的短板它按像素平均而分割任务的评价指标mIoU是按类别平均。两者对哪个像素重要的定义完全不同。所以经常出现loss降得漂亮、mIoU纹丝不动的奇特现象。排查思路很简单别只看总loss把每个类别的IoU单独打印出来绝大多数情况下你会发现是一个或几个少数类的IoU几乎为0。3.2 用class weight修正最直接的修正方式是给少样本类别更大的惩罚权重也就是weight参数。权重通常用逆频率计算def build_class_weight(labels, num_classes, ignore_index255): mask labels ! ignore_index counts torch.bincount(labels[mask].flatten(), minlengthnum_classes).float() freq counts / counts.sum() weight 1.0 / (freq 1e-6) # 逆频率 # 归一化让权重均值在1附近方便调lr weight weight / weight.sum() * num_classes return weight实际操作中要注意三点权重要在训练集全量统计而不是每个batch算一次。batch小的时候某个batch里可能完全没有某个类别bincount会算出0除以0就inf了。权重不要无限大。逆频率算出来的权重可能高到上百训练会剧烈震荡。我一般会把权重上限clamp到10~20之间。加权重后loss量级会变原来lr 1e-3可能就偏大了需要重新按loss绝对值调lr。如果逆频率效果不好可以试试中频平衡median frequency balancing权重 全类别频率的中位数 / 该类频率。它比逆频率温和对小类别没那么激进在不少分割任务上更稳。3.3 加权还不够就上Focal Loss如果类别不平衡极端到某个类只有0.1%像素加权交叉熵也容易在训练后期被已经很好分类的背景主导梯度。这时候可以换成Focal Loss它的思路是降低易分类样本的权重让模型聚焦难样本class FocalLoss(nn.Module): def __init__(self, gamma2.0, alphaNone, ignore_index255): super().__init__() self.gamma gamma self.alpha alpha self.ignore_index ignore_index def forward(self, logits, target): mask target ! self.ignore_index target_c target.clamp(min0).unsqueeze(1) log_pt F.log_softmax(logits, dim1) log_pt log_pt.gather(1, target_c).squeeze(1) pt log_pt.exp() focal -((1 - pt) ** self.gamma) * log_pt if self.alpha is not None: focal focal * self.alpha[target.clamp(min0)] return focal[mask].mean()gamma通常取2.0alpha可以根据逆频率来。我的经验是先试weighted CE效果不够再换Focal不要一上来就Focal。Focal对难样本过于敏感如果数据标签本身有噪声会把噪声当难样本反复强化。4. 第三个大坑ignore_index、weight、全ignore区域三个坑连着踩4.1 ignore_index的正确用法分割数据集的标注里通常会有一些不参与训练的区域比如街景任务里的天空边界、物体边缘的模糊区域公开数据集往往把它们标成255或者-1。这时要用ignore_index255让这些像素不产生losscriterion nn.CrossEntropyLoss(ignore_index255)看着简单但坑就藏在数据预处理里。有些流程为了处理方便会把标签整体减1于是原本的255变成了254而原本的-1还在。如果你设置了ignore_index255但标签里实际是-1PyTorch会直接报Target -1 is out of bounds因为-1不在[0, C-1]范围内。正确做法是在数据加载阶段统一改标签把所有的-1和255都归一到同一个ignore值上再进loss。千万别想着通过给ignore单独设一个类别来规避那样等于凭空加了一类会让模型去学无意义区域mIoU反而更难看。4.2 算weight时忘了排除ignore像素这个坑很隐蔽。有些人算类别权重时直接用整张标签图统计bincount把255也算进去了。如果ignore区域的像素占比很大有些街景数据集能到20%以上这些无意义像素会人为抬高某个类别的频率导致那个类别的权重被压低训练出来的模型在真实类别上就系统性偏差了。正确的统计逻辑必须和loss一致先用mask把ignore像素剔除再统计各类别频率。上面build_class_weight里的mask labels ! ignore_index就是干这个的。4.3 空类和全ignore区域的边界情况当某个batch里前景像素特别少甚至一个batch里所有像素都是ignore时loss会退化成空集操作直接NaN。PyTorch内置的CE在全部ignore时会返回0这个行为我没完全确认不同版本可能有差异但如果你自己写了loss[mask].mean()这种自定义逻辑mask全False时reduce mean就会NaN。稳妥做法是加一个保护loss F.cross_entropy(logits, target, reductionnone, ignore_index255) mask target ! 255 loss loss[mask] if loss.numel() 0: return 0.0 * logits.sum() # 保持梯度通路 return loss.mean()还有空类问题某个类别在训练集里存在但在当前batch里没有bincount算出来是0后面除权重时就出inf。统计时记得clamp最小值或者用torch.where(count 0, 1, count)兜底。4.4 ignore_index和weight同时用顺序别搞反内置的nn.CrossEntropyLoss(weightw, ignore_index255)是支持两者同时使用的但权重张量的长度必须是C且不包含ignore类——也就是说255对应的权重位置是什么值都没关系因为255的loss本身就被置零了。我自己见过有人把权重长度写成了C1多了一位结果loss行为变得很奇怪还百思不得其解。报错虽然不会报但输出就是不对这种不报错的逻辑错误是最耗时的建议在初始化后先打印weight的shape和数值确认一遍。5. 第四个大坑维度、dtype和设备不一致四个报错一个对策5.1 shape错误清单CrossEntropyLoss对输入shape的要求非常严格logits是[N, C, H, W]target是[N, H, W]target不带通道维。实际项目里常见的低级错误我列个表错误用法报错或现象正确写法target是[N, 1, H, W]Expected target size [N, H, W], got [N, 1, H, W]target.squeeze(1)target是one-hot [N, C, H, W]CE不认报target维度和logits不一致转成索引target.argmax(dim1)logits是[N, H, W]忘了类别维报shape不匹配检查最后一层是否输出C个通道logits是[N, H, W, C]报错或结果完全错误logits.permute(0, 3, 1, 2)这个表里最坑的就是one-hot格式。很多人做分割之前先做过多标签分类习惯了one-hot target到了分割任务还保持这个习惯结果CE的target要的是类别索引而不是one-hot向量两者一混就出岔子。5.2 dtype问题target必须是torch.long类型这是最容易被忽略的。从numpy读进来的标签通常是np.int64转成tensor时用torch.from_numpy(labels)默认就是long没问题但有些人用torch.tensor(labels)如果labels是np.uint8得到的tensor就是uint8喂给CE直接报错。正确习惯是统一做一次强制转换target torch.tensor(labels, dtypetorch.long, devicelogits.device)另外混合精度训练下如果你手动拆开写-torch.log(torch.softmax(logits, dim1))在fp16下很容易因为log(0)或softmax溢出产生NaN。而内置的CrossEntropyLoss是经过数值稳定性优化的一般不会有问题。所以我的建议很朴素能用内置函数就别手搓。5.3 设备不一致多卡或混合设备训练时最经典的报错是RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!这个其实好解决关键是定位是哪个tensor不在同一个设备上。调试技巧是在进loss之前打印print(logits.device, target.device)一目了然。但要注意DataParallel这种封装会让问题变得隐蔽有时候单个GPU跑没问题一上多卡就报设备错误那就得检查target有没有跟着模型一起to(device)。5.4 NaN排查链路最后说NaN。loss突然变NaN几乎都是因为logits里出现了inf或NaNsoftmax一算就崩。排查顺序我固定是这几步在loss之前打印torch.isnan(logits).any()和torch.isinf(logits).any()先确认是不是logits本身炸了。如果是logits炸了大概率是lr太大导致梯度更新步长过大把权重推到inf。先降lr到1/10试跑通常就好了。如果不是lr检查混合精度设置fp16下某些自定义loss会溢出。实在找不到开torch.autograd.set_detect_anomaly(True)PyTorch会在第一次产生NaN的backward位置给你提示能缩短很多排查时间。6. 一次完整的排查实录loss降、val也降mIoU就是不动6.1 现象与初始判断回到A同学那个案例。他做的是遥感建筑分割二分类前景占比5%。训练交叉熵从0.68降到0.12验证集从0.72降到0.18像素准确率97.2%但mIoU稳定在0.40~0.42之间从第30个epoch开始就再也不涨了。我第一反应是让他先别动模型把下面三样东西打印出来每个类别的IoU、每个类别被预测的像素占比、若干张预测可视化图。结果很典型背景IoU 0.92前景IoU 0.08模型预测中建筑类像素面积占比只有2%而真实标签里建筑占5%可视化里只有大片、规则的建筑能分割出来小建筑基本漏检。判断基本明确模型已经学会了把大概率是背景的像素全判成背景靠几块大建筑勉强维持住一点前景IoU小目标全部放弃。这就是类别不平衡下交叉熵优化的必然结果loss越低模型越懒。6.2 逐层排查的完整过程第一步先检查评估口径。mIoU的计算必须和训练时的类别对齐别把ignore区域算进去也别把预测缩放到错误尺寸。A同学这个项目里评估代码没问题排除。第二步打印训练集每个batch的前景像素比例。这一步很容易发现另一个隐藏坑他的数据采样是随机的很多batch里前景像素占比不到1%甚至有个别batch全是背景。在这种batch上加权CE的梯度也几乎全是背景的梯度。第三步加class weight。按逆频率计算权重前景权重约20背景权重约1clamp到10。加了之后mIoU从0.42涨到0.55但还是卡着不动。第四步加数据采样修正。改用带类别感知的随机采样保证每个batch里前景像素不低于某个阈值比如2%同时加了小目标友好的增强操作。这一步之后mIoU到0.58。第五步换损失组合。最终方案用weighted CE dice loss联合两个loss等权相加同时把学习率从1e-3降到3e-4。mIoU最后到0.63比初始方案提升了20个点。6.3 根因与复盘这个案例里真正的根因有两个一是纯CE无法对抗极度不平衡二是数据采样让不平衡雪上加霜。光加权重不够因为梯度仍然被背景主导光改采样不够因为没有损失函数层面的矫正。两者结合起来才解决问题。复盘时还有个小细节值得单独记val loss从0.18继续降到0.16的时候mIoU反而掉了。原因是val loss是逐像素平均背景占比高意味着它对背景过拟合不敏感而mIoU按类平均对少数类极其敏感。所以分割任务里必须把mIoU当主指标loss只能当辅助参考别被loss的漂亮曲线迷惑。7. 什么时候别硬撑CrossEntropyLoss三条不得不换的路线7.1 每个像素可能属于多个类改用BCEWithLogitsLoss语义分割的标签通常是互斥的一个像素只能属于一个类用CE没问题。但有些分割任务是真·多标签比如一个像素同时属于建筑和阴影两个类别或者医学影像里一个像素同时属于器官和病灶。这时候softmax加CE的思路天然就不成立因为你强行让所有类别概率加起来等于1一个像素既是建筑又是阴影时概率分布根本无法表达。正确做法是每个类别独立做二分类用BCEWithLogitsLosscriterion nn.BCEWithLogitsLoss() loss criterion(logits, target_one_hot_float) # target: [N, C, H, W] 的float张量注意target格式和CE完全不同CE要的是索引BCE要的是每个通道的0/1浮点标签。如果少数类别特别少还可以给BCE加pos_weight参数按类别给正样本加权。这个pos_weight一般取负样本数/正样本数的比例效果和CE的weight类似。7.2 只关心少数几个难分的小目标Focal Loss更直接Focal Loss和第3.3节代码一样这里不重复。要补充的是我自己的选型经验当加权CE已经把mIoU推到接近预期、但少数类仍差几个点时换Focal或者给少数类加大gamma往往有奇效当少数类压根没被识别出来时先解决采样和mIoU评估不要急着上Focal。Focal的另一个副作用是训练初期loss值会比CE大不少因为难样本的加权比例高了。这会导致你看起来loss降不下去实际上是量纲变了不是模型有问题。有心理准备就好。7.3 边界糊、小目标漏检和Dice Loss做组合拳CE是逐像素独立计算的完全不关心预测区域的形状、连通性、边界质量。这是它的结构性缺陷。而Dice Loss直接优化预测区域和真实区域的交并比对小目标非常友好边界也通常会因为区域约束而更干净def dice_loss(logits, target, ignore_index255, eps1e-6): probs torch.softmax(logits, dim1) mask target ! ignore_index target_c target.clamp(min0) target_onehot F.one_hot(target_c, num_classesprobs.shape[1]) target_onehot target_onehot.permute(0, 3, 1, 2).float() probs probs * mask.unsqueeze(1).float() target_onehot target_onehot * mask.unsqueeze(1).float() inter (probs * target_onehot).sum(dim(0, 2, 3)) cardinal probs.sum(dim(0, 2, 3)) target_onehot.sum(dim(0, 2, 3)) dice (2 * inter eps) / (cardinal eps) return 1 - dice.mean()融合方式我建议初期就固定为简单加权求和total ce_weighted dice两个loss量级不同按1:1往往会让dice主导导致训练震荡。我的习惯是先各跑几个epoch看量级然后让CE和Dice大致等权比如total 1.0 * ce 0.5 * dice。具体系数没有金标准但记住一个原则换损失函数等于换了优化目标学习率必须重新调。最后分享一个我固化下来的习惯任何分割项目开跑前先打印三个数字——首迭代CE理论值(ln C)、训练集各类别像素占比、模型输出的max logit范围。三个数对上了才开始调参。别问我为什么问就是我曾经因为漏掉第一个检查在二次softmax那个坑里耗了两天。