测试时增强 (TTA)

📚 本指南解释了如何在测试和推理过程中使用测试时增强（TTA）来提高YOLOv5的mAP和召回率 🚀。

开始之前

克隆仓库并在Python>=3.8.0环境中安装requirements.txt，包括PyTorch>=1.8。模型和数据集会自动从最新的YOLOv5发布版本下载。

git clone https://github.com/ultralytics/yolov5  # 克隆
cd yolov5
pip install -r requirements.txt  # 安装

正常测试

在尝试TTA之前，我们希望建立一个基准性能以进行比较。此命令在图像大小为640像素的COCO val2017上测试YOLOv5x。yolov5x.pt是最大且最准确的可用模型。其他选项包括yolov5s.pt、yolov5m.pt和yolov5l.pt，或者您自己从自定义数据集训练的检查点./weights/best.pt。有关所有可用模型的详细信息，请参阅我们的README表格。

python val.py --weights yolov5x.pt --data coco.yaml --img 640 --half

输出:

val: data=./data/coco.yaml, weights=['yolov5x.pt'], batch_size=32, imgsz=640, conf_thres=0.001, iou_thres=0.65, task=val, device=, single_cls=False, augment=False, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=True, project=runs/val, name=exp, exist_ok=False, half=True
YOLOv5 🚀 v5.0-267-g6a3ee7c torch 1.9.0+cu102 CUDA:0 (Tesla P100-PCIE-16GB, 16280.875MB)

Fusing layers...
Model Summary: 476 layers, 87730285 parameters, 0 gradients

val: Scanning '../datasets/coco/val2017' images and labels...4952 found, 48 missing, 0 empty, 0 corrupted: 100% 5000/5000 [00:01<00:00, 2846.03it/s]
val: New cache created: ../datasets/coco/val2017.cache
               Class     Images     Labels          P          R     mAP@.5 mAP@.5:.95: 100% 157/157 [02:30<00:00,  1.05it/s]
                 all       5000      36335      0.746      0.626       0.68       0.49
Speed: 0.1ms pre-process, 22.4ms inference, 1.4ms NMS per image at shape (32, 3, 640, 640)  # <--- 基准速度

Evaluating pycocotools mAP... saving runs/val/exp/yolov5x_predictions.json...
...
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.504  # <--- 基准mAP
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.688
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.546
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.351
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.551
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.644
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.382
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.628
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.681  # <--- 基准mAR
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.524
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.735
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.826

使用TTA进行测试

在任何现有的val.py命令后附加--augment以启用TTA，并将图像大小增加约30%以获得更好的结果。请注意，启用TTA的推理通常会比正常推理多花费2-3倍的时间，因为图像会进行左右翻转并在3种不同的分辨率下处理，然后在NMS之前合并输出。速度下降的部分原因是图像尺寸变大（832 vs 640），部分原因是实际的TTA操作。

python val.py --weights yolov5x.pt --data coco.yaml --img 832 --augment --half

输出:

val: data=./data/coco.yaml, weights=['yolov5x.pt'], batch_size=32, imgsz=832, conf_thres=0.001, iou_thres=0.6, task=val, device=, single_cls=False, augment=True, verbose=False, save_txt=False, save_hybrid=False, save_conf=False, save_json=True, project=runs/val, name=exp, exist_ok=False, half=True
YOLOv5 🚀 v5.0-267-g6a3ee7c torch 1.9.0+cu102 CUDA:0 (Tesla P100-PCIE-16GB, 16280.875MB)

融合层...
/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:718: UserWarning: 命名张量及其所有相关API是实验性功能，可能会发生变化。请不要在重要场合使用它们，直到它们被发布为稳定版本。（在 /pytorch/c10/core/TensorImpl.h:1156. 内部触发）
  return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
模型摘要: 476层, 87730285参数, 0梯度
val: 扫描 '../datasets/coco/val2017' 图像和标签...找到4952个, 缺失48个, 空0个, 损坏0个: 100% 5000/5000 [00:01<00:00, 2885.61it/s]
val: 新缓存已创建: ../datasets/coco/val2017.cache
               类别     图像数     标签数          P          R     mAP@.5 mAP@.5:.95: 100% 157/157 [07:29<00:00,  2.86s/it]
                 全部      5000      36335      0.718      0.656      0.695      0.503
速度: 预处理0.2ms, 推理80.6ms, NMS 2.7ms每张图像在形状 (32, 3, 832, 832)  # <--- TTA 速度

评估pycocotools mAP... 保存 runs/val/exp2/yolov5x_predictions.json...
...
 平均精度 (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.516  # <--- TTA mAP
 平均精度 (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.701
 平均精度 (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.562
 平均精度 (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.361
 平均精度 (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.564
 平均精度 (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.656
 平均召回率 (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.388
 平均召回率 (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.640
 平均召回率 (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.696  # <--- TTA mAR
 平均召回率 (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.553
 平均召回率 (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.744
 平均召回率 (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.833

使用TTA进行推理

detect.py TTA推理与val.py TTA的操作方式相同：只需在任何现有的detect.py命令后附加--augment：

python detect.py --weights yolov5s.pt --img 832 --source data/images --augment

输出：

YOLOv5 🚀 v5.0-267-g6a3ee7c torch 1.9.0+cu102 CUDA:0 (Tesla P100-PCIE-16GB, 16280.875MB)

Downloading https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt to yolov5s.pt...
100% 14.1M/14.1M [00:00<00:00, 81.9MB/s]

Fusing layers...
Model Summary: 224 layers, 7266973 parameters, 0 gradients
image 1/2 /content/yolov5/data/images/bus.jpg: 832x640 4 persons, 1 bus, 1 fire hydrant, Done. (0.029s)
image 2/2 /content/yolov5/data/images/zidane.jpg: 480x832 3 persons, 3 ties, Done. (0.024s)
Results saved to runs/detect/exp
Done. (0.156s)

YOLOv5测试时间增强

PyTorch Hub TTA

TTA已自动集成到所有YOLOv5 PyTorch Hub模型中，可以通过在推理时传递augment=True来访问。

import torch

# 模型
model = torch.hub.load("ultralytics/yolov5", "yolov5s")  # 或 yolov5m, yolov5x, custom

# 图像
img = "https://ultralytics.com/images/zidane.jpg"  # 或文件, PIL, OpenCV, numpy, 多个

# 推理
results = model(img, augment=True)  # <--- TTA推理

# 结果
results.print()  # 或 .show(), .save(), .crop(), .pandas(), 等.