YoloV8 TensorRT CPP
使用TensorRT的YoloV8 C++实现
支持目标检测、语义分割和人体姿态估计。
入门指南
本项目演示了如何使用TensorRT C++ API来运行YoloV8的GPU推理。 它使用了我的另一个项目tensorrt-cpp-api来在后台运行推理,所以请确保你熟悉那个项目。
先决条件
- 在Ubuntu 20.04和22.04上测试并工作(目前不支持Windows)
- 安装CUDA,说明在这里。
- 推荐 >= 12.0
- 安装cudnn,说明在这里。
- 推荐 >= 8
sudo apt install build-essential
sudo apt install python3-pip
pip3 install cmake
- 安装支持cuda的OpenCV。要从源代码编译OpenCV,请运行这里提供的
build_opencv.sh
脚本。- 推荐 >= 4.8
- 从这里下载TensorRT 10。
- 要求 >= 10.0
- 解压后,导航到
CMakeLists.txt
文件,将TODO
替换为你的TensorRT安装路径。
安装
git clone https://github.com/cyrusbehr/YOLOv8-TensorRT-CPP --recursive
- 注意: 确保使用
--recursive
标志,因为这个仓库使用了git子模块。
将模型从PyTorch转换为ONNX
- 导航到官方YoloV8仓库并下载你想要的模型版本(例如YOLOv8x)。
- 该代码还支持开箱即用的语义分割模型(例如YOLOv8x-seg)和姿态估计模型(例如yolov8x-pose.onnx)。
pip3 install ultralytics
- 导航到
scripts/
目录并运行以下命令: python3 pytorch2onnx.py --pt_path <你的pt文件路径>
- 运行此命令后,你应该已成功从PyTorch转换为ONNX。
- 注意:如果使用不同的脚本转换模型,请确保禁用
end2end
。此标志将直接将bbox解码和nms添加到模型中,而我的实现使用纯C++在模型外部执行这些步骤。
构建项目
mkdir build
cd build
cmake ..
make -j
运行可执行文件
- 注意:第一次运行任何脚本时,可能需要相当长的时间(5分钟以上),因为TensorRT必须从onnx模型生成优化的TensorRT引擎文件。这将保存到磁盘并在后续运行时加载。
- 注意:所有可执行文件都可以与Ultralytic的预训练目标检测、分割和姿态估计模型开箱即用。
- 要运行基准测试脚本,请运行:
./benchmark --model /path/to/your/onnx/model.onnx --input /path/to/your/benchmark/image.png
- 要对图像进行推理并将带注释的图像保存到磁盘,请运行:
./detect_object_image --model /path/to/your/onnx/model.onnx --input /path/to/your/image.jpg
- 你可以使用
images/
目录中的图像进行测试
- 你可以使用
- 要使用网络摄像头进行推理并实时显示结果,请运行:
./detect_object_video --model /path/to/your/onnx/model.onnx --input 0
- 要查看完整的参数列表,请不提供任何参数运行任何可执行文件。
INT8推理
启用INT8精度可以进一步加快推理速度,但由于动态范围减小,会牺牲一定的准确性。 对于INT8精度,必须提供代表模型将看到的真实数据的校准数据。 建议使用1000+张校准图像。要使用YoloV8模型进行INT8推理,需要执行以下步骤:
- 下载并解压COCO验证数据集,或获取代表你的推理数据的数据:
wget http://images.cocodataset.org/zips/val2017.zip
- 运行可执行文件时提供额外的命令行参数:
--precision INT8 --calibration-data /path/to/your/calibration/data
- 如果遇到"out of memory in function allocate"错误,则必须减小
Options.calibrationBatchSize
,使整个批次可以适应你的GPU内存。
基准测试
- 在运行基准测试之前,请确保你的GPU未被占用。
- 使用
/images/640_640.jpg
图像运行可执行文件benchmark
。 - 如果你想对每个组件(
preprocess
、inference
、postprocess
)进行基准测试,请将ENABLE_BENCHMARKS
标志设置为ON
重新编译:cmake -DENABLE_BENCHMARKS=ON ..
。- 然后你可以重新运行可执行文件
基准测试在NVIDIA GeForce RTX 3080笔记本GPU,Intel(R) Core(TM) i7-10870H CPU @ 2.20GHz上运行,使用GPU内存中的640x640 BGR图像和FP16精度。
模型 | 总时间 | 预处理时间 | 推理时间 | 后处理时间 |
---|---|---|---|---|
yolov8n | 3.613 ms | 0.081 ms | 1.703 ms | 1.829 ms |
yolov8n-pose | 2.107 ms | 0.091 ms | 1.609 ms | 0.407 ms |
yolov8n-seg | 15.194 ms | 0.109 ms | 2.732 ms | 12.353 ms |
模型 | 精度 | 总时间 | 预处理时间 | 推理时间 | 后处理时间 |
---|---|---|---|---|---|
yolov8x | FP32 | 25.819 ms | 0.103 ms | 23.763 ms | 1.953 ms |
yolov8x | FP16 | 10.147 ms | 0.083 ms | 7.677 ms | 2.387 ms |
yolov8x | INT8 | 7.32 ms | 0.103 ms | 4.698 ms | 2.519 ms |
待办:需要使用CUDA内核改进后处理时间。
如何调试
- 如果你在从onnx模型创建TensorRT引擎文件时遇到问题,请导航到
libs/tensorrt-cpp-api/src/engine.cpp
并将日志级别的严重性级别更改为kVERBOSE
,然后重新构建并重新运行。这应该能给你更多关于构建过程失败的具体位置的信息。
表示你的感谢
如果这个项目对你有帮助,我会很感谢你给它一个星星。这将鼓励我确保它保持最新并快速解决问题。