Skip to content

Latest commit

 

History

History
278 lines (188 loc) · 5.56 KB

README.md

File metadata and controls

278 lines (188 loc) · 5.56 KB
Raw

YOLO Series TensorRT Python/C++

Support

YOLOv10、YOLOv9、YOLOv8、YOLOv7、YOLOv6、 YOLOX、 YOLOV5、YOLOv3

  • YOLOv10
  • YOLOv9
  • YOLOv8
  • YOLOv7
  • YOLOv6
  • YOLOX
  • YOLOv5
  • YOLOv3

Update

  • 2024.6.16 Support YOLOv9, YOLOv10, changing the TensorRT version to 10.0
  • 2023.8.15 Support cuda-python
  • 2023.5.12 Update
  • 2023.1.7 support YOLOv8
  • 2022.11.29 fix some bug thanks @JiaPai12138
  • 2022.8.13 rename reop、 public new version、 C++ for end2end
  • 2022.8.11 nms plugin support ==> Now you can set --end2end flag while use export.py get a engine file
  • 2022.7.8 support YOLOv7
  • 2022.7.3 support TRT int8 post-training quantization

Prepare TRT Env

Install via Python

pip install tensorrt
pip install cuda-python

Install via C++

By Docker

YOLOv10

Generate TRT File

python export.py  -o yolov10n.onnx -e yolov10.trt --end2end --v10 -p fp32

Inference

python trt.py -e yolov10.trt  -i src/1.jpg -o yolov10-1.jpg --end2end

YOLOv9

Generate TRT File

python export.py  -o yolov9-c.onnx -e yolov9.trt --end2end --v8 -p fp32

Inference

python trt.py -e yolov9.trt  -i src/1.jpg -o yolov9-1.jpg --end2end

Python Demo

Expand
  1. YOLOv5
  2. YOLOx
  3. YOLOv6
  4. YOLOv7
  5. YOLOv8

YOLOv8

Install && Download Weights

pip install ultralytics

Export ONNX

from ultralytics import YOLO

model = YOLO("yolov8s.pt")
model.fuse()  
model.info(verbose=False)  # Print model information
model.export(format='onnx')  # TODO:

Generate TRT File

python export.py -o yolov8n.onnx -e yolov8n.trt --end2end --v8 --fp32

Inference

python trt.py -e yolov8n.trt  -i src/1.jpg -o yolov8n-1.jpg --end2end

YOLOv5

!git clone https://github.com/ultralytics/yolov5.git
!wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5n.pt
!python yolov5/export.py --weights yolov5n.pt --include onnx --simplify --inplace

include NMS Plugin

!python export.py -o yolov5n.onnx -e yolov5n.trt --end2end
!python trt.py -e yolov5n.trt  -i src/1.jpg -o yolov5n-1.jpg --end2end

exclude NMS Plugin

!python export.py -o yolov5n.onnx -e yolov5n.trt 
!python trt.py -e yolov5n.trt  -i src/1.jpg -o yolov5n-1.jpg

YOLOX

!git clone https://github.com/Megvii-BaseDetection/YOLOX.git
!wget https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.pth
!cd YOLOX && pip3 install -v -e . --user
!cd YOLOX && python tools/export_onnx.py --output-name ../yolox_s.onnx -n yolox-s -c ../yolox_s.pth --decode_in_inference

include NMS Plugin

!python export.py -o yolox_s.onnx -e yolox_s.trt --end2end
!python trt.py -e yolox_s.trt  -i src/1.jpg -o yolox-1.jpg --end2end

exclude NMS Plugin

!python export.py -o yolox_s.onnx -e yolox_s.trt 
!python trt.py -e yolox_s.trt  -i src/1.jpg -o yolox-1.jpg

YOLOv6

!wget https://github.com/meituan/YOLOv6/releases/download/0.1.0/yolov6s.onnx

include NMS Plugin

!python export.py -o yolov6s.onnx -e yolov6s.trt --end2end
!python trt.py -e yolov6s.trt  -i src/1.jpg -o yolov6s-1.jpg --end2end

exclude NMS Plugin

!python export.py -o yolov6s.onnx -e yolov6s.trt 
!python trt.py -e yolov6s.trt  -i src/1.jpg -o yolov6s-1.jpg

YOLOv7

!git clone https://github.com/WongKinYiu/yolov7.git
!wget https://github.com/WongKinYiu/yolov7/releases/download/v0.1/yolov7-tiny.pt
!pip install -r yolov7/requirements.txt
!python yolov7/export.py --weights yolov7-tiny.pt --grid  --simplify

include NMS Plugin

!python export.py -o yolov7-tiny.onnx -e yolov7-tiny.trt --end2end
!python trt.py -e yolov7-tiny.trt  -i src/1.jpg -o yolov7-tiny-1.jpg --end2end

exclude NMS Plugin

!python export.py -o yolov7-tiny.onnx -e yolov7-tiny-norm.trt
!python trt.py -e yolov7-tiny-norm.trt  -i src/1.jpg -o yolov7-tiny-norm-1.jpg

C++ Demo

support NMS plugin show in C++ Demo

Citing

If you use this repo in your publication, please cite it by using the following BibTeX entry.

@Misc{yolotrt2022,
  author =       {Jian Lin},
  title =        {YOLOTRT: tensorrt for yolo series},
  howpublished = {\url{[https://github.com/Linaom1214/TensorRT-For-YOLO-Series]}},
  year =         {2022}
}

Sponsor

Buy me a cup of coffee