Awesome

This repo only focuses on NMS speed improvement based on https://github.com/ultralytics/yolov5.

See non_max_suppression function of utils/general.py for our Cluster-NMS implementation.

Batch mode Cluster-NMS

Torchvision NMS has the fastest speed but fails to run in batch mode.

Batch mode Cluster-NMS is made for this.

Our goal is that when using TTA for getting better performance, NMS no longer becomes a potential time-consuming growth factor.

Some Pretrained Weights

For more details, please refer to https://github.com/ultralytics/yolov5.

Requirements

Python 3.8 or later with all requirements.txt dependencies installed, including torch>=1.6. To install run:

$ pip install -r requirements.txt

Evaluation for Batch Mode Weighted Cluster-NMS

Hardware

• 1 RTX 2080 Ti

Evaluation command: python test.py --weights yolov5s.pt --data coco.yaml --img 640 --augment --merge --batch-size 32

YOLOv5s.pt

YOLOv5m.pt

YOLOv5x.pt python test.py --weights yolov5s.pt --data coco.yaml --img 832 --augment --merge --batch-size 32

Details:

• AP reports on coco 2017val.
• TTA denotes Test-Time Augmentation.
• max-box denotes maximum number of boxes processed in Batch Mode Cluster-NMS.
• weighted threshold denotes the threshold used in weighted coordinates.
• time reports model inference / NMS.
• To avoid randomness, NMS runs three times here. See test.py.

# Run NMS
t = time_synchronized()
output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, max_box=max_box, merge=merge)
output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, max_box=max_box, merge=merge)
output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, max_box=max_box, merge=merge)
t1 += time_synchronized() - t

Conclusion

• Batch mode Weighted Cluster-NMS will have comparable speed with Torchvision merge NMS when batchsize>=16 and without TTA.
• When using TTA, the time of torchvision NMS will increase significantly, because the model predicts much more boxes. Especially when using multi-scale testing or more TTA means.
• Observed from experience, when using TTA, max-box = 1500 will be good. And when TTA is turned off, max-box = 1000.

Related issues

• Test-Time Augmentation (TTA)
• Model Ensembling
• INCREASING NMS SPEED
• TESTING/INFERENCE AUGMENTATION

Environments

YOLOv5 may be run in any of the following up-to-date verified environments (with all dependencies including CUDA/CUDNN, Python and PyTorch preinstalled):

• Google Colab Notebook with free GPU: <a href="https://colab.research.google.com/github/ultralytics/yolov5/blob/master/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
• Kaggle Notebook with free GPU: https://www.kaggle.com/ultralytics/yolov5
• Google Cloud Deep Learning VM. See GCP Quickstart Guide
• Docker Image https://hub.docker.com/r/ultralytics/yolov5. See Docker Quickstart Guide

Citation

This is the code for our paper:

• Distance-IoU Loss: Faster and Better Learning for Bounding Box Regression
• Enhancing Geometric Factors into Model Learning and Inference for Object Detection and Instance Segmentation

@Inproceedings{zheng2020diou,
  author    = {Zheng, Zhaohui and Wang, Ping and Liu, Wei and Li, Jinze and Ye, Rongguang and Ren, Dongwei},
  title     = {Distance-IoU Loss: Faster and Better Learning for Bounding Box Regression},
  booktitle = {The AAAI Conference on Artificial Intelligence (AAAI)},
  year      = {2020},
}

@Article{zheng2021ciou,
  author    = {Zheng, Zhaohui and Wang, Ping and Ren, Dongwei and Liu, Wei and Ye, Rongguang and Hu, Qinghua and Zuo, Wangmeng},
  title     = {Enhancing Geometric Factors in Model Learning and Inference for Object Detection and Instance Segmentation},
  booktitle = {IEEE Transactions on Cybernetics},
  year      = {2021},
}