第一次用CrossEntropyLoss()做图像分割,是在一个街道场景的语义分割实验里。网络结构用的还是当时很常见的编码-解码结构,第一版代码跑起来之后,loss从0.6慢慢降到0.3,我当时还挺高兴,结果验证集的mIoU一出来直接傻眼——背景类占了接近90%,准确率很高,但车辆、行人、建筑这些真正关心的类几乎全部没学到。后来花了好几个晚上排查,发现问题基本都出在我对CrossEntropyLoss()这条函数本身的理解上:输入格式、标签编码、类别分布、ignore_index,每一个都能让训练“看似正常、实则报废”。这篇就把我用CrossEntropyLoss()做图像分割踩过的坑按类别整理一遍,每个坑后面附上可照做的检查方式和修复思路,希望能帮你少熬几个夜。
1. 为什么图像分割任务绕不开CrossEntropyLoss
1.1 CrossEntropyLoss到底在算什么
首先回到一个最基本的认知:图像分割的本质是逐像素分类。模型对每个像素输出C个类别的分数,CrossEntropyLoss()的作用就是衡量“模型预测的分数分布”和“真实类别”之间的差异。PyTorch里的nn.CrossEntropyLoss()并不是一个复杂的黑盒,它内部等价于F.log_softmax(input, dim=1)加F.nll_loss(),也就是说它会先把logits在类别维度上做log softmax,再取正确类别对应的负对数概率作为loss。
把公式拆开就是:对某个像素,如果真实类别是c,模型对该类的得分是z_c,那损失就是-log(exp(z_c) / sum_j exp(z_j))。得分越高,loss越小;得分越低,loss越大,而且梯度会因为log函数的性质变得更加明显。这也是它在分类类任务里比MSE好用的原因:MSE对输出概率的错误惩罚是二次的,到了饱和区梯度几乎消失;交叉熵在错误严重时梯度很大,优化起来效率高得多。
有一个点容易被忽略:F.log_softmax里的dim=1意味着它默认输入的第一维是batch,第二维是通道,也就是类别维度。这个约定直接影响后面所有输入格式问题,只要把通道维度放错位置,整个loss计算就会在错误的方向上“丝滑”地跑下去。
1.2 分割场景下它比BCE、MSE更合适的原因
很多做分割的新手会因为“分割也是输出mask”就想到二分类,于是直接拿BCE或者MSE去套,这种做法在二前景二背景的简单任务里能跑通,但一旦类别变多就立刻出问题。BCE是逐类别独立计算损失的,它不关心类别之间是不是互斥的,相当于把每个类都当成一个独立的二元分类问题;而语义分割的类别通常是互斥的,一个像素只能是某一个类,用CE这种带softmax归一化的损失,才能天然保证所有类别的概率加起来等于1,模型输出的是真正的概率分布。
MSE的问题更明显。如果把分割当成回归任务,模型输出的向量被强行拉向one-hot一样的分布,数值在饱和区梯度极小,训练速度慢,而且对类别判别边界非常不敏感。简单说就是:MSE问“你预测的数字和真实数字差多少”,CE问“你给正确类别的信心够不够高”,后者更贴合分类任务的本质。
所以主流分割项目的第一步基本都是criterion = nn.CrossEntropyLoss()。但要注意,默认情况下它给每个像素的权重是相同的,也就是没考虑类别分布差异。这在类别极不平衡的数据集上是致命问题,后面第4章会展开讲。
1.3 使用它的前提约束:logits与target的三条约定
使用CrossEntropyLoss()前,必须明确三条硬性约定:第一,输入必须是未经过softmax的原始logits,不是概率;第二,logits的shape一般是(N, C, H, W),N是batch,C是类别数,H和W是空间尺寸;第三,target必须是整数索引形式的标签,shape是(N, H, W),取值必须在[0, C-1]范围内,或者用ignore_index指定一个特殊值来忽略某些区域。
这三条约定错一条,训练就会出现各种奇怪现象。最常见的错误就是有人先写了prob = torch.softmax(logits, dim=1),再把prob当作输入扔给CrossEntropyLoss(),结果loss没有爆掉,但模型越训越差。还有人是把标签做成了one-hot形式,然后发现shape对不上直接报错。后面所有踩坑内容,本质上都围绕这三条约定展开。
2. 最容易踩的维度与形状坑
2.1 target的shape永远比logits少一维
这是我身边同事问过最多的问题。模型输出logits是(B, C, H, W),而mask标签往往是(B, H, W),有些人会习惯性把mask也做成(B, H, W, 1)或者(B, 1, H, W),然后直接丢进loss,接着就会看到类似Expected target size [B, H, W], got [B, 1, H, W]的报错。
为什么PyTorch要这么设计?因为C通道是类别维度,对每个像素来说,真实类别就一个整数索引,不需要展开成one-hot。如果你习惯把mask做成(B,1,H,W),在放进loss之前需要做一次squeeze(1)。如果你从图像处理库读出来的是(H,W,1)类似的多余通道,也要先去掉,目的就是让target保持(B,H,W)。
我的建议是在dataloader的返回值里就把格式固化。图像返回(3,H,W)的float tensor,标签返回(H,W)的long tensor,dataloader自动堆叠后就是(B,3,H,W)和(B,H,W),直接在loss里对齐。不要在训练循环里临时做各种view、squeeze、reshape,少一步排错步骤,就会少一个潜在坑。
2.2 通道维度在第二维,别对最后一维做softmax
PyTorch的图像张量默认是NCHW格式,也就是说类别通道维度在第二维,也就是dim=1。但这个约定太容易被忽略,尤其是从TensorFlow迁移过来的同学,TF习惯把通道放在最后一维,PyTorch里错误地写torch.softmax(logits, dim=-1)非常常见。
表面上看dim=-1也能运行,不会报错,但它在W方向上做归一化,而不是在类别方向上。举个例子,如果logits的shape是(2, 19, 256, 256),dim=-1会在最后一个维度的256个空间位置做softmax,结果是每个通道内部的空间值被归一化了,类别之间完全没有竞争关系。这种情况下loss还能正常计算,因为CrossEntropyLoss()内部还会在dim=1再做一次log_softmax,但输入已经变成了一堆被空间归一化过的数字,模型基本学不出有效的类别边界。
诊断这类问题有一个很笨但有效的办法:在loss计算前打印一行logits.shape,然后自己默念一遍“第二维是类别数”,再检查softmax的dim是1。如果嫌麻烦,就把softmax这类操作封装成一个函数,统一传dim=1,避免到处手写。
2.3 提前对logits做softmax再送进loss,等于白做
这个坑我踩过,而且是在一轮已经很深的训练之后才发现的。当时我为了在训练日志里输出“置信度”,顺手在loss之前写了一句out = torch.softmax(out, dim=1),然后才把out传给CrossEntropyLoss()。结果模型训练一直很慢,loss数值偏小,验证集几乎不动。
原因是CrossEntropyLoss()内部已经有log_softmax这一步。如果你先把softmax结果传进去,它等于在[0,1]范围内的概率值上再取一次log,得到的分布已经不再是模型原始分数的对数概率。尤其在概率接近0或1时,数值会发生奇怪的偏移,梯度方向虽然不至于完全反向,但优化过程被严重扭曲。
正确做法是:模型最后一层输出logits,直接丢给CrossEntropyLoss();需要看概率时,单独在inference阶段做softmax。如果你非要在训练过程中监视概率,就用with torch.no_grad()的代码块,只算softmax做日志,别让它进入损失计算图。
2.4 上采样、padding、裁剪引起的尺寸错位
分割模型里,编码-解码结构非常常见,很多模型在解码器末尾直接输出了下采样后的特征图,而不是原图大小。此时如果直接把logits和原始分辨率mask去算CrossEntropyLoss(),会报shape不匹配的错误。但比这个更隐蔽的情况是:你用某种方式强行把它们对齐了,但空间位置是错的。
比如DeepLab系模型最后特征图是输入尺寸的1/8或1/16,如果没做上采样,shape对不上会立刻报错;如果你写了logits = F.interpolate(logits, size=(H, W), mode='bilinear', align_corners=True),尺寸对了,但因为mask和原图经过的裁剪、padding不一样,实际像素错位,loss还在继续算,只是结果全程没有意义。
我的排查顺序是:先看原始图像进模型前的尺寸变换,再看mask有没有做过相同的resize和crop。特别是在使用随机裁剪时,必须保证image和mask使用同一个crop参数。如果你在dataloader里用了图像增强库,一定检查是不是对mask也同步做了增强,而且mask的插值方式要选最近邻。
3. mask标签处理中隐藏的坑
3.1 RGB的标注图不能直接当target
现在很多公开分割数据集给的标签是RGB伪彩色图,而不是每个像素一个索引的PNG。比如VOC格式的mask是一张彩图,背景是黑色,行人是红色,汽车是蓝色。如果直接把RGB图读进来变成(H,W,3)或(3,H,W)再扔给CrossEntropyLoss(),第一个崩溃点就是target维度不对,第二个崩溃点是数值范围不对——CrossEntropyLoss希望target是单通道整数索引,不是一个三元组。
正确做法是建一个颜色映射表,把每个RGB值映射到类别索引。速度上别用纯Python逐像素循环,数据量一大就慢得致命。可以用NumPy一行搞定每一类:先做一个全零的索引矩阵,然后对每个类别,用np.all(mask_rgb == color, axis=-1)得到一个布尔矩阵,再把对应位置填成类别id。最后转成torch.from_numpy(index).long()。
这个转换过程最好放在预处理阶段,提前批量生成索引版mask保存下来,而不是每次训练都在dataloader里现场转换。否则训练时每读一个样本都做一次全图颜色匹配,CPU负载高到无法接受。
3.2 标签值必须连续且从0开始
这个坑极其隐蔽。很多数据集的原始标签里,忽略区域用255,背景用0,但某些类别的ID可能是1、2、3之后直接跳到7、8、9,中间没有4、5、6。人的直觉是“既然是类别ID,那就直接用”,但CrossEntropyLoss()内部做索引查表时,是按target里的具体数值去取logits对应通道的。
比如logits有10个通道,代表类别0到9,target里如果出现值12,它会直接访问越界,报错;如果target里的值不连续但都没有越界,比如某些类从0到9中间缺了几个,那模型的一个输出通道可能永远不被用到,另一个类别却对上了错误通道,最终预测结果混乱。
所以在训练前要做一件看起来多余但救命的事:对训练集和验证集分别执行np.unique(mask),看看mask到底有哪些值。如果发现不连续或者存在超出类别数范围的值,就建立一个重映射字典,把所有有效值映射成0,1,2,...,C-1,忽略值映射成255。这样CrossEntropyLoss的通道和标签才能一一对应。
3.3 ignore_index是双刃剑,别指望它自动处理所有无效区域
ignore_index的设计初衷是忽略某些位置的像素,比如VOC的边界白线、标注不全的区域等。常见做法是nn.CrossEntropyLoss(ignore_index=255),然后在target里把无效区域的值设为255。这个机制很好用,但它有几个副作用需要留意。
第一个副作用是:当reduction='mean'时,分母只统计有效像素数,不是所有像素。如果一个batch里有效像素特别少,loss值会表现出不正常的剧烈波动。更极端的情况是,整个batch里的像素全是255,loss直接变成nan。这个问题在小batch训练、mask标注大量缺失时尤其明显。
第二个副作用是:如果你同时设置了类别权重weight和ignore_index,weight的长度必须等于真实类别数C,不要把忽略区域当成一个类别放进weight里。否则要么报维度不匹配,要么权重语义错乱。我习惯在计算weight时只统计非忽略区域的像素,确保每个类别都有合理的频率估计。
3.4 数据增强后mask错位和插值问题
做分割的人都用过随机翻转、随机旋转、随机缩放这些增强。很多人会把图像增强做好,但mask增强忘掉,或者虽然做了,用了双线性插值。双线性插值会产生介于两个整数之间的浮点值,比如0.2、0.8,这样的target直接传给CrossEntropyLoss()会报dtype错误,或者被某些操作截断成0和1,导致类别完全错乱。
mask的resize必须用最近邻插值,这在PyTorch的F.interpolate里就是mode='nearest',在图像增强库里有类似参数。翻转和旋转这种不改变像素值的几何变换要特别确认边界填充值不要生成垃圾类别。随机旋转时如果边界填充用了0,那么补齐区域会被当成背景类,这往往不是你想要的结果,需要显式用一个填充值,然后在target里把它映射成ignore_index。
另外一个容易被忽略的点:如果对图像做了归一化,mask不要做同样的归一化。图像需要减去均值除以标准差,mask只需要做类型转换,否则会引入浮点误差。很多新手在同一个transform里对图像和mask套了同一个标准化函数,最后target直接变成了负数或者0.几的小数,loss瞬间变成nan。
4. 训练时loss表现异常的排查
4.1 loss突然变成nan,先从数据和target查起
loss变成nan是分割训练里最常见的异常之一。很多人第一反应是调整学习率、换优化器,但我经手的案例里,大部分nan其实来自target和输入格式的问题,而不是优化器的问题。排查顺序应该是:先打印target.min()和target.max(),确认没有超出类别数范围;再打印logits.min()和logits.max(),确认没有inf或者nan;最后看loss.item()是不是一开始就是nan,如果是,那大概率是数据问题。
比如target里有值200,而你的模型输出通道只有19个,CrossEntropyLoss在取log_softmax对应通道时会越界,轻则报错,重则在某些版本里直接产生nan。还有如果标签是浮点型,必须转成long()。我建议在dataloader里就把mask = mask.long()固定好,不要等到训练循环里再改,更不要依赖自动类型推断。
排查时在loss.backward()前面加一行断言也很有用:assert loss.item() == loss.item(),这能快速发现nan。加assert不能解决根本问题,但它能把问题卡在最早出现的时刻,避免训了半天才发现loss早就坏了。
4.2 loss不下降,多半是大类别把模型带偏了
这个问题表面上跟CrossEntropyLoss关系不大,但它确实是使用默认参数最容易踩的大坑。默认的weight=None意味着每个像素对loss的贡献相同。当背景像素占95%、前景只占5%时,模型只需要把所有像素都判成背景,loss就能降到很低,mIoU却惨不忍睹。
我比较推荐的做法是:先跑一个epoch,统计各类别像素占比,然后计算中位数频率平衡权重。公式思路很简单:weight_i = median_freq / freq_i,其中freq_i是第i类像素在所有有效像素中的比例,median_freq是这些比例的中位数。这样做的效果是让出现少的类别获得更大权重,同时不至于让极稀有类权重过大导致训练震荡。
设置weight时还有一个细节:weight的每个值对应logits的通道索引,不是类别标签本身。如果你重映射过mask,那么mask的值和weight的顺序都要按照重映射后的索引来。否则即使维度匹配,权重也会“配错对象”。
4.3 整个batch全是ignore区域时的空张量问题
这个坑我调试了很久。当时用的数据集标注很稀疏,一个batch里大部分像素都是未标注区域,全是255。CrossEntropyLoss的ignore_index=255把所有像素都过滤掉之后,reduction='mean'要对0个元素求平均,结果就是nan。这个问题在batch size比较小、类别又多又稀疏时很容易发生。
修复思路有几个方向。最直接的是增大batch size,或者过滤掉那些有效标注太少的数据样本。另一个通用做法是在训练循环里检测一下有效像素数量,如果为0就跳过这个batch。但更治本的是检查数据pipeline,确认255这个值到底是“忽略区域”,还是“背景类0”被强行转成了255。我见过有人把mask填充值设成255但忘记标注它是ignore,导致所有有效像素都被丢掉了。
还有一个小技巧:如果你发现训练集里某些类别的像素总数少得可怜,与其硬扛CrossEntropyLoss,不如先做数据清洗,把标注明显错误的样本剔除。损失函数能缓解类别不平衡,但补不了标注质量的问题。
4.4 损失看起来太小或为0的几种情况
loss异常“好”不一定代表训练正常。最常见的一种情况是忘记喂logits,直接拿softmax后的概率值传给CrossEntropyLoss,此时loss数值会比正常情况小很多,因为它是在概率值上再取log,而概率值的范围天然被压缩在0到1之间,log后的负数绝对值非常小。这时候模型不一定完全学废,但跑很久也很难收敛到理想状态。
还有一种情况是target退化成了常数,比如所有mask都是背景类0。此时模型只需要输出一个常数分布就能把loss压得很低。验证集mIoU可能不为0,但仔细看每类IoU,只有一个类有数值。要排查这个问题,可以在训练开始前可视化几个样本的mask,看一眼是不是存在不同类别的像素。
如果验证时loss很低但mIoU不高,还有一个隐藏原因是模型处于model.train()而不是model.eval()模式,BatchNorm的统计量在验证时还在不断变化,导致输出不稳定。这种问题在分割任务里很常见,尤其在用验证集计算loss时,必须用with torch.no_grad()加上model.eval()。
5. 一次完整最小可复现的分割训练实操
5.1 最小数据读取与标签格式转换
下面给一个可以直接改来跑的流程,重点不是模型多强,而是让CrossEntropyLoss的输入输出严格合规。
图像读取后统一转成float32的(3,H,W)张量,mask读取后先做颜色映射或直接读索引图,然后转成int64的(H,W)张量。
def load_sample(image_path, mask_path): image = read_image(image_path).astype(np.float32) # H,W,3 mask_raw = read_mask(mask_path) # 可能是RGB或索引图 # 如果mask是RGB图,这里要完成颜色到索引的映射 h, w = mask_raw.shape[:2] mask = np.zeros((h, w), dtype=np.int64) for cls, color in color_mapping.items(): mask[np.all(mask_raw == color, axis=-1)] = cls # 统一格式:图像CHW,标签HW image = torch.from_numpy(image.transpose(2, 0, 1)).float() mask = torch.from_numpy(mask).long() return image, mask这里color_mapping要自己根据数据集定义,比如{0: (0,0,0), 1: (128,0,0)}。映射完成后做一次np.unique(mask),确认最大值小于类别数,或者把不需要的区域改成255。这个步骤虽然看起来啰嗦,但能拦住后面至少一半的问题。
5.2 一个最小分割模型与loss初始化
模型不需要太复杂,核心是最后输出logits。
class SegmentNet(nn.Module): def __init__(self, in_channels=3, num_classes=19): super().__init__() self.backbone = nn.Sequential( nn.Conv2d(in_channels, 32, 3, padding=1), nn.ReLU(inplace=True), nn.Conv2d(32, 64, 3, padding=1), nn.ReLU(inplace=True), ) self.head = nn.Conv2d(64, num_classes, 1) def forward(self, x): return self.head(self.backbone(x))损失函数我一般先这样初始化:
criterion = nn.CrossEntropyLoss(ignore_index=255)如果已经统计好了类别频率,就加上weight参数,比如weight=class_weight_tensor。注意这个tensor要放到和模型一样的device上,并且数据类型是torch.float32。不要一开始就上focal loss、dice loss之类的复杂变体,先用默认CE跑通基线,后续再根据问题换。
5.3 训练循环中损失计算的常见写法对照
训练循环的核心就几行,但错误写法非常多。正确写法是:
for images, targets in dataloader: images = images.to(device) targets = targets.to(device) # (B,H,W), long logits = model(images) # (B,C,H,W) loss = criterion(logits, targets) optimizer.zero_grad() loss.backward() optimizer.step()常见错误写法我列几个,如果你在代码里看到类似写法,直接改:
loss = criterion(torch.softmax(logits, dim=1), targets)——错误,重复softmax。loss = criterion(logits, targets.unsqueeze(1))——错误,target多了通道维度。loss = criterion(logits, targets.float())——错误,target必须是long。loss = criterion(logits.view(-1, C), targets.view(-1))——这种做法本身没问题,但前提是原始空间对应关系没乱。很多人是在HWC和CHW之间换得晕头转向,结果像素顺序全乱了。与其手动view,不如直接保持(B,C,H,W)和(B,H,W)传给loss。
5.4 验证与保存指标时要注意的细节
验证阶段很多人会直接pred = logits.argmax(dim=1),然后和target比较。这里最隐蔽的问题是argmax(dim=1)是类别维度,是对的;但如果你因为之前习惯了TF的NHWC,写成argmax(dim=-1),那你的pred会在W方向上取最大,得到的结果完全不是类别索引。
计算mIoU时,也要过滤掉ignore_index对应的像素。最简单的方法是只计算有效像素的混淆矩阵,然后逐类计算IoU。混淆矩阵的维度是C x C,但一定要先确认target的值都在0到C-1之间,否则索引直接越界。
另外,保存可视化结果时,索引mask直接保存成灰度图会很难看,最好映射回原始RGB颜色再保存。但保存和可视化用的颜色映射不能影响训练数据,它只是后处理,别把它混进dataloader里。
6. 问题排查速查表与实际心得
6.1 常见报错与修复对照表
这里整理一份我现在几乎每次排查都会对照的速查表,能覆盖大多数分割任务中的CrossEntropyLoss问题。
| 现象 | 常见原因 | 修复方式 |
|---|---|---|
| 报错Expected target size | target多了一维,或传了one-hot | 统一为(B,H,W)的long型标签 |
| loss是nan | target含越界值、ignore_index未设置、lr过高 | 打印target min/max,设置ignore_index,降低lr |
| loss数值异常小 | 对softmax结果再算CrossEntropyLoss | 改为直接传logits |
| 训练acc高但mIoU低 | 类别不平衡 | 加类别权重,或换focal/dice loss |
| label和模型通道错位 | 使用了不连续的类别ID | 重映射标签为0到C-1 |
| mask增强后错位 | mask用了双线性插值或未同步变换 | 用最近邻插值,图像和mask走同一变换 |
| 验证时loss波动大 | 忘记model.eval(),BatchNorm失效 | 验证阶段固定model.eval() |
这张表每次排查都能用上,但真正解决问题的往往不是最后一条,而是前几条数据链路的问题。所以我的习惯是:先检查数据,再去动损失函数。
6.2 几个我一直在用的诊断习惯
我现在的固定流程是:在第一次训练前,单独写一个“冒烟测试”脚本,把第一个batch的数据取出来,手动前向一次,打印以下信息:logits.shape、target.shape、target.dtype、target.min()、target.max()、logits.min()、logits.max()。这一堆信息看一眼,能拦住绝大多数低级错误。
如果训练过程中loss出现异常,我还会把当前batch的image和target可视化拼图保存下来。直接看图片比看数字快得多,尤其是mask是否错位、是否有莫名奇妙的填充值、以及255和0到底哪个是背景,一张图就能讲清楚。
还有一个小习惯是给训练脚本设置固定随机种子,并且使用确定性的dataloader。这样排查问题时能稳定复现,不会因为数据顺序变化导致现象时有时无。分割任务本来训练周期就长,能复现的问题最好在早期就控制住。
6.3 回到最初的那个项目:最后是怎么修好的
回到我开头说的那个街道场景项目,最后其实只改了三处:把RGB mask转成了连续索引,给CrossEntropyLoss设置了ignore_index=255,再按各类别像素频率计算了权重。没换模型结构,没换优化器,mIoU直接从不到20涨到接近50。
那次之后我养成了一个习惯:每次想把损失函数换成更复杂的东西之前,先问自己几个问题——target是不是long型?通道维度是不是在第2维?数据增强是不是同步了?mask里的值是不是在0到C-1之间?如果这些答案有任何一个是“没确认”,那就先别换loss,因为CrossEntropyLoss本身很简单,复杂的永远是它前后那堆数据。
最后再分享一个小技巧:如果你调来调去还是觉得loss表现诡异,先临时把数据集缩小到一个batch,拟合一个或几个样本,看loss能不能降到接近0。如果连单样本都拟合不了,那问题一定在数据链路或模型前向,而不是损失函数。这个办法我在很多项目里都用过,比反复调整优化器参数有效得多。