训练完UNet只看loss曲线等模型跑完才发现预测效果不行这大概是我见过新手最容易踩的坑。尤其是多类别语义分割loss降得漂亮不代表你的模型真的分得准类别不平衡时更是如此。我刚开始做分割项目的时候也犯过这个错误盯着训练损失从0.8降到0.1沾沾自喜结果一测mIoU只有0.3模型几乎把前景类别全忽略了。后来老老实实把评估指标补上才发现问题出在哪里。这篇东西就围绕PyTorch下训练UNet后如何计算并可视化mIoU来写核心是把原理讲透、把代码给全、把坑都填上适合已经有PyTorch基础、在跑语义分割模型但还没搞清楚怎么科学评估的开发者参考。1. 为什么训练完不能只盯loss这是整篇文章的起点也是我踩坑最多的地方。很多PyTorch初学者训练UNet时习惯性看loss曲线来判定模型的收敛情况认为loss越低代表模型越好。这个思路在分类任务里勉强能用但在多类别语义分割里会出大问题。1.1 loss下降并不代表分割质量提升拿我自己跑过的一个三类别分割任务来举例前景目标特别小背景占了95%以上像素另外两个类别加起来不到5%。用交叉熵损失训练UNet模型很快就把所有像素都预测成背景因为这样做loss也很低——背景类别占绝对主导只要背景预测对了loss就被拉下来了。表现在训练曲线上loss从0.7一路降到0.2曲线平滑又漂亮但模型实际上什么都没学会。Focal loss这类改进损失函数能缓解这个问题因为它会降低易分样本的权重让模型更关注难分的少数类别。但即使用了加权损失仍然需要一个独立的评估指标来客观衡量分割效果。loss是训练过程中的优化目标mIoU才是最终业务上关心的质量指标两者不能直接划等号。1.2 mIoU到底在算什么mIoU的中文全称是平均交并比来自计算机视觉领域的Intersection over Union概念。对任意一个类别CIoU的计算方式是模型预测为C且真实标签也是C的像素数量交集除以模型预测为C或真实标签为C的所有像素数量并集。打个比方你预测了一个区域真实标注也有一个区域两个区域重叠的部分就是预测对的。把预测区域和标注区域合起来看重叠占整体的比例越高说明你的预测越准。这个比例的取值在0到1之间1表示完美重合0表示彻底没对上。多类别语义分割会有多个类别每个类别都能算出一个IoU把所有类别的IoU取平均就是mIoU。这也是国际通用的语义分割评估标准相关论文中提到y的标准也基本都基于这个指标。mIoU的公式可以写成先对第i个类别算出TP、FP、FNIoU_i TP / (TP FP FN)然后mIoU就是对所有类别的IoU_i求平均。需要注意这里用的是类别平均不是像素平均所以类别不平衡被直接暴露出来——小类别的IoU低了整体mIoU会立刻下降而loss可能完全看不出这个问题。2. mIoU计算原理与核心代码实现原理清楚了落地的挑战在于代码怎么写。多类别语义分割的标签形状一般是H×W每个位置是一个类别编号模型输出形状是N×C×H×WN是批量大小C是类别总数。要算mIoU必须把模型输出和标签对齐逐类别统计TP、FP、FN。2.1 核心指标计算逻辑计算过程中有两个关键点第一个是预测标签的提取需要从模型输出的每个通道概率中取最大概率对应的索引作为预测类别第二个是混淆矩阵的构建语义分割中mIoU通过混淆矩阵计算更方便矩阵的每一行代表真实类别每一列代表预测类别矩阵[i][j]表示真实标签为i、预测为j的像素数。有了混淆矩阵后第i类的TP就是矩阵[i][i]FP是第i列所有非对角元素之和FN是第i行所有非对角元素之和。这样逐类计算再平均逻辑清晰、代码简洁也方便扩展新的评估指标。2.2 完整可运行的代码下面是我自用的一个计算函数完整代码可以直接跑输入分别是模型输出logits形状为N×C×H×W、真实标签mask形状为N×H×W、类别数量和需要忽略的类别索引比如背景或边缘输出是每个类别的IoU列表和平均mIoUimport numpy as np import torch def compute_mIoU(pred_logits, true_masks, num_classes, ignore_indexNone): 计算多类别语义分割的mIoU Args: pred_logits: 模型输出形状 (N, C, H, W)每个位置是未经过softmax的logits true_masks: 真实标签形状 (N, H, W)每个位置是类别索引 num_classes: 类别总数包含背景类 ignore_index: 需要忽略的类别索引列表可省略 Returns: iou_list: 每个类别的IoU长度为 num_classes miou: 所有类别IoU的平均值 # 将logits转为预测类别索引 preds torch.argmax(pred_logits, dim1).cpu().numpy() masks true_masks.cpu().numpy() # 初始化混淆矩阵 conf_matrix np.zeros((num_classes, num_classes), dtypenp.int64) # 展平并统计每个像素 preds_flat preds.ravel() masks_flat masks.ravel() # 过滤掉需忽略的像素 if ignore_index is not None: valid_mask ~np.isin(masks_flat, ignore_index) preds_flat preds_flat[valid_mask] masks_flat masks_flat[valid_mask] # 用np.add.at避免索引叠加时精度丢失 np.add.at(conf_matrix, (masks_flat, preds_flat), 1) iou_list [] for cls in range(num_classes): tp conf_matrix[cls, cls] fp conf_matrix[:, cls].sum() - tp fn conf_matrix[cls, :].sum() - tp union tp fp fn if union 0: # 真实标签和预测都不包含该类别时按1处理 iou 1.0 else: iou tp / union iou_list.append(iou) miou np.mean(iou_list) return iou_list, miou这段代码写起来很简单但有三个细节必须注意。第一个细节是np.add.at的使用。直接写成conf_matrix[masks_flat, preds_flat] 1在索引重复时是错的numpy会把重复位置的更新在原数组上直接覆盖只保留最后一次导致统计严重丢失。np.add.at是专门的ufunc.at操作专治这种需要重复累加的索引场景。第二个细节是union为0的情况。某个类别在验证集里压根没出现又没有预测出来说明模型对这类没有感知学术上通常将其IoU计为1因为不存在误判也会计为0或直接跳过。这里我选了1.0理由是未出现类对整体效果的惩罚不应与预测错了却想掩盖等同不同论文实现有差异实际跑分时要注明。第三个细节是torch.argmax是在logits上直接做的。很多人习惯先加softmax再取argmax数学上等价因为softmax是单调函数不会改变最大值位置但少了这步能省下不少显存和计算时间。2.3 基于混淆矩阵的等价写法上面那种逐行逐列统计的方式比较好理解还有一个更紧凑的等价写法对所有类别的TP、FP、FN用矩阵运算一次性算出来适合写进训练代码里加速def compute_iou_from_conf_matrix(conf_matrix): 给定混淆矩阵直接算每个类别的IoU。 tp np.diag(conf_matrix) fp conf_matrix.sum(axis0) - tp fn conf_matrix.sum(axis1) - tp union tp fp fn # union为0时用eps保证数值稳定 iou_list np.where(union 0, tp / np.maximum(union, 1e-10), 1.0) return iou_list, iou_list.mean()我实际在项目里用的是这种写法速度快、代码短需要可视化的类别IoU直接从iou_list里拿即可。等下要做的可视化柱状图数据来源就是它。3. 验证集上的评估流程与训练结合函数写好了下一步是把它接进训练流程。评估mIoU不能随手拿几张图测完就当最终结果那样存在很大的偶然性必须建立标准的评估流程。3.1 验证集评估的规范流程我的实践经验是严格按照以下几个步骤组织评估过程第一步把数据集划分好。常见做法是训练集和验证集按8:2或9:1划分验证集要保证包含所有类别的样本并且与训练集数据分布一致不能出现空间重叠导致的数据泄漏。第二步训练过程中每个epoch结束后在验证集上跑一次评估。把模型切到eval模式关闭梯度计算用当前的模型权重对整个验证集做推理累积每一批的混淆矩阵验证集全部算完后再统一计算mIoU。第三步不仅记录mIoU同时记录每个类别的IoU。这一步很关键只看整体mIoU会掩盖具体哪一类崩了的问题。按这个流程走每个epoch的产物就是一张表列出各类IoU和整体mIoU模型什么时候变好、哪一个类别拖后腿一目了然。下面是我在训练循环中集成mIoU评估的代码片段def evaluate(model, val_loader, num_classes, device, ignore_indexNone): model.eval() conf_matrix np.zeros((num_classes, num_classes), dtypenp.int64) with torch.no_grad(): for images, masks in val_loader: images images.to(device) masks masks.to(device) outputs model(images) preds torch.argmax(outputs, dim1) preds_np preds.cpu().numpy().ravel() masks_np masks.cpu().numpy().ravel() if ignore_index is not None: valid_mask ~np.isin(masks_np, ignore_index) preds_np preds_np[valid_mask] masks_np masks_np[valid_mask] np.add.at(conf_matrix, (masks_np, preds_np), 1) iou_list, miou compute_iou_from_conf_matrix(conf_matrix) model.train() return iou_list, miou调用时在epoch结束处加上记录逻辑保存最佳mIoU的模型权重best_miou 0.0 for epoch in range(epochs): train_one_epoch(model, train_loader, optimizer, criterion) iou_list, miou evaluate(model, val_loader, num_classes, device) print(fEpoch {epoch}: mIoU {miou:.4f}) for i, iou in enumerate(iou_list): print(f Class {i}: IoU {iou:.4f}) if miou best_miou: best_miou miou torch.save(model.state_dict(), best_unet.pth)这里有个容易忽略的细节评估结束后记得调回model.train()模式否则接下来训练时BatchNorm和Dropout仍在推理模式训练会莫名其妙变慢甚至效果下降。我见过有同事漏掉这行排查了半天才发现是模型没切回训练模式。3.2 训练曲线与mIoU曲线的对照价值把loss曲线和mIoU曲线画在同一张图里能看出很多名堂。最常见的情况是loss持续下降但mIoU停滞说明模型在过拟合或类别不平衡问题严重。训练集上的loss降了但验证集上的泛化能力没跟上mIoU就是最灵敏的警报器。另一种情况是loss和mIoU同时上升某些正则化方法下会出现说明模型可能陷入局部最优解需要调整学习率或优化器参数。我把这两条曲线叠加展示的经验是loss曲线适合看收敛状态mIoU曲线适合看模型质量两个一起看比只看一个要靠谱得多。之前提到的某个项目就是靠mIoU曲线发现了loss下降但指标不涨的异常现象及时切了Focal loss和加权采样才救回来。4. 可视化让评估结果一目了然数值列表能帮助判断但视觉效果更直观。mIoU可视化有三个层次类别IoU的柱状图、混淆矩阵热力图、预测结果的对照图。4.1 类别IoU柱状图这个是基础每个类别画一根柱子一眼看出哪个类别最难分。import matplotlib.pyplot as plt def plot_iou_bar(iou_list, class_names): plt.figure(figsize(10, 6)) x range(len(class_names)) plt.bar(x, iou_list, color#4C72B0) plt.xticks(x, class_names) plt.ylim(0, 1.0) plt.ylabel(IoU) plt.title(Per-class IoU on Validation Set) for i, v in enumerate(iou_list): plt.text(i, v 0.02, f{v:.3f}, hacenter, fontsize9) plt.tight_layout() plt.savefig(per_class_iou.png, dpi150) plt.show()柱子上标注具体数值很重要国际论文里的IoU表也惯例保留三位小数方便其他研究者复现对比。4.2 混淆矩阵热力图柱状图告诉你哪类差混淆矩阵告诉你差在哪。模型是把A类误分成B类了还是把B类误分成A类了从混淆矩阵上能直接看出来。import seaborn as sns def plot_confusion_matrix(conf_matrix, class_names, save_pathconfusion_matrix.png): plt.figure(figsize(8, 7)) sns.heatmap(conf_matrix, annotTrue, fmtd, cmapBlues, xticklabelsclass_names, yticklabelsclass_names) plt.xlabel(Predicted) plt.ylabel(Ground Truth) plt.tight_layout() plt.savefig(save_path, dpi150) plt.show()使用混淆矩阵时要注意当类别数量较多时热力值的绝对数量差异很大一个背景类别可能占了近百万像素一个细长类别只有几千。直接看原始数字往往小的类别颜色很淡。解决办法是对矩阵每行做归一化重点关注行方向的误分模式norm_conf conf_matrix / (conf_matrix.sum(axis1, keepdimsTrue) 1e-10)归一化后的值表示真实标签为i的像素有多大比例被预测为j这样处理比较能突出细小的误分类问题。4.3 预测结果可视化评估指标之外真正直观的还是图像。把原始图、真实标注、模型预测三张图并排放视觉上判断模型什么地方边界模糊、哪里漏检比任何数值都来得快。def visualize_prediction(image, mask, prediction, class_names, save_pathNone): fig, axes plt.subplots(1, 3, figsize(15, 5)) axes[0].imshow(image) axes[0].set_title(Image) mask_vis np.zeros((mask.shape[0], mask.shape[1], 3), dtypenp.uint8) for cls, color in enumerate(class_names): mask_vis[mask cls] color axes[1].imshow(mask_vis) axes[1].set_title(Ground Truth) pred_vis np.zeros((prediction.shape[0], prediction.shape[1], 3), dtypenp.uint8) for cls, color in enumerate(class_names): pred_vis[prediction cls] color axes[2].imshow(pred_vis) axes[2].set_title(Prediction) for ax in axes: ax.axis(off) plt.tight_layout() if save_path: plt.savefig(save_path, dpi150) plt.show()这里class_names建议直接定义为RGB颜色列表比如背景是黑色[0,0,0]道路是白色[255,255,255]建筑是蓝色[0,0,255]这样可视化时直接映射清晰也好看。4.4 可视化结果驱动模型迭代多次实践下来我发现可视化不只是为了写报告展示对调试也有很大价值。举个例子如果某类别IoU低且混淆矩阵显示它和背景互相混淆那么大概率是前景目标太小、边界不清晰可以考虑在loss上给这个小类加权。如果预测结果整体偏模糊、边界有一圈空隙则可能是上采样层的问题需要检查UNet的跳跃连接是否有信息丢失。另一个实用的技巧是保存困难样本也就是真实标注和预测差距最大的那几张图。我在每个epoch评估后都会顺手存几份当前epoch最差的预测图训练完一翻整个模型演进的脉络都清楚了初期为什么误检、后期如何改正一目了然。5. 常见问题与排查技巧实录mIoU计算与可视化过程中我自己反复踩过不少坑这里整理成速查表每一项都是实际发生过的问题。5.1 典型问题速查表问题现象可能原因解决方案mIoU数值固定在很低范围如0.1~0.2类别极度不平衡模型偏向背景使用类别权重、Focal loss或对少数类做过采样mIoU高但实际效果肉眼看起来差验证集与测试集分布不一致检查数据划分确保类别分布一致某个类别的IoU总是0该类别样本没进验证集或类别索引错位确认数据集类别定义和模型输出类别数一一对应loss正常下降但mIoU震荡剧烈学习率过大验证集样本量太少降低学习率扩充验证集混淆矩阵里对角线很亮但某一行整体偏暗该类别像素量太少统计意义不足对该类别做可视化检查确认标注质量计算mIoU时程序报维度错误预测输出和标签的H、W尺寸不一致确认Unet的输出尺寸与原图一致必要时插值回原尺寸两个类别的颜色在可视化里几乎一样class_names颜色定义相近使用对比明显的颜色列表5.2 类别索引不对齐的问题一个很典型的坑是类别编号错位。常见的数据集标注文件里可能有0、1、2三个类别但模型输出是4通道多出来的那个通道永远预测不到任何正确像素mIoU全程会被拖低。排查方法很简单生成一次混淆矩阵把每行每列总和分别打出来看看哪个类别完全没有像素标记或完全没有预测基本就是索引对齐的问题。5.3 忽略不想评估的类别某些项目里会有不想评估的类别比如标注质量很差的边缘、或者标注的未知类别可以给compute_mIoU传入ignore_index[0]把背景类从计算中剔除。这样算出来的mIoU是除背景外全部类别的平均IoU有时候更贴合业务需求。如果不加这个参数背景类占大头小类别的表现容易被稀释。5.4 resize带来的像素错位问题在语义分割里很多UNet为了提高分辨率会做上采样但如果在数据预处理时把图片resize了、标签没同步resize或者resize的插值方式不一致预测结果与标签在边界处会有半像素级错位mIoU会在边界部分被低估。我自己一次漏了给标签做nearest插值结果mIoU直接掉了0.15训练集上完全没发现就是从评估阶段查出来的。所有数据增强操作必须对图像和标签使用同样的几何变换而且标签的插值方式必须用最近邻不能用双线性否则类别边界会出现不存在的类别编号。5.5 数值类型与设备问题还有一个容易忽略的问题是数据类型的差异。PyTorch里标签通常是LongTensor模型输出是FloatTensor直接用torch.argmax没问题但是如果标签从文件里读出来没有转成long或者经过了torch.squeeze之类的操作导致维度多了一维混淆矩阵的索引就会出错。我建议在compute_mIoU函数入口处统一做一次true_masks.long()省去一堆排查时间。类不均衡时类别IoU统计也需要小心。有些数量极少的类别训练集和验证集分布稍有偏差模型可能倾向于全预测成背景该类别IoU为0。此时不能只调loss得考虑给验证集做类别均衡抽样或者用加权mIoU按每个类别的像素占比给IoU加权重这样对少数类更友善。5.6 训练过程中的动态监控技巧最后分享一个我常用来监控模型状态的技巧在训练过程中每N个iteration就在一小批验证样本上快速算一次mIoU不需要跑全部验证集只用所有类别的样本各取一两张拼成一个小batch。这样能从训练日志里直接看到mIoU的动态走向避免跑完整个epoch才发现模型早就发散。这个小batch mIoU不适合作为最终报告数据因为样本太少统计噪声大但作为训练过程的参考手环完全够用。正式评估还是要按第3节流程对整个验证集做一次。结尾的几句大实话说真的我在实际项目里最大的教训是评估指标不是训练结束后的附加题而是模型开发过程中的导航仪。如果从一开始就把mIoU计算和可视化嵌入训练流程很多模型问题在发生初期就能看出来而不是等训练完做了一大堆可视化才发现方向错了。另外还要提醒一点mIoU是一个需要明确约定细节的指标——含不含背景、忽略哪些标签、union为0时怎么处理、用什么插值方式统一尺寸所有这些都会影响最终数值的可比性写报告或跟别人对比模型时一定要把这些条件写清楚。如果后续还想继续深入有几个可以扩展的方向一是把mIoU之外的边界指标如Boundary IoU加进来二是针对小目标类别引入类别加权IoU三是把整个评估流程写成独立的评估脚本每次训练完自动生成一份包含所有指标的评估报告。这些内容等下次有机会我再专门写一篇展开聊。