https://www.bilibili.com/video/BV1ka411p7u4/
YOLOv6是美团发布的Anchor-free目标检测模型,本仓库实现YOLOv6端到端的Tensorrt部署。
- Tensorrt 8.4.1.5
- Cuda 10.2 Cudnn 8.4.1
- onnx 1.8.0
- onnx-simplifier 0.3.10
- onnx-graphsurgeon 0.3.19
- Torch 1.10.0
Model | Size | mAPval 0.5:0.95 | GTX1650 FP16(ms) | GTX1650 FP32(ms) |
---|---|---|---|---|
YOLOv6n | 640 | 35.0 | 6.6 | 8.7 |
YOLOv6-tiny | 640 | 41.3 | 10.6 | 16.7 |
YOLOv6s | 640 | 43.1 | 11.6 | 18.4 |
说明:此处FP16,fp32预测时间包含preprocess+inference+nms,测速方法为warmup10次,预测100次取平均值,并未使用trtexec测速,与官方测速不同;mAPval为原始模型精度,转换后精度未测试。
下载YOLOv6仓库,并导出含有EfficientNMS的端到端ONNX模型。
git clone https://github.com/meituan/YOLOv6
cd YOLOv6
pip install -r requirements.txt
导出FP32 ONNX模型
python ./deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 640 --batch 1 --end2end
导出FP16 ONNX模型
python ./deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 640 --batch 1 --end2end --half
使用Tensorrt官方工具转换engine模型,FP32预测删除--fp16
参数即可。CUDA10.2+Tensorrt8.4.1存在BUG,需要指定--tacticSources=-cublasLt,+cublas
,否则预测时报错,CUDA11+请忽略该参数。
trtexec --onnx=./yolov6s_half.onnx --saveEngine=./yolov6s_half.engine --fp16 --workspace=30 --tacticSources=-cublasLt,+cublas
等待生成序列化模型后,修改本仓库infer.py模型路径和图片路径。
trt_engine = TRT_engine("./trt_model/yolov6s_half.engine")
img1 = cv2.imread("./pictures/zidane.jpg")
python infer.py
https://github.com/meituan/YOLOv6
https://github.com/ultralytics/yolov5