075、多尺度推理与 TTA 源码：测试时增强的 Flip和Scale 先求平均再 NMS 的代码实现-编程阁

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 的框，如果直接 NMS，0.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_flip=True# 是否水平翻转

为什么选 0.83、1.0、1.2？经验值。0.83 约等于 5/6，1.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_scale=letterbox(im,new_shape,stride=32,auto=False)[0]# 推理pred=model(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_flip=im_scale.flip(-1)# 别写成 flip(0)，那是垂直翻转pred_flip=model(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_center，y_center 和宽高不变。如果你用 xyxy 格式，需要同时调整 x1 和 x2，容易算错。

先求平均再 NMS：核心逻辑

# 收集所有尺度和翻转的结果all_preds=[]forpredin[pred_scale1,pred_scale2,pred_scale3,pred_flip1,...]:# 过滤低置信度框，这里阈值设低一点，因为后面要平均pred=pred[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_preds=torch.cat(cls_preds,dim=0)# 如果该类别没有检测框，跳过ifcls_preds.shape[0]==0:continue# 用 NMS 合并重复框（不同尺度的检测）# 注意：这里 NMS 的 IoU 阈值要设高一点，比如 0.5keep=nms(cls_preds[:,:4],cls_preds[:,4],iou_thres=0.5)cls_preds=cls_preds[keep]# 对保留的框求平均（坐标和置信度）# 这里踩过坑：直接 mean 会丢失框的多样性# 正确做法：用加权平均，置信度高的框权重更大weights=cls_preds[:,4]/cls_preds[:,4].sum()avg_box=(cls_preds[:,:4]*weights.unsqueeze(1)).sum(dim=0)avg_conf=cls_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_preds=torch.stack(final_preds)# 这里 iou_thres 用常规值，比如 0.45keep=nms(final_preds[:,:4],final_preds[:,4],iou_thres=0.45)final_preds=final_preds[keep]