PyTorch语义分割mIoU计算:从混淆矩阵到多卡同步的工程实现 简介这份资源聚焦语义分割任务中的精度评估环节面向正在使用 PyTorch 训练分割模型、需要量化各类别表现的学生与算法工程师。它解决的是「模型整体指标尚可、却难以定位具体类别分割效果」的问题通过脚本演示如何从预测结果出发逐类计算 mIoU 并汇总整个测试集的 mIoU帮助读者建立从预测图到评估指标的完整链路。压缩包共 2 个文件均为 py 脚本体积约 4KB其中一个负责生成 8 位预测图另一个负责与 mask 比对并输出各类别及整体 mIoU 值两者在 PyCharm 中直接运行即可。目前已有 7160 人学习下载说明该评估思路在分割实践中具有较高参考价值。读者可借此掌握 mIoU 的计算流程与实现细节理解预测图与标注 mask 的对应关系并可将脚本迁移到自己的数据集上快速获得可解释的类别级评估结果。1. 语义分割里的 mIoU为什么你的类别得分总是被背景吞掉跑完一个语义分割模型pixel_acc看着挺漂亮0.92、0.94 都有可一到验收环节甲方或者导师盯着mIoU问“小目标那几类怎么才 0.3”这时候才发现自己被像素准确率骗了。语义分割里真正能反映模型对每个类别分割质量的指标是 mIoUmean Intersection over Union它先算每一类的 IoU再对所有类别取平均。问题在于很多开源训练框架默认只打印一个总的 mIoU或者干脆把背景类也算进去导致你根本不知道是哪一类拖了后腿。这篇笔记就围绕 mIoU 的 PyTorch 实现展开从混淆矩阵的构建、逐类 IoU 的计算到多卡、多尺度、忽略标签这些工程细节把一套能直接抄进自己项目里的评估代码拆开讲清楚。适合正在做分割任务、需要自己写验证脚本或者想搞明白框架里mean_iou到底怎么算的从业者。2. 从混淆矩阵到逐类 IoUPyTorch 里的计算链路2.1 为什么不能直接对预测和标签做逐像素比较新手最容易写出的 mIoU 计算是这样的把预测结果和标签展平对每个类别算(pred label c).sum()当作交集再除以(pred c).sum() (label c).sum()。这个写法在单张图上勉强能跑但一旦进入验证循环就会出问题。首先逐像素比较无法处理“预测为 c 但标签是 d”这种交叉情况你算出来的交集只是对角线上的命中数而 IoU 的分母需要的是预测为 c 的全体像素加上标签为 c 的全体像素再减去交集。其次这种写法每张图都要重新遍历所有类别验证集大了以后速度会明显拖慢。常见做法是先把整个验证集的预测结果和标签累积到一个num_classes × num_classes的混淆矩阵里矩阵的[i, j]位置表示“标签为 i、预测为 j”的像素总数。有了这个矩阵每一类的 IoU 就是矩阵[i, i] / (矩阵[i, :].sum() 矩阵[:, i].sum() - 矩阵[i, i])。这样只需要在验证循环里更新矩阵最后一次性算指标既准确又快。2.2 混淆矩阵的 PyTorch 实现下面这段代码是我在多个分割项目里反复用过的混淆矩阵更新逻辑核心是用bincount把二维的标签-预测对映射到一维索引上。import torch def update_confusion_matrix(conf_matrix, pred, label, num_classes, ignore_index255): conf_matrix: (num_classes, num_classes) 的累积矩阵 pred: 模型输出经过 argmax 后的预测形状 (N, H, W) 或 (N,) label: 对应的标签形状与 pred 一致 num_classes: 类别数 ignore_index: 需要忽略的标签值常见做法是 255 # 把 pred 和 label 展平成一维 pred pred.flatten() label label.flatten() # 过滤掉 ignore_index 的像素这些像素不参与任何类别的统计 valid_mask (label ! ignore_index) pred pred[valid_mask] label label[valid_mask] # 把 (label, pred) 组合成一个唯一索引label * num_classes pred # 这样每个组合落在 [0, num_classes*num_classes) 区间内 indices label * num_classes pred # 用 bincount 统计每个组合出现的次数minlength 保证长度固定 count torch.bincount(indices, minlengthnum_classes * num_classes) # 累加到混淆矩阵reshape 回二维 conf_matrix count.reshape(num_classes, num_classes).to(conf_matrix.dtype) return conf_matrix逻辑说明bincount是这里的关键它把“标签-预测”对当成一个整数索引来计数避免了 Python 层面的循环。参数上ignore_index通常设为 255因为很多分割数据集用 255 表示边界或未标注区域这些像素如果参与计算会严重拉低某些类的 IoU。minlength必须显式指定否则当某个组合在 batch 里没出现时bincount返回的长度会短于num_classes * num_classesreshape 就会报错。2.3 从混淆矩阵算逐类 IoU 和 mIoU拿到混淆矩阵后逐类 IoU 的计算就是一行的事但有几个边界情况要处理。def compute_iou_from_confusion(conf_matrix): conf_matrix: (num_classes, num_classes) 的累积矩阵 返回: 每一类的 IoU 列表以及 mIoU # 对角线是每一类被正确预测的像素数 intersection torch.diag(conf_matrix) # 每一类的真实像素总数按行求和 label_sum conf_matrix.sum(dim1) # 每一类的预测像素总数按列求和 pred_sum conf_matrix.sum(dim0) # 并集 真实总数 预测总数 - 交集 union label_sum pred_sum - intersection # 防止除零某一类在验证集里完全没出现时union 为 0 # 常见做法是给 union 加一个极小值或者把这类标记为 nan 后跳过 iou intersection / (union 1e-10) # 只对在验证集中出现过的类别求平均避免 nan 拉低 mIoU valid label_sum 0 miou iou[valid].mean().item() return iou.tolist(), miou参数说明valid这个掩码很重要。如果验证集里某个类别一个像素都没有它的label_sum为 0union也为 0算出来的 IoU 是 0 或者 nan。直接对所有类别取平均会把这种“不存在的类”也算进去导致 mIoU 偏低。我一般会同时打印valid为 True 的类别数和总类别数方便判断验证集是否覆盖了所有类。2.4 在验证循环里串起来把上面两个函数放进验证循环整体结构大概是这样def validate(model, dataloader, num_classes, device): model.eval() conf_matrix torch.zeros(num_classes, num_classes, dtypetorch.int64, devicedevice) with torch.no_grad(): for images, labels in dataloader: images images.to(device) labels labels.to(device) outputs model(images) # 输出通常是 (N, C, H, W)取 argmax 得到 (N, H, W) preds outputs.argmax(dim1) conf_matrix update_confusion_matrix( conf_matrix, preds, labels, num_classes, ignore_index255 ) iou_list, miou compute_iou_from_confusion(conf_matrix) return iou_list, miou这里有个容易忽略的点conf_matrix的 dtype 用int64因为bincount返回的是 int64如果矩阵是 float 类型累加时会有精度损失。另外如果用了多卡验证每个卡上的混淆矩阵需要先all_reduce求和再算指标否则你拿到的是单卡的局部结果。3. 多卡、多尺度与忽略标签工程实现里的参数细节3.1 多卡验证时混淆矩阵怎么同步用DataParallel或者DistributedDataParallel做验证时每张卡各自维护一个混淆矩阵验证结束后需要跨卡求和。常见做法是用torch.distributed.all_reduceimport torch.distributed as dist def sync_confusion_matrix(conf_matrix): 在多卡环境下把各卡的混淆矩阵求和返回同步后的矩阵 if dist.is_available() and dist.is_initialized(): # 确保矩阵在 GPU 上且为连续内存 conf_matrix conf_matrix.cuda() dist.all_reduce(conf_matrix, opdist.ReduceOp.SUM) return conf_matrix注意all_reduce是原地操作调用后所有卡的conf_matrix都会变成全局求和的结果。如果你只想在 rank 0 上算指标可以在all_reduce之后判断dist.get_rank() 0再执行compute_iou_from_confusion。另外all_reduce要求张量在 GPU 上如果混淆矩阵一直放在 CPU 上需要先.cuda()否则会报错或者走很慢的 CPU 通信路径。3.2 多尺度推理对 mIoU 的影响很多分割模型在验证时会做多尺度推理multi-scale inference把输入缩放到不同尺寸分别预测再融合。这个操作对 mIoU 的影响不是线性的我实测下来有几个经验点缩放尺度一般选[0.75, 1.0, 1.25, 1.5]这几档再大对小目标反而有害融合方式用平均概率图再 argmax比直接对预测结果投票更稳多尺度会让验证时间增加 3 到 5 倍如果只是日常跑验证单尺度就够了只有在最终汇报或者对比实验时才值得开。def multi_scale_inference(model, image, scales[0.75, 1.0, 1.25, 1.5]): image: (1, C, H, W) 的单张输入 返回: 融合后的预测 (1, H, W) _, _, H, W image.shape prob_sum None for scale in scales: new_h, new_w int(H * scale), int(W * scale) scaled torch.nn.functional.interpolate( image, size(new_h, new_w), modebilinear, align_cornersFalse ) with torch.no_grad(): out model(scaled) # 把输出缩回原尺寸 out torch.nn.functional.interpolate( out, size(H, W), modebilinear, align_cornersFalse ) prob torch.softmax(out, dim1) prob_sum prob if prob_sum is None else prob_sum prob return (prob_sum / len(scales)).argmax(dim1)参数上align_cornersFalse是分割任务里的常见选择和训练时的 resize 方式保持一致否则插值带来的偏移会累积到 mIoU 上。modebilinear对概率图融合足够如果追求更平滑可以用bicubic但速度会慢一些。3.3 ignore_index 和类别权重的取舍ignore_index设成 255 是很多数据集的惯例但有些数据集用 0 表示背景、用 255 表示忽略区域这时候如果你把ignore_index设成 0背景类就整个被丢掉了mIoU 会虚高。我一般会在数据加载阶段先统计一下标签里各个值的分布确认哪些值是真实类别、哪些是忽略标记。另一个坑是类别权重。有些框架在算 mIoU 时会按类别频率加权但标准的 mIoU 是不加权的每一类权重相同。如果你在论文或者报告里写的是 mIoU就不要偷偷加权否则和别人的结果没法比。如果确实想关注小目标可以额外报告frequency weighted IoU或者per-class IoU表格把每一类的数值都列出来。4. 避坑与排查mIoU 算不对时先查这五处4.1 现象mIoU 比框架自带指标高出一大截原因最常见的是ignore_index没设对把本该忽略的边界像素算进了某些类的分母里或者反过来把背景类整个排除了。另一个可能是验证集里某些类根本没出现你的代码把这些类的 IoU 当 0 算进平均而框架可能直接跳过了。解决先打印混淆矩阵的对角线和每类的label_sum确认哪些类有像素、哪些类没有。然后检查ignore_index的值是否和数据集的标注约定一致。如果框架用的是mmsegmentation或者segmentation_models_pytorch直接看它mean_iou的实现里ignore_index默认值是多少和自己的代码对齐。4.2 现象单卡验证正常多卡验证 mIoU 掉点原因多卡时每张卡的混淆矩阵没有同步你拿到的只是 rank 0 上的局部统计。如果验证集划分不均匀rank 0 上的类别分布和全局差异很大算出来的 mIoU 自然对不上。解决在验证循环结束后、算指标之前对所有卡的混淆矩阵做all_reduce求和。注意all_reduce之后所有卡上的矩阵都是全局值如果只想在 rank 0 上打印加一个 rank 判断即可。4.3 现象mIoU 在训练过程中波动很大原因验证集的 batch size 太小或者验证时没有用model.eval()导致 BatchNorm 层还在用 batch 统计量。另外如果验证时开了数据增强比如随机裁剪每次验证看到的样本分布不一样指标也会抖。解决验证时固定model.eval()和torch.no_grad()验证集的预处理只保留 resize 和 normalize去掉随机翻转、随机裁剪这些增强。如果显存够把验证 batch size 调到 4 或 8让每批覆盖更多类别。4.4 现象某一类的 IoU 始终为 0原因可能是这一类在验证集里样本极少模型完全没学到也可能是标签映射错了比如训练时类别 3 是“车”验证时类别 3 变成了“人”混淆矩阵的对角线自然对不上。解决先确认训练和验证的类别映射表是同一份。然后单独把这一类的小图拿出来可视化看模型到底预测成了什么。如果确实是样本太少考虑在损失函数里给这一类加权重或者用过采样。4.5 现象混淆矩阵累加后数值溢出原因用了int32或者float16来存混淆矩阵验证集大了以后像素总数超过类型上限。int32上限约 21 亿一般分割验证集不至于超但如果做了多尺度融合或者测试时增强累加次数会翻好几倍。解决混淆矩阵统一用int64显存占用可以忽略不计。如果用了混合精度训练注意bincount的输出是 int64累加时不要被自动转成 float16。5. 进阶技巧把 per-class IoU 做成可追踪的验证报告5.1 用表格记录每一轮的逐类 IoU只盯着一个 mIoU 数字很难判断模型到底在哪些类上进步了。我习惯在验证脚本里把每一类的 IoU 存成一个字典训练结束后导出成表格方便对比不同 epoch 或者不同模型。import pandas as pd def build_iou_report(iou_list, class_names, epoch): iou_list: 每一类的 IoU 列表 class_names: 类别名称列表长度与 iou_list 一致 epoch: 当前轮次 df pd.DataFrame({ epoch: [epoch] * len(class_names), class: class_names, iou: iou_list }) return df # 假设有 5 个类别 class_names [background, road, car, person, tree] all_reports [] for epoch in range(num_epochs): # ... 训练和验证 ... iou_list, miou validate(model, val_loader, num_classes, device) report build_iou_report(iou_list, class_names, epoch) all_reports.append(report) final_report pd.concat(all_reports, ignore_indexTrue) final_report.to_csv(iou_per_class.csv, indexFalse)这个表格跑完训练后直接丢进 Excel 或者 pandas 里画折线图哪一类在涨、哪一类在跌一目了然。参数上class_names的顺序必须和混淆矩阵的行列顺序严格一致否则表格里的类别名会对错行。5.2 用 per-class IoU 反推数据问题有一次我跑一个道路分割的任务road类的 IoU 一直在 0.6 左右上不去但car和person都正常。把 per-class IoU 表格拉出来看发现road的label_sum比预期少了一个数量级。回头查数据加载发现标注图里road的像素值被错误地映射成了 255也就是被当成忽略区域丢掉了。修正映射后road的 IoU 直接跳到 0.85。这件事之后我每次拿到新数据集第一件事就是统计标签里各个值的像素占比和类别定义表对一遍确认没有类别被误标成忽略值。5.3 验证脚本的检查清单在把验证脚本交给别人用之前我会强制走一遍下面这几项检查项确认内容ignore_index和数据集标注约定一致常见为 255混淆矩阵 dtypeint64避免累加溢出多卡同步验证结束后 all_reduce 求和类别映射训练和验证使用同一份类别表空类别处理验证集中未出现的类不参与 mIoU 平均可视化抽查随机抽 3 张图叠加预测和标签对比这套流程跑下来mIoU 的数值基本不会出现“玄学”波动。从那以后我每次新写一个分割验证脚本都会先把混淆矩阵打印出来确认对角线上的数值分布合理再去看 mIoU 的数字。希望帮到你。本文还有配套的精品资源点击获取