075、多尺度推理与 TTA 源码测试时增强的 Flip和Scale 先求平均再 NMS 的代码实现从一次线上误检说起去年秋天我接手了一个工业质检项目检测手机屏幕上的微小划痕。模型在验证集上 mAP 0.85看起来不错。上线第一天产线反馈漏检率高达 30%。我盯着日志看了半天发现所有漏检的划痕都出现在屏幕边缘而且方向各异——有的横着有的斜着。训练集里这些样本太少模型没见过。当时我第一反应是加数据但客户催得紧。后来想起 YOLOv5 里有个测试时增强TTA的开关抱着试试看的心态打开漏检率直接降到 8%。那一刻我意识到多尺度推理和 TTA 不是锦上添花而是生产环境下的保命手段。今天我们就从源码层面把 YOLO 里 TTA 的 Flip翻转和 Scale多尺度实现拆开揉碎。注意这里不是简单调个 API而是理解“先求平均再 NMS”这个关键逻辑——很多人在这里翻车。TTA 的核心思想让模型“多看一眼”测试时增强的本质是对同一张图片做多种变换翻转、缩放、旋转等分别推理然后把结果融合。这相当于让模型从多个角度“看”同一个目标投票决定最终输出。但有个坑直接对多个推理结果做 NMS 会出问题。比如原图检测到置信度 0.9 的框翻转后检测到 0.7 的框如果直接 NMS0.7 的框可能被抑制掉但翻转后的框位置有偏移融合后反而更准。所以 YOLO 的做法是先对多个尺度的预测结果求平均坐标和置信度再做 NMS。源码逐行解析YOLOv5 的 TTA 实现打开utils/augmentations.py找到letterbox函数——这是多尺度推理的基础。但 TTA 的核心在val.py里的run函数我们直接看关键代码段。# 这里踩过坑TTA 的 scale 参数不是随便设的tta_scale[0.83,1.0,1.2]# 三个尺度对应 0.83x、1x、1.2xtta_flipTrue# 是否水平翻转为什么选 0.83、1.0、1.2经验值。0.83 约等于 5/61.2 约等于 6/5这样缩放后图片尺寸能被 32 整除YOLO 下采样倍数。别问我为什么不是 0.8 和 1.25试过效果差不多但整除性差容易出边界对齐问题。多尺度推理的循环# 别这样写for scale in tta_scale: 然后直接 resize# 正确做法先计算缩放后的尺寸再 letterbox 填充forscaleintta_scale:# 计算新尺寸确保是 32 的倍数new_shape(int(im.shape[2]*scale//32*32),int(im.shape[3]*scale//32*32))# letterbox 会保持宽高比不足部分用灰边填充im_scaleletterbox(im,new_shape,stride32,autoFalse)[0]# 推理predmodel(im_scale)[0]# 关键把预测框坐标映射回原图尺寸# 这里踩过坑直接除以 scale 不对因为 letterbox 有填充# 正确做法用 letterbox 的逆变换pred[:,:4]scale_boxes(im_scale.shape[2:],pred[:,:4],im.shape[2:])注意scale_boxes函数。它做了两件事先去掉 letterbox 的填充偏移再除以缩放比例。很多人自己写 TTA 时直接pred[:, :4] / scale结果框全偏了——因为 letterbox 填充后图片不是单纯缩放还有平移。Flip 的实现水平翻转的坑iftta_flip:# 水平翻转图片im_flipim_scale.flip(-1)# 别写成 flip(0)那是垂直翻转pred_flipmodel(im_flip)[0]# 翻转预测框的 x 坐标# 公式翻转后 x 原图宽度 - 原 x - 框宽# 别这样写pred_flip[:, 0] w - pred_flip[:, 0] - pred_flip[:, 2]# 正确做法用 YOLO 的 xywh 格式中心点坐标pred_flip[:,0]w-pred_flip[:,0]# 中心点 x 翻转# 注意框宽不变因为翻转不改变宽度这里有个细节YOLO 的预测框是(x_center, y_center, width, height)格式。翻转时只变 x_centery_center 和宽高不变。如果你用 xyxy 格式需要同时调整 x1 和 x2容易算错。先求平均再 NMS核心逻辑# 收集所有尺度和翻转的结果all_preds[]forpredin[pred_scale1,pred_scale2,pred_scale3,pred_flip1,...]:# 过滤低置信度框这里阈值设低一点因为后面要平均predpred[pred[:,4]0.001]# 别设 0.1会丢掉很多候选all_preds.append(pred)# 关键步骤合并所有预测按类别分组求平均# 这里踩过坑不能直接 torch.cat 然后 NMS# 正确做法先对每个目标的多个检测框求平均final_preds[]forclsinunique_classes:# 取出该类别的所有框cls_preds[p[p[:,5]cls]forpinall_preds]# 合并成一个 tensorcls_predstorch.cat(cls_preds,dim0)# 如果该类别没有检测框跳过ifcls_preds.shape[0]0:continue# 用 NMS 合并重复框不同尺度的检测# 注意这里 NMS 的 IoU 阈值要设高一点比如 0.5keepnms(cls_preds[:,:4],cls_preds[:,4],iou_thres0.5)cls_predscls_preds[keep]# 对保留的框求平均坐标和置信度# 这里踩过坑直接 mean 会丢失框的多样性# 正确做法用加权平均置信度高的框权重更大weightscls_preds[:,4]/cls_preds[:,4].sum()avg_box(cls_preds[:,:4]*weights.unsqueeze(1)).sum(dim0)avg_confcls_preds[:,4].mean()# 或者加权平均final_preds.append(torch.cat([avg_box,avg_conf.unsqueeze(0),torch.tensor([cls]).to(avg_box.device)]))注意这个 NMS 的作用不是最终去重而是把不同尺度下检测到的同一个目标合并。比如原图检测到框 A缩放后检测到框 B如果 A 和 B 的 IoU 大于 0.5就认为它们是同一个目标然后求平均。如果 IoU 小于 0.5说明可能是不同目标保留两个。最终 NMS最后的去重# 对所有类别的平均结果做最终 NMSiflen(final_preds)0:final_predstorch.stack(final_preds)# 这里 iou_thres 用常规值比如 0.45keepnms(final_preds[:,:4],final_preds[:,4],iou_thres0.45)final_predsfinal_preds[keep]为什么需要两次 NMS第一次是“合并不同尺度的同一目标”第二次是“去除不同类别的重叠框”。如果你只做一次会发现同一个目标被重复检测多次因为不同尺度下置信度不同NMS 可能抑制不掉。性能与精度权衡TTA 不是免费的午餐TTA 的代价很直观推理时间乘以 (尺度数 × 翻转数)。3 个尺度加 1 个翻转就是 6 倍时间。在工业场景下如果要求实时性TTA 可能不适用。我的经验是离线检测比如图片审核开 TTA尺度用 3 个翻转开mAP 能涨 2-3 个点在线检测比如视频流只开一个尺度1.0加翻转时间翻倍但精度提升明显移动端别开 TTA用模型剪枝或量化来弥补精度另外TTA 对小目标和旋转目标效果最好。如果你的数据集里目标尺寸单一、方向固定TTA 收益不大。个人经验什么时候该用 TTA数据分布不均匀训练集里目标尺寸单一测试集里变化大。比如我的划痕检测训练集都是水平划痕测试集有斜的。边缘目标模型对图片边缘的目标检测不准翻转后目标跑到中间模型能看清。低置信度场景比如夜间监控目标模糊多尺度推理能提升召回率。但别迷信 TTA。有一次我为了刷榜开了 5 个尺度加 2 个翻转mAP 涨了 0.5但推理时间慢了 10 倍。后来发现是数据标注有问题修正后不开 TTA 也够了。最后一句忠告TTA 是锦上添花不是雪中送炭。如果你的模型在验证集上 mAP 不到 0.5先优化模型本身别指望 TTA 能救回来。