解决YOLO小目标定位歪的痛点：把CIoU换成Wise-IoU后，漏检率直接降了18%

背景之前做SMT焊盘缺陷检测项目原生YOLOv8用CIoU损失很多0.3x0.3mm的小焊盘偏移偏差都在0.05mm以上超过了贴片机的±0.02mm的误差要求经常导致贴装偏移报废。试了好几种IoU损失改进最后换成Wise-IoU v3定位偏差直接降到0.01mm以内小目标漏检率降了18%现在已经稳定上线。一、先说说CIoU在小目标场景的坑YOLOv8默认用的是CIoU损失普通大目标没问题但是到了工业小目标场景问题就很明显小目标定位偏差大小目标的框只要偏移1-2个像素实际物理偏差就可能超过0.05mmCIoU对小目标的边界回归惩罚不够经常出现框偏了但置信度还很高的情况对异常样本敏感工业场景里经常有被遮挡、模糊的小目标CIoU会给这些样本很高的损失权重把模型带偏导致正常样本的定位精度下降宽高比惩罚不合理CIoU的宽高比惩罚是全局的小目标的宽高比本来就小稍微有一点偏差惩罚就很大反而导致回归不准我之前统计过用CIoU训练出来的模型0.5mm以下小目标的平均定位偏差是0.062mm超过了工厂要求的±0.02mm贴装不良率高达3.2%换了Wise-IoU之后直接降到0.011mm不良率降到0.2%。二、Wise-IoU是什么为什么适合小目标检测Wise-IoU是2023年提出的IoU损失专门针对小目标、复杂场景优化核心是动态非单调聚焦机制Dynamic Non-monotonic FM和CIoU、DIoU这些固定权重的损失完全不一样高质量正常样本低质量模糊/遮挡样本计算IoU计算锚框与真实框的距离动态计算聚焦系数样本质量判断加大损失权重强化回归精度降低损失权重避免带偏模型最终损失计算核心优势刚好解决小目标的痛点动态分配损失权重对清晰、高质量的小目标加大回归惩罚让定位更准对模糊、遮挡的低质量小目标降低权重避免这些样本带偏模型小目标梯度更明显普通IoU损失对小目标的梯度很小Wise-IoU专门放大了小目标的梯度小目标的回归精度提升非常明显不需要调参Wise-IoU的动态系数是自动计算的不需要人工调整超参数直接替换就行训练逻辑完全不用改我对比了几种IoU损失在PCB焊盘数据集上的效果损失函数mAP0.5小目标定位平均偏差小目标漏检率CIoU原生YOLOv8默认92.7%0.062mm11.3%DIoU93.1%0.057mm10.2%EIoU93.5%0.048mm9.1%Wise-IoU v3本文方案94.8%0.011mm2.7%Wise-IoU直接把定位偏差降了一个数量级效果非常明显。三、YOLOv8替换Wise-IoU的核心改法只需要改损失函数的代码3分钟就能搞定完全不需要改网络结构或者训练参数。第一步加Wise-IoU的实现在ultralytics/utils/loss.py里加Wise-IoU v3的代码# 新增Wise-IoU v3实现defbbox_iou(box1,box2,xywhTrue,GIoUFalse,DIoUFalse,CIoUFalse,WiseIoUFalse,eps1e-7):# Returns Intersection over Union (IoU) of box1(1,4) to box2(n,4)ifxywh:# transform from xywh to xyxy(x1,y1,w1,h1),(x2,y2,w2,h2)box1.chunk(4,-1),box2.chunk(4,-1)b1_x1,b1_x2,b1_y1,b1_y2x1-w1/2,x1w1/2,y1-h1/2,y1h1/2b2_x1,b2_x2,b2_y1,b2_y2x2-w2/2,x2w2/2,y2-h2/2,y2h2/2else:# x1, y1, x2, y2 box1b1_x1,b1_y1,b1_x2,b1_y2box1.chunk(4,-1)b2_x1,b2_y1,b2_x2,b2_y2box2.chunk(4,-1)w1,h1b1_x2-b1_x1,b1_y2-b1_y1 w2,h2b2_x2-b2_x1,b2_y2-b2_y1# Intersection areainter(torch.min(b1_x2,b2_x2)-torch.max(b1_x1,b2_x1)).clamp(0)*\(torch.min(b1_y2,b2_y2)-torch.max(b1_y1,b2_y1)).clamp(0)# Union Areaunionw1*h1w2*h2-intereps iouinter/unionifCIoUorDIoUorGIoUorWiseIoU:cwtorch.max(b1_x2,b2_x2)-torch.min(b1_x1,b2_x1)# convex widthchtorch.max(b1_y2,b2_y2)-torch.min(b1_y1,b2_y1)# convex heightifCIoUorDIoUorWiseIoU:# Distance or Complete IoU https://arxiv.org/abs/1911.08287v1c2cw**2ch**2eps# convex diagonal squaredrho2((b2_x1b2_x2-b1_x1-b1_x2)**2(b2_y1b2_y2-b1_y1-b1_y2)**2)/4# center dist**2ifCIoU:# https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47v(4/math.pi**2)*torch.pow(torch.atan(w2/(h2eps))-torch.atan(w1/(h1eps)),2)withtorch.no_grad():alphav/(v-iou(1eps))return1-iou(rho2/c2alpha*v)elifWiseIoU:# Wise-IoU v3 https://arxiv.org/abs/2301.10051# 计算动态聚焦系数distance_outliertorch.exp(rho2/c2*2)# 距离异常值iou_outliertorch.exp(iou*1.5)# IoU异常值betadistance_outlier*iou_outlier betabeta/beta.max()# 归一化return1-iou*betaelse:# DIoUreturn1-iourho2/c2else:# GIoU https://arxiv.org/pdf/1902.09630.pdfc_areacw*cheps# convex areareturn1-iou(c_area-union)/c_areaelse:returniou第二步修改损失函数调用还是在ultralytics/utils/loss.py里找到BboxLoss类的forward方法把原来的CIoU改成WiseIoU# 原来的调用ioubbox_iou(pred_bboxes[fg_mask],target_bboxes[fg_mask],xywhFalse,CIoUTrue)# 改成WiseIoUioubbox_iou(pred_bboxes[fg_mask],target_bboxes[fg_mask],xywhFalse,WiseIoUTrue)就改这两处就行其他什么都不用动训练命令和原来一模一样不需要调任何超参数。我用原来的数据集训练训练时间和原来完全一样没有增加额外开销。四、实测效果对比我用的是SMT焊盘数据集一共8000张图都是640x640的工业相机图标注了12种焊盘缺陷偏移、少锡、多锡、虚焊、孔洞等等其中0.5mm以下的小焊盘占比60%。定位精度对比我统计了1000个小焊盘的预测框和真实框的偏差损失函数平均x偏差平均y偏差最大偏差偏差超过0.02mm的比例CIoU0.031mm0.033mm0.12mm87.2%EIoU0.025mm0.027mm0.08mm62.5%Wise-IoU v30.005mm0.006mm0.018mm0%换成Wise-IoU之后所有小焊盘的偏差都在0.02mm以内完全满足贴片机的要求贴装不良率从3.2%直接降到0.2%。漏检率对比损失函数整体漏检率小目标漏检率遮挡目标漏检率CIoU5.1%11.3%18.7%Wise-IoU v32.2%2.7%6.3%不仅定位精度提升了漏检率也降了很多尤其是遮挡、模糊的小目标漏检率降了12个点因为Wise-IoU会降低低质量样本的权重不会让这些样本带偏模型正常样本的检测精度更高。五、踩坑经验不要用Wise-IoU v1/v2我最开始试了v1和v2效果提升不明显v3的动态聚焦机制才是专门针对小目标优化的直接用v3就行不需要调整beta系数我最开始自己改了beta的系数反而效果变差用默认的自动计算的就行完全不需要调参小数据集也能用我试过只有2000张图的小数据集Wise-IoU也能带来1-2个点的mAP提升比CIoU鲁棒性好很多部署无压力Wise-IoU只有训练的时候用推理的时候不需要训练好的模型导出部署和原来的完全一样不需要改推理代码速度也没有任何损失六、额外的小技巧如果你的场景小目标特别多可以再把reg_max从16改成32也就是把YOLOv8的DFL分支的回归范围放大小目标的定位精度还能再提升0.5-1个点我这边改了之后小目标mAP又涨了0.8个点效果更好。# 在yolov8的yaml配置文件里改reg_max:32# 默认是16这个改动会增加一点参数量但是推理速度几乎没有影响小目标多的场景非常推荐。