English | 简体中文
)
TensorRT-YOLO 是一个支持 YOLOv5、YOLOv8、YOLOv9、PP-YOLOE 和 PP-YOLOE+ 的推理加速项目，使用 NVIDIA TensorRT 进行优化。项目不仅集成了 EfficientNMS TensorRT 插件以增强后处理效果，还使用了 CUDA 核函数来加速前处理过程。TensorRT-YOLO 提供了 C++ 和 Python 推理的支持，旨在提供快速而优化的目标检测解决方案。
- 支持 YOLOv5、YOLOv8、YOLOv9、PP-YOLOE 和 PP-YOLOE+
- 支持 ONNX 静态、动态导出以及 TensorRT 推理
- 集成 EfficientNMS TensorRT 插件加速后处理
- 利用 CUDA 核函数加速前处理 (V1.0)
- 支持 C++ 和 Python 推理（C++ 实现中）
- 推荐 CUDA 版本 >= 11.7
- 推荐 TensorRT 版本 >= 8.6
安装
克隆 repo，并要求在 Python>=3.8.0 环境中安装 requirements.txt，且要求 PyTorch>=1.8（导出 YOLOv5、YOLOv8 与 YOLOv9）、PaddlePaddle>=2.5（导出 PP-YOLOE 与 PP-YOLOE+）。
git clone https://github.com/laugh12321/TensorRT-YOLO # clone
cd TensorRT-YOLO
pip install -r requirements.txt # install
pip install ultralytics # Optional, export YOLOv5, YOLOv8 and YOLOv9
pip install paddle2onnx # Optional, export PP-YOLOE and PP-YOLOE+模型导出
使用下面的命令将导出 ONNX 模型并添加 EfficientNMS 插件进行后处理。
注意： 导出 PP-YOLOE 与 PP-YOLOE+ 的 ONNX 模型，只会对 batch 维度进行修改，height 与 width 维度无法被更改，需要在PaddleDetection中设置，默认为 640。
YOLOv5, v8, v9
# Static
python python/export/{yolo version}/export.py -w your_model_path.pt -o output -b 8 --img 640 -s
# Dynamic
python python/export/{yolo version}/export.py -w your_model_path.pt -o output -s --dynamicPP-YOLOE 与 PP-YOLOE+
# Static
python python/export/ppyoloe/export.py --model_dir modeldir --model_filename model.pdmodel --params_filename model.pdiparams -o output -b 8 -s
# Dynamic
python python/export/ppyoloe/export.py --model_dir modeldir --model_filename model.pdmodel --params_filename model.pdiparams -o output -s --dynamic生成的 ONNX 模型使用 trtexec 工具导出 TensorRT 模型。
# Static
trtexec --onnx=model.onnx --saveEngine=model.engine --fp16
# Dynamic
trtexec --onnx=model.onnx --saveEngine=model.engine --minShapes=images:1x3x640x640 --optShapes=images:4x3x640x640 --maxShapes=images:8x3x640x640 --fp16使用 detect.py 推理
detect.py 目前支持对单张图片进行推理或批量推理整个目录，可通过 --inputs 参数指定推理数据。推理结果可使用 --output 参数指定保存路径，默认为 None，表示不保存。有关详细指令描述，请运行python detect.py -h查看。
python detect.py -e model.engine -o output -i img.jpg # image
path/ # directoryTensorRT-YOLO采用 GPL-3.0许可证，这个OSI 批准的开源许可证非常适合学生和爱好者，可以推动开放的协作和知识分享。请查看LICENSE 文件以了解更多细节。
感谢您选择使用 TensorRT-YOLO，我们鼓励开放的协作和知识分享，同时也希望您遵守开源许可的相关规定。
对于 TensorRT-YOLO 的错误报告和功能请求，请访问 GitHub Issues！