测试时增强 (TTA)#
📚 本指南介绍了如何在测试和推理过程中使用测试时增强(TTA),以改善 YOLOv5 🚀 的 mAP 和 Recall。
开始之前#
克隆仓库并在包含 PyTorch>=1.8 的 Python>=3.8.0 环境中安装 requirements.txt。模型和数据集会从最新的 YOLOv5 发布版本中自动下载。
git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # install常规测试#
在尝试 TTA 之前,我们需要建立一个基准性能来进行比较。此命令在图像大小为 640 像素的 COCO val2017 上测试 YOLOv5x。yolov5x.pt 是可用最大、最准确的模型。其他选项包括 yolov5s.pt、yolov5m.pt 和 yolov5l.pt,或者你训练自定义数据集得到的检查点 ./weights/best.pt。有关所有可用模型的详细信息,请参阅我们的 YOLOv5 documentation。
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_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) # <--- baseline speed
Evaluating pycocotools mAP... saving runs/val/exp/yolov5x_predictions.json...
...
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.504 # <--- baseline 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 # <--- baseline 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 进行测试#
将 --augment 追加到任何现有的 val.py 命令中以启用 TTA,并将图像大小增加约 30% 以获得更好的结果。请注意,启用 TTA 的推理通常需要正常推理约 2-3 倍的时间,因为图像会被左右翻转并在 3 种不同的分辨率下进行处理,输出在 NMS 之前进行合并。速度下降的部分原因仅仅是因为图像尺寸更大(832 对比 640),部分原因则是由于实际的 TTA 操作,因此在增加 --img 之前,请确保你的 GPU 有足够的内存余量。
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_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...
/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py:718: UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at /pytorch/c10/core/TensorImpl.h:1156.)
return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
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, 2885.61it/s]
val: New cache created: ../datasets/coco/val2017.cache
Class Images Labels P R mAP@.5 mAP@.5:.95: 100% 157/157 [07:29<00:00, 2.86s/it]
all 5000 36335 0.718 0.656 0.695 0.503
Speed: 0.2ms pre-process, 80.6ms inference, 2.7ms NMS per image at shape (32, 3, 832, 832) # <--- TTA speed
Evaluating pycocotools mAP... saving runs/val/exp-2/yolov5x_predictions.json...
...
Average Precision (AP) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.516 # <--- TTA mAP
Average Precision (AP) @[ IoU=0.50 | area= all | maxDets=100 ] = 0.701
Average Precision (AP) @[ IoU=0.75 | area= all | maxDets=100 ] = 0.562
Average Precision (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.361
Average Precision (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.564
Average Precision (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.656
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 1 ] = 0.388
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets= 10 ] = 0.640
Average Recall (AR) @[ IoU=0.50:0.95 | area= all | maxDets=100 ] = 0.696 # <--- TTA mAR
Average Recall (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.553
Average Recall (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.744
Average Recall (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.833使用 TTA 进行推理#
detect.py TTA 推理的操作与 val.py TTA 完全相同:只需将 --augment 追加到任何现有的 detect.py 命令中:
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)
PyTorch Hub TTA#
TTA 已自动集成到所有 YOLOv5 PyTorch Hub 模型中,可以通过在推理时传入 augment=True 来访问。
import torch
# Model
model = torch.hub.load("ultralytics/yolov5", "yolov5s") # or yolov5m, yolov5x, custom
# Images
img = "https://ultralytics.com/images/zidane.jpg" # or file, PIL, OpenCV, numpy, multiple
# Inference
results = model(img, augment=True) # <--- TTA inference
# Results
results.print() # or .show(), .save(), .crop(), .pandas(), etc.自定义#
你可以自定义在 YOLOv5 forward_augment() method 中应用的 TTA 操作。
测试时增强的好处#
测试时增强为对象检测任务带来了几个关键优势:
- 更高的准确性:如上述结果所示,TTA 将 mAP 从 0.504 提高到了 0.516,将 mAR 从 0.681 提高到了 0.696。
- 更好的小目标检测:TTA 特别增强了对小目标的检测能力,小目标 AP 从 0.351 提高到了 0.361。
- 更强的鲁棒性:通过测试每张图像的多种变体,TTA 降低了视角、光照和其他环境因素的影响。
- 简单实现:只需向现有命令添加
--augment标志即可。
其代价是增加了推理时间,这使得 TTA 更适合那些优先考虑准确性而非速度的应用场景。
支持的环境#
Ultralytics 提供了各种开箱即用的环境,每个环境都预先安装了诸如 CUDA、CUDNN、Python 和 PyTorch 等基本依赖项,以帮助你快速启动项目。
- 免费 GPU 笔记本:
- Google Cloud:GCP 快速入门指南
- Amazon:AWS 快速入门指南
- Azure:AzureML 快速入门指南
- Docker:Docker 快速入门指南
项目状态#
此徽章表示所有 YOLOv5 GitHub Actions 持续集成 (CI) 测试均已成功通过。这些 CI 测试严格检查了 YOLOv5 在各个核心方面的功能和性能:训练、验证、推理、导出和基准测试。它们确保在 macOS、Windows 和 Ubuntu 上保持一致且可靠的运行,测试每 24 小时以及每次新提交时进行。