最简洁的端到端实现方案,无需第三方库
支持 YOLOv5,6,7, YOLOX
Get This repo
pip install --upgrade setuptools pip --user
pip install nvidia-pyindex
pip install --upgrade nvidia-tensorrt
pip install pycuda
git clone https://github.com/Linaom1214/TensorRT-For-YOLO-Series.git
cd TensorRT-For-YOLO-Series
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