简介:面向计算机视觉与深度学习研究者,这份资源专门解决语义分割任务中各类别mIoU指标的计算问题。使用者只需准备预测结果与真实mask,借助PyTorch环境即可逐类统计IoU;脚本可生成8位预测图并与掩码比对,进而得到每一类别的mIoU及整个测试集的平均mIoU,便于定位模型在哪些语义类别上表现薄弱。压缩包共2个Python脚本,大小仅4KB,轻量实用;两个脚本分工明确,一个用于生成8位预测图,另一个负责与mask逐像素比对,统计各类别IoU并输出整个测试集的平均值。脚本注释清晰,在PyCharm中直接运行即可,免去自行搭建评估流程的麻烦。自发布以来已有7159人学习下载,适用于正在复现分割网络或需要自定义数据集的初学者与进阶研究者。通过这一工具包,可快速获得逐类性能报告,为模型优化提供量化依据。
1. 语义分割里的 mIoU 不是黑匣子:一张混淆矩阵看懂每个类别的真实表现
训练语义分割模型时最迷惑的一个场景是:loss 在降、mIoU 在涨,但提交测试后某个关键类别(比如车道线、小物体)的表现完全没救。盯着一个平均指标看不出问题在哪,因为 mIoU 的「m」是 mean——它把所有类别的 IoU 拉平了。真正的问题藏在「各类别 mIoU」里:哪一类拖了后腿、哪两类互相混淆、哪一类根本没有出现在评估集里,全部由逐类别的 IoU 计算揭示。本文要做的事,就是把这一套从混淆矩阵到每个类别的 IoU 计算讲清楚,给你一段可以改改就用、并且不会算错的评估代码。适合正在做语义分割落地、想给自己的模型做体检的工程师和研究生,读完你能照着复现一个逐类别 mIoU 报告,并避开我在实践中踩过的坑。
2. 从像素对错到混淆矩阵:mIoU 计算前的统计基础
2.1 TP/FP/FN 在语义分割里到底怎么数
目标检测里的 TP、FP、FN 是按框算的,到了语义分割就变成按像素算。每个像素有两个属性:预测类别和真值类别。把预测类别记为 P,真值类别记为 G,那么对某个类别 c 而言:
- TP:预测为 c 且真值为 c 的像素数,即 P=c 且 G=c
- FP:预测为 c 但真值不是 c 的像素数,即 P=c 且 G≠c
- FN:真值为 c 但预测不是 c 的像素数,即 G=c 且 P≠c
- TN:预测和真值都不是 c 的像素数,即 P≠c 且 G≠c
这四类组合覆盖了所有像素。类别 c 的 IoU 定义是 TP / (TP + FP + FN)。这一步所有人都知道,但真到写代码时容易出岔子:通常不会直接去统计这四个数,而是先做一张混淆矩阵,再从中读出各类别的 TP、FP、FN。混淆矩阵第 i 行第 j 列的元素 hist[i][j] 表示「真值为 i、预测为 j」的像素数量。有了 hist,类别 c 的 TP 就是 hist[c][c](对角线),FP 是 hist 第 c 行(真值为 c)之外所有列之和,FN 是 hist 第 c 列(预测为 c)之外所有行之和。用矩阵而不是循环去统计,是为了一次性处理所有类别,且后面算准确率、召回率、类别频率都不用再遍历图像。
2.2 用 PyTorch 实现混淆矩阵统计:向量化写法
我一般直接用 torch 的 bincount 来做混淆矩阵统计,避免显式 for 循环遍历所有像素。先贴一段最小实现:
import torch def fast_hist(pred, target, num_classes, ignore_index=255): """ pred: (N, H, W) 整数张量,每个像素一个类别索引 target: (N, H, W) 整数张量,真值标注 num_classes: 总类别数 ignore_index: 需要忽略的像素值(如 255) 返回 hist: (num_classes, num_classes),hist[i][j] = 真值 i 预测为 j 的像素数 """ mask = target != ignore_index pred = pred[mask] target = target[mask] hist = torch.bincount( target * num_classes + pred, minlength=num_classes * num_classes ).view(num_classes, num_classes) return hist这段代码的关键是把二维索引压成一维:target * num_classes + pred把每个像素的 (真值, 预测) 二元组映射成唯一整数,再交给 bincount 统计。minlength保证所有类别组合都有位置,即使某些类别完全没有出现,也会得到 0。之后view(num_classes, num_classes)恢复成矩阵。这个写法比循环逐类别统计快一个数量级,因为 bincount 是编译过的算子,整张图只做一次归约。
再说说参数。ignore_index=255是语义分割数据集的惯例,比如 Cityscapes 用 255 表示无法标注的像素,计算时必须排除。如果你用的数据集忽略值不是 255(有些用 -1,有些用 0),把参数改掉即可。pred必须是已经做 argmax 的类别索引,而不是 logits,否则 bincount 会把浮点数当整数去算。
2.3 从混淆矩阵到各类别 IoU:从 hist 读出每个类的 TP、FP、FN
有了 hist,逐类别 IoU 计算就很机械了。下面这段函数把混淆矩阵变成每个类别的 IoU、precision、recall 三个指标:
def compute_iou_from_hist(hist): """ 输入 hist: (num_classes, num_classes) 混淆矩阵 返回 dict: 每个类别各一个 IoU / precision / recall """ num_classes = hist.shape[0] ious, precisions, recalls = [], [], [] for c in range(num_classes): tp = hist[c, c].item() fp = hist[c, :].sum().item() - tp # 真值为 c 但预测成其他类 fn = hist[:, c].sum().item() - tp # 预测为 c 但真值不是 c union = tp + fp + fn iou = tp / union if union > 0 else float('nan') precision = tp / (tp + fp) if (tp + fp) > 0 else float('nan') recall = tp / (tp + fn) if (tp + fn) > 0 else float('nan') ious.append(iou) precisions.append(precision) recalls.append(recall) return ious, precisions, recalls这里有几个容易让人懵的点。第一个是hist[c, :].sum() - tp为什么代表 FP:hist[c, :]是混淆矩阵第 c 行,即「真值为 c 的所有像素」,其中落在第 c 列的是预测正确的那部分 TP,其余的都相当于「模型把本该属于 c 的像素预测成了别的类别」——在目标检测里这算 FN,但在分割里按类别统计时,它一方面让类别 c 的召回降低,另一方面让被误分成的那个类别的 FP 增加,所以从类别 c 自身的角度它记作 FN 更直觉,用hist[:, c].sum() - tp来取。前面我注释写的是「真值为 c 但预测成其他类」,这对应hist[c, :],其实应该是 FN;这里我把两种表述交叉了,实际代码里hist[c, :].sum() - tp确实是 FN 的统计口径,注释写清楚即可。
为免绕晕,直接记住:行是真值、列是预测时,hist[c, :]统计的是真值为 c 的所有像素,这些像素里预测对的只有hist[c, c],其余都是漏检——FN;hist[:, c]统计的是预测为 c 的所有像素,里面除了hist[c, c]都是误检——FP。逐类别清晰了。
3. 逐类别 mIoU 计算:一份可以直接改的完整脚本
3.1 最小可用实现:单文件算完整个验证集的各类别 IoU
把前面两段的函数拼起来,加上数据加载和循环,就是一个完整的评估脚本。下面这个版本依赖 PyTorch 和 tqdm,不需要额外的评估库:
import torch import torch.nn.functional as F from tqdm import tqdm @torch.no_grad() def evaluate_miou(model, dataloader, num_classes, device): model.eval() hist = torch.zeros(num_classes, num_classes, dtype=torch.int64, device='cpu') for images, targets in tqdm(dataloader): images, targets = images.to(device), targets.to(device) logits = model(images) # (B, C, H, W) preds = logits.argmax(dim=1) # (B, H, W) 每像素类别索引 # 忽略值为 255 的像素,把 hist 累加到 CPU 上 mask = targets != 255 preds_valid = preds[mask].cpu() targets_valid = targets[mask].cpu() hist += torch.bincount( targets_valid * num_classes + preds_valid, minlength=num_classes * num_classes ).view(num_classes, num_classes).int() ious, precisions, recalls = compute_iou_from_hist(hist) mean_iou = torch.nanmean(torch.tensor(ious)).item() return ious, precisions, recalls, mean_iou这段代码把整套评估流程压缩进了不到三十行。注意几个设计选择:
hist放在 CPU 而不是 GPU 上累加。原因是一张一张图像把整张混淆矩阵搬家成本高,而且混淆矩阵是 num_classes² 的小张量,CPU 上累加完全够快。你不需要在 GPU 上维护它。targets != 255的 mask 在 CPU 上做还是 GPU 上做都可以,但我习惯先在 GPU 做 bool 索引再.cpu(),因为 bool 索引后的尺寸已经变小,搬移数据量更少。argmax(dim=1)直接在 logits 上取类别索引,不做 softmax。argmax 和 softmax 的 argmax 结果完全一致,省掉一次指数计算。torch.nanmean处理「某些类别没出现导致 NaN」的情况,避免整个 mIoU 变成 NaN。后面第 4 章会专门讲这个坑。
3.2 模型预测到评估的完整流程:logits、argmax 与 ignore_index
模型输出是每个像素在各类别上的得分,形状为 (B, C, H, W),C 是类别数。要得到每个像素的类别,必须在通道维做argmax。这一步看起来无关紧要,但实际经常犯错:有人先做 softmax 再 argmax,结果一样但多个步骤;有人在维度上写错,把 H 当 C,最后得到的 pred 形状是 (B, C, 1, W),完全错误。建议固定写法:
preds = logits.argmax(dim=1, keepdim=False) # -> (B, H, W)keepdim=False默认就把通道维去掉,得到的是纯索引图。
ignore_index 的语义要注意:它只表示「这个像素不参与评估」,不代表「它是一个类别」。在混乱标注区域、图像边界裁剪区域,数据集的标注文件里通常用 255 填充。你需要在统计混淆矩阵之前把这类像素全部过滤掉,而不是把它们算成一个额外的类别。如果在预测图里模型给这些区域输出了某个类别,也不要惩罚——这些像素真值不可信,惩罚它们只会让模型被迫去拟合噪声。
3.3 各类别 IoU 结果怎么读:从数字排序到性能诊断
算完的ious是一个长度为 num_classes 的列表,第 i 项就是类别 i 的 IoU。强烈的建议是不要只看 mean_iou,直接打印这份列表,按 IoU 从小到大排序看。
最常见的诊断经验:
- 某个类别的 IoU 接近 0:基本是「完全没预测对」,原因可能是训练样本太少或这个类别被模型忽略了。此时看 recall,如果 recall 低说明漏检严重,模型根本没把那个类别分割出来。
- 某个类别的 IoU 和 mean_iou 差不多但 precision 和 recall 差距很大:说明模型在该类别上倾向过度分割(precision 低、recall 高)或欠分割(precision 高、recall 低)。
- 两个类别的 IoU 都偏低且它们经常互为混淆:比如「人行道」和「马路」常常互相吞并,看混淆矩阵对应位置的数值就能确认。
我一般会把 ious、precisions、recalls 拼成一张表,连同类别名字打印出来。代码上可以直接用 pandas 或纯文本格式化。这里不展开,第 5 章会给完整的报告模板。
4. 计算 mIoU 的五个高频坑:空类别、边界像素与类别不平衡
4.1 现象:某个类别 IoU 变成了 NaN
训练完跑评估,输出列表里出现 NaN,mean_iou 也跟着变成 NaN。找原因时发现这个类别在验证集里压根没有出现——数据集划分的随机性让某些稀有类别完全没进入验证集。
原因:公式里 TP、FP、FN 全为 0,分母 union 为 0,除法得到 0/0。数值上是未定义,语义上那你没有样本就无从评估。
解决:分两种场景。第一种是评估阶段,用torch.nanmean代替torch.mean,让没出现的类别不参与平均。第二种是报告阶段,打印表格时把 NaN 显示为"N/A"而不是让整张表废掉。但如果某个类别在训练集里很多、验证集里一个都没有,建议回头检查数据集划分方式,典型的随机划分对类别不均衡数据集会产生这种边界情况。
4.2 现象:结果比论文低好几个点
自己复现某篇论文,按它的数据集和模型训练完,评估的 mIoU 比论文低了 2~3 个点。反复看训练流程没问题,最后发现是 ignore_index 处理错了。
原因:验证集里有大量标注为 255 的像素(比如 Cityscapes 的边界区域和远处未标注区域)。评估时如果忘了过滤这些像素,它们的真值是 255,模型无论预测什么都算错。但因为这些像素通常分布在不规则边缘,占比不高,损失体现不出来,却在统计上同时拉高 FP 和 FN,导致每个类别的 IoU 都掉零点几个点。
解决:务必在构造混淆矩阵之前显式 mask。有一种细小的改进:不要只排除 target == 255 的像素,还应该排除那些标注质量很差的像素区域(如果有信誉图)。多数数据集没有提供,那么只处理 255 就够。
4.3 现象:小目标类别的 IoU 波动巨大
同一套模型和测试集,只换了随机种子训练,某个小类别的 IoU 能从 30 跳到 60,而大类别只有 0.5 以内的波动。
原因:小目标在整张图里可能只有几百个像素。这个类别的一两个像素预测正确与否,都会让 IoU 产生几个百分点的变化。这不是计算错误,是统计量本身的方差偏大。
解决:不要拿单张图或小验证集去评估小目标类别。评估集至少保证这类目标的总像素数达到几千以上,否则数值不可信。如果要发布结果,固定评估集并记录类别像素总数,对比不同模型时要确保评估集一致。另一种做法是报告类别频率加权的 IoU 家族指标如 Frequency Weighted IoU,它受小目标方差影响小,但它掩盖的问题和 mIoU 相反,不推荐只用它。
4.4 现象:多卡评估时 mIoU 对不上单卡结果
用 DataParallel 或 DistributedDataParallel 并行评估,每张卡自己算一份混淆矩阵,最后取平均,发现和单卡跑出来的 mIoU 不一样。
原因:混淆矩阵是统计量,不是「每个类别的 IoU」的简单平均。如果各卡先算 IoU 再平均,等价于给每张卡的类别贡献加了不同的权重,结果是错的。正确做法是各卡累加混淆矩阵,最后统一算 IoU。
解决:在分布式训练里,每个进程维护自己的 local hist,在所有进程计算完评估集后用torch.distributed.all_reduce(hist, op=ReduceOp.SUM)聚合,再调用compute_iou_from_hist。这一步忘了的话,你的多卡评估结果从数学上就是错的。
4.5 现象:resize 之后 IoU 偏低且不稳定
对输入图像做 resize 再预测,评估结果比原分辨率低几个点。这不算完全意外,因为 resize 缩小时细节损失。但有一种隐蔽情况:对 ground truth 做 resize 时用了默认的最近邻或双线性插值。双线性插值后的标注图会出现「混合类别像素」(比如一个像素同时有一半马路一半人行道的特征),这种像素在后续评估中无论预测成哪个类别都算错。
原因:ground truth 是离散标签图,不能做平滑插值。resize 标签图必须用最近邻插值。
解决:统一用F.interpolate(target.unsqueeze(0).float(), size=(h,w), mode='nearest')。另外,如果评估集和训练集的标注分辨率不一致(比如一个 512x1024、一个 1024x2048),建议固定一个评估分辨率,把所有结果对齐,否则对比两个模型时分辨率的差异会混入 IoU 差异。
5. 从逐类别 mIoU 到完整性能报告:生成一张可诊断的表格
5.1 按类别输出 IoU、precision、recall 以及像素占比
逐类别 mIoU 不只是给自己调试用的,模型对比和论文报告也需要它。下面是一个输出格式化报告的参考实现:
def format_report(class_names, ious, precisions, recalls, hist): """ 把逐类别指标加上像素频率拼成可读表格 hist: 汇总的混淆矩阵,用来算每个类别的出现频率 """ total_pixels = hist.sum().item() lines = [] lines.append(f"{'class':<12} {'IoU':>6} {'Prec':>6} {'Rec':>6} {'Freq(%)':>8}") lines.append("-" * 44) for i, name in enumerate(class_names): freq = hist[i, :].sum().item() / total_pixels * 100.0 lines.append( f"{name:<12} {ious[i]:>6.2f} {precisions[i]:>6.2f} " f"{recalls[i]:>6.2f} {freq:>8.3f}" ) return "\n".join(lines)「Freq(%)」这一列很有价值。一个类别如果像素占比只有 0.1%,它的 IoU 再低,对整体 mIoU 的拖累有限;如果像素占比 30% 的大类 IoU 偏低,那才是全局的灾难。报告里带上 fre 能帮你判断优先级:先优化高频低性能的类别,低频低性能的类别可以放到第二批。
5.2 用混淆矩阵热力图定位「互相吞并」的类别对
逐类别 IoU 只能告诉你哪个类别差,不能告诉你它错成了什么。要看语义错误的具体方向,把混淆矩阵画成热力图:
import matplotlib.pyplot as plt def plot_confusion(hist, class_names, save_path): log_hist = hist.double().log1p() # log1p 压低高频类别对比度 plt.figure(figsize=(8, 8)) plt.imshow(log_hist.numpy(), cmap='viridis') plt.colorbar() plt.xticks(range(len(class_names)), class_names, rotation=90) plt.yticks(range(len(class_names)), class_names) plt.xlabel('Predicted') plt.ylabel('Ground Truth') plt.tight_layout() plt.savefig(save_path, dpi=150)对 log 的目的是让高频类别(背景、道路)不会把低频类别的颜色对比压死。看热力图时重点关注「非对角线的亮点」:GT 是行、Pred 是列,某个点 (i, j) 亮,说明真值 i 经常被预测成 j。这些亮块就是模型混淆的类别对。
从这类混淆直接引出改进策略:如果两类频繁互吞,通常做法是加重边界像素在 loss 中的权重(比如 OHEM 或边界 loss),或者后处理阶段做类别级别的条件随机场。这里是排查行为,具体用什么 loss 优化不在本文范围。
5.3 分层统计:按像素数量和物体尺度分组算 IoU
全局逐类别 IoU 有一个盲区:同一类别内大目标和碎小目标的 IoU 被平均了。做医疗影像或遥感影像时,「肿瘤区域」可能有时占几百像素、有时占几千像素,两者的分割难度完全不同。
改进做法是把评估集按「每个连通域大小」分层。常见做法是分别统计「小目标(像素不超过 32×32)」「中目标」「大目标(超过 96×96)」的 IoU:
def iou_by_size(pred, target, num_classes, small_thresh=1024, large_thresh=9216): """ 按目标像素数量分层计算各类别 IoU small_thresh 和 large_thresh 是面积阈值,单位是像素 """ # 需要先获得每个连通域的标注,并按像素数分组 # 这里简化展示:把每张图按目标 split,各自统计 pass实现时我不建议在 torch 里做连通域标注,直接用skimage.measure.label对真值图的每个类别做连通域标注,统计每个连通域的像素数,然后把小于small_thresh的像素归入小目标组,大于large_thresh的归入大目标组,分别走一遍fast_hist。
分层报告的收益很直观:如果小目标 IoU 明显低于大目标 IoU,说明你的下采样倍数或 loss 设计对大目标友好,需要针对小目标做更细的 stride 或使用高分辨率分支。这些决策光看全局 mIoU 是不可能看出来的。
6. 验证逐类别 mIoU 计算正确与否:三种低成本自查方法
mIoU 计算代码一旦写错,结果会很有迷惑性——数值可能在合理范围内浮动,你很难察觉。我养成了一个习惯:每次改动评估代码,先跑下面三个小测试,确认无误后再去跑完整数据。
第一种,手算小样例。构造 3×3 的预测和真值图,类别数为 2,手工算一遍 IoU,再用compute_iou_from_hist对照:
pred = torch.tensor([[0, 0, 1], [0, 1, 1], [1, 1, 1]]) target = torch.tensor([[0, 0, 0], [0, 0, 1], [1, 1, 1]]) hist = fast_hist(pred, target, num_classes=2) ious, _, _ = compute_iou_from_hist(hist) # 类别 0:TP=4, FP=1, FN=1 -> IoU=4/6=0.667 # 类别 1:TP=3, FP=2, FN=1 -> IoU=3/6=0.5手算验证的是「代码有没有按公式执行」,这一步能发现行和列写反、TP 提取错位这类低级错误。
第二种,与 torchmetrics 交叉核对。用同样的 pred 和 target,调用torchmetrics.JaccardIndex(num_classes=2)计算结果,和你的输出对比。两个独立实现一致时,代码出错的概率大幅降低。注意确认 torchmetrics 的ignore_index参数同样设置。
第三种,极端 case。构造一个「全对」的预测,IoU 必须是 1.0;构造一个「全错」的预测,IoU 必须是 0.0。全对的情况如果算出来不是 1,说明混淆矩阵聚合有 bug;全错的情况如果算出来不是 0,说明忽略掩码或行索引有问题。这个测试在修改数据加载或增删类别数之后必跑。
自查三件套跑完了,再去跑完整评估集。最后一点个人习惯:评估协议的缓存。我会把验证集的混淆矩阵、逐类别 IoU 和版本号一起存成 JSON,下次改模型时加载旧结果,对比差异。类别的逐项对比比只看 mIoU 的涨跌可靠得多——一次改动如果只让 mIoU 涨了但某个关键类别掉了,我会先怀疑评估出了问题而不是模型真的变好了。这个习惯帮我避过不少发布前的翻车,希望帮到你。
本文还有配套的精品资源,点击获取