# tensorRT_Pro-YOLOv8
**Repository Path**: zyb314/tensorRT_Pro-YOLOv8
## Basic Information
- **Project Name**: tensorRT_Pro-YOLOv8
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: MIT
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No
## Statistics
- **Stars**: 0
- **Forks**: 3
- **Created**: 2024-10-12
- **Last Updated**: 2024-10-12
## Categories & Tags
**Categories**: Uncategorized
**Tags**: None
## README
## 简介
该仓库基于 [shouxieai/tensorRT_Pro](https://github.com/shouxieai/tensorRT_Pro),并进行了调整以支持 YOLOv8 的各项任务。
* 目前已支持 YOLOv8、YOLOv8-Cls、YOLOv8-Seg、YOLOv8-Pose 高性能推理!!!🚀🚀🚀
* 基于 tensorRT8.x,C++ 高级接口,C++ 部署,服务器/嵌入式使用

## CSDN文章同步讲解
- 🔥 [YOLOv8推理详解及部署实现](https://blog.csdn.net/qq_40672115/article/details/134276907)
- 🔥 [YOLOv8-Cls推理详解及部署实现](https://blog.csdn.net/qq_40672115/article/details/134277392)
- 🔥 [YOLOv8-Seg推理详解及部署实现](https://blog.csdn.net/qq_40672115/article/details/134277752)
- 🔥 [YOLOv8-Pose推理详解及部署实现](https://blog.csdn.net/qq_40672115/article/details/134278117)
- 🔥 [RT-DETR推理详解及部署实现](https://blog.csdn.net/qq_40672115/article/details/134356250)
## 环境配置
该项目依赖于 cuda、cudnn、tensorRT、opencv、protobuf 库,请在 CMakeLists.txt 或 Makefile 中手动指定路径配置
* 服务器
* CUDA >= 10.2
* cuDNN >= 8.x
* TensorRT >= 8.x
* protobuf == 3.11.4
* 软件安装请参考:[Ubuntu20.04软件安装大全](https://blog.csdn.net/qq_40672115/article/details/130255299)
* 嵌入式
* jetpack >= 4.6
* protobuf == 3.11.4
克隆该项目
```shell
git clone https://github.com/Melody-Zhou/tensorRT_Pro-YOLOv8.git
```
CMakeLists.txt 编译
1. 修改库文件路径
```cmake
# CMakeLists.txt 13 行, 修改 opencv 路径
set(OpenCV_DIR "/usr/local/include/opencv4/")
# CMakeLists.txt 15 行, 修改 cuda 路径
set(CUDA_TOOLKIT_ROOT_DIR "/usr/local/cuda-11.6")
# CMakeLists.txt 16 行, 修改 cudnn 路径
set(CUDNN_DIR "/usr/local/cudnn8.4.0.27-cuda11.6")
# CMakeLists.txt 17 行, 修改 tensorRT 路径
set(TENSORRT_DIR "/opt/TensorRT-8.4.1.5")
# CMakeLists.txt 20 行, 修改 protobuf 路径
set(PROTOBUF_DIR "/home/jarvis/protobuf")
```
2. 编译
```shell
mkdir build
cd build
cmake ..
make -j64
```
Makefile 编译
1. 修改库文件路径
```makefile
# Makefile 4 行,修改 protobuf 路径
lean_protobuf := /home/jarvis/protobuf
# Makefile 5 行,修改 tensorRT 路径
lean_tensor_rt := /opt/TensorRT-8.4.1.5
# Makefile 6 行,修改 cudnn 路径
lean_cudnn := /usr/local/cudnn8.4.0.27-cuda11.6
# Makefile 7 行,修改 opencv 路径
lean_opencv := /usr/local
# Makefile 8 行,修改 cuda 路径
lean_cuda := /usr/local/cuda-11.6
```
2. 编译
```shell
make -j64
```
## 各项任务支持
YOLOv3支持
1. 下载 YOLOv3
```shell
git clone https://github.com/ultralytics/yolov3.git
```
2. 修改代码, 保证动态 batch
```python
# ========== export.py ==========
# yolov3/export.py第160行
# output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output0']
# if dynamic:
# dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640)
# if isinstance(model, SegmentationModel):
# dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
# dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
# elif isinstance(model, DetectionModel):
# dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
# 修改为:
output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output']
if dynamic:
dynamic = {'images': {0: 'batch'}} # shape(1,3,640,640)
if isinstance(model, SegmentationModel):
dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
elif isinstance(model, DetectionModel):
dynamic['output'] = {0: 'batch'} # shape(1,25200,85)
```
3. 导出 onnx 模型
```shell
cd yolov3
python export.py --weights=yolov3.pt --dynamic --simplify --include=onnx --opset=11
```
4. 复制模型并执行
```shell
cp yolov3/yolov3.onnx tensorRT_Pro-YOLOv8/workspace
cd tensorRT_Pro-YOLOv8
# 修改代码在 src/application/app_yolo.cpp: app_yolo 函数中, 使用 V3 的方式即可运行
# test(Yolo::Type::V3, TRT::Mode::FP32, "yolov3");
make yolo -j64
```
YOLOX支持
1. 下载 YOLOX
```shell
git clone https://github.com/Megvii-BaseDetection/YOLOX.git
```
2. 导出 onnx 模型
```shell
cd YOLOX
export PYTHONPATH=$PYTHONPATH:.
python tools/export_onnx.py -c yolox_s.pth -f exps/default/yolox_s.py --output-name=yolox_s.onnx --dynamic --decode_in_inference
```
3. 复制模型并执行
```shell
cp YOLOX/yolox_s.onnx tensorRT_Pro-YOLOv8/workspace
cd tensorRT_Pro-YOLOv8
# 修改代码在 src/application/app_yolo.cpp: app_yolo 函数中, 使用 X 的方式即可运行
# test(Yolo::Type::X, TRT::Mode::FP32, "yolox_s");
make yolo -j64
```
YOLOv5支持
1. 下载 YOLOv5
```shell
git clone https://github.com/ultralytics/yolov5.git
```
2. 修改代码, 保证动态 batch
```python
# ========== export.py ==========
# yolov5/export.py第160行
# output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output0']
# if dynamic:
# dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640)
# if isinstance(model, SegmentationModel):
# dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
# dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
# elif isinstance(model, DetectionModel):
# dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
# 修改为:
output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output']
if dynamic:
dynamic = {'images': {0: 'batch'}} # shape(1,3,640,640)
if isinstance(model, SegmentationModel):
dynamic['output0'] = {0: 'batch', 1: 'anchors'} # shape(1,25200,85)
dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
elif isinstance(model, DetectionModel):
dynamic['output'] = {0: 'batch'} # shape(1,25200,85)
```
3. 导出 onnx 模型
```shell
cd yolov5
python export.py --weights=yolov5s.pt --dynamic --simplify --include=onnx --opset=11
```
4. 复制模型并执行
```shell
cp yolov5/yolov5s.onnx tensorRT_Pro-YOLOv8/workspace
cd tensorRT_Pro-YOLOv8
# 修改代码在 src/application/app_yolo.cpp: app_yolo 函数中, 使用 V5 的方式即可运行
# test(Yolo::Type::V5, TRT::Mode::FP32, "yolov5s");
make yolo -j64
```
YOLOv6支持
1. 下载 YOLOv6
```shell
git clone https://github.com/meituan/YOLOv6.git
```
2. 修改代码, 保证动态 batch
```python
# ========== export_onnx.py ==========
# YOLOv6/deploy/ONNX/export_onnx.py第84行
# output_axes = {
# 'outputs': {0: 'batch'},
# }
# 修改为:
output_axes = {
'output': {0: 'batch'},
}
# YOLOv6/deploy/ONNX/export_onnx.py第106行
# torch.onnx.export(model, img, f, verbose=False, opset_version=13,
# training=torch.onnx.TrainingMode.EVAL,
# do_constant_folding=True,
# input_names=['images'],
# output_names=['num_dets', 'det_boxes', 'det_scores', 'det_classes']
# if args.end2end else ['outputs'],
# dynamic_axes=dynamic_axes)
# 修改为:
torch.onnx.export(model, img, f, verbose=False, opset_version=13,
training=torch.onnx.TrainingMode.EVAL,
do_constant_folding=True,
input_names=['images'],
output_names=['num_dets', 'det_boxes', 'det_scores', 'det_classes']
if args.end2end else ['output'],
dynamic_axes=dynamic_axes)
```
3. 导出 onnx 模型
```shell
cd YOLOv6
python deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 640 --dynamic-batch --simplify
```
4. 复制模型并执行
```shell
cp YOLOv6/yolov6s.onnx tensorRT_Pro-YOLOv8/workspace
cd tensorRT_Pro-YOLOv8
# 修改代码在 src/application/app_yolo.cpp: app_yolo 函数中, 使用 V6 的方式即可运行
# test(Yolo::Type::V6, TRT::Mode::FP32, "yolov6s");
make yolo -j64
```
YOLOv7支持
1. 下载 YOLOv7
```shell
git clone https://github.com/WongKinYiu/yolov7.git
```
2. 导出 onnx 模型
```shell
python export.py --dynamic-batch --grid --simplify --weights=yolov7.pt
```
3. 复制模型并执行
```shell
cp yolov7/yolov7.onnx tensorRT_Pro-YOLOv8/workspace
cd tensorRT_Pro-YOLOv8
# 修改代码在 src/application/app_yolo.cpp: app_yolo 函数中, 使用 V7 的方式即可运行
# test(Yolo::Type::V7, TRT::Mode::FP32, "yolov7");
make yolo -j64
```
YOLOv8支持
1. 下载 YOLOv8
```shell
git clone https://github.com/ultralytics/ultralytics.git
```
2. 修改代码, 保证动态 batch
```python
# ========== head.py ==========
# ultralytics/nn/modules/head.py第72行,forward函数
# return y if self.export else (y, x)
# 修改为:
return y.permute(0, 2, 1) if self.export else (y, x)
# ========== exporter.py ==========
# ultralytics/engine/exporter.py第323行
# output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output0']
# dynamic = self.args.dynamic
# if dynamic:
# dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640)
# if isinstance(self.model, SegmentationModel):
# dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400)
# dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
# elif isinstance(self.model, DetectionModel):
# dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 84, 8400)
# 修改为:
output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output']
dynamic = self.args.dynamic
if dynamic:
dynamic = {'images': {0: 'batch'}} # shape(1,3,640,640)
if isinstance(self.model, SegmentationModel):
dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400)
dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
elif isinstance(self.model, DetectionModel):
dynamic['output'] = {0: 'batch'} # shape(1, 84, 8400)
```
3. 导出 onnx 模型, 在 ultralytics-main 新建导出文件 `export.py` 内容如下:
```python
# ========== export.py ==========
from ultralytics import YOLO
model = YOLO("yolov8s.pt")
success = model.export(format="onnx", dynamic=True, simplify=True)
```
```shell
cd ultralytics-main
python export.py
```
4. 复制模型并执行
```shell
cp ultralytics/yolov8s.onnx tensorRT_Pro-YOLOv8/workspace
cd tensorRT_Pro-YOLOv8
make yolo -j64
```
YOLOv8-Cls支持
1. 下载 YOLOv8
```shell
git clone https://github.com/ultralytics/ultralytics.git
```
2. 修改代码, 保证动态 batch
```python
# ========== exporter.py ==========
# ultralytics/engine/exporter.py第323行
# output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output0']
# dynamic = self.args.dynamic
# if dynamic:
# dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640)
# if isinstance(self.model, SegmentationModel):
# dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400)
# dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
# elif isinstance(self.model, DetectionModel):
# dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 84, 8400)
# 修改为:
output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output']
dynamic = self.args.dynamic
if dynamic:
dynamic = {'images': {0: 'batch'}} # shape(1,3,640,640)
dynamic['output'] = {0: 'batch'}
if isinstance(self.model, SegmentationModel):
dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400)
dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
elif isinstance(self.model, DetectionModel):
dynamic['output'] = {0: 'batch'} # shape(1, 84, 8400)
```
3. 导出 onnx 模型, 在 ultralytics-main 新建导出文件 `export.py` 内容如下:
```python
# ========== export.py ==========
from ultralytics import YOLO
model = YOLO("yolov8s-cls.pt")
success = model.export(format="onnx", dynamic=True, simplify=True)
```
```shell
cd ultralytics-main
python export.py
```
4. 复制模型并执行
```shell
cp ultralytics/yolov8s-cls.onnx tensorRT_Pro-YOLOv8/workspace
cd tensorRT_Pro-YOLOv8
make yolo_cls -j64
```
YOLOv8-Seg支持
1. 下载 YOLOv8
```shell
git clone https://github.com/ultralytics/ultralytics.git
```
2. 修改代码, 保证动态 batch
```python
# ========== head.py ==========
# ultralytics/nn/modules/head.py第106行,forward函数
# return (torch.cat([x, mc], 1), p) if self.export else (torch.cat([x[0], mc], 1), (x[1], mc, p))
# 修改为:
return (torch.cat([x, mc], 1).permute(0, 2, 1), p) if self.export else (torch.cat([x[0], mc], 1), (x[1], mc, p))
# ========== exporter.py ==========
# ultralytics/engine/exporter.py第323行
# output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output0']
# dynamic = self.args.dynamic
# if dynamic:
# dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640)
# if isinstance(self.model, SegmentationModel):
# dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400)
# dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
# elif isinstance(self.model, DetectionModel):
# dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 84, 8400)
# 修改为:
output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output0']
dynamic = self.args.dynamic
if dynamic:
dynamic = {'images': {0: 'batch'}} # shape(1,3,640,640)
if isinstance(self.model, SegmentationModel):
dynamic['output0'] = {0: 'batch'} # shape(1, 116, 8400)
dynamic['output1'] = {0: 'batch'} # shape(1,32,160,160)
elif isinstance(self.model, DetectionModel):
dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 84, 8400)
```
3. 导出 onnx 模型, 在 ultralytics-main 新建导出文件 `export.py` 内容如下:
```python
# ========== export.py ==========
from ultralytics import YOLO
model = YOLO("yolov8s-seg.pt")
success = model.export(format="onnx", dynamic=True, simplify=True)
```
```shell
cd ultralytics-main
python export.py
```
4. 复制模型并执行
```shell
cp ultralytics/yolov8s-seg.onnx tensorRT_Pro-YOLOv8/workspace
cd tensorRT_Pro-YOLOv8
make yolo_seg -j64
```
YOLOv8-Pose支持
1. 下载 YOLOv8
```shell
git clone https://github.com/ultralytics/ultralytics.git
```
2. 修改代码, 保证动态 batch
```python
# ========== head.py ==========
# ultralytics/nn/modules/head.py第130行,forward函数
# return torch.cat([x, pred_kpt], 1) if self.export else (torch.cat([x[0], pred_kpt], 1), (x[1], kpt))
# 修改为:
return torch.cat([x, pred_kpt], 1).permute(0, 2, 1) if self.export else (torch.cat([x[0], pred_kpt], 1), (x[1], kpt))
# ========== exporter.py ==========
# ultralytics/engine/exporter.py第323行
# output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output0']
# dynamic = self.args.dynamic
# if dynamic:
# dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640)
# if isinstance(self.model, SegmentationModel):
# dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400)
# dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
# elif isinstance(self.model, DetectionModel):
# dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 84, 8400)
# 修改为:
output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output']
dynamic = self.args.dynamic
if dynamic:
dynamic = {'images': {0: 'batch'}} # shape(1,3,640,640)
dynamic['output'] = {0: 'batch'}
if isinstance(self.model, SegmentationModel):
dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400)
dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
elif isinstance(self.model, DetectionModel):
dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 84, 8400)
```
3. 导出 onnx 模型, 在 ultralytics-main 新建导出文件 `export.py` 内容如下:
```python
# ========== export.py ==========
from ultralytics import YOLO
model = YOLO("yolov8s-pose.pt")
success = model.export(format="onnx", dynamic=True, simplify=True)
```
```shell
cd ultralytics-main
python export.py
```
4. 复制模型并执行
```shell
cp ultralytics/yolov8s-pose.onnx tensorRT_Pro-YOLOv8/workspace
cd tensorRT_Pro-YOLOv8
make yolo_pose -j64
```
RT-DETR支持
1. 前置条件
- **tensorRT >= 8.6**
2. 下载 YOLOv8
```shell
git clone https://github.com/ultralytics/ultralytics.git
```
3. 修改代码, 保证动态 batch
```python
# ========== exporter.py ==========
# ultralytics/engine/exporter.py第323行
# output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output0']
# dynamic = self.args.dynamic
# if dynamic:
# dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}} # shape(1,3,640,640)
# if isinstance(self.model, SegmentationModel):
# dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400)
# dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
# elif isinstance(self.model, DetectionModel):
# dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 84, 8400)
# 修改为:
output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output']
dynamic = self.args.dynamic
if dynamic:
dynamic = {'images': {0: 'batch'}} # shape(1,3,640,640)
if isinstance(self.model, SegmentationModel):
dynamic['output0'] = {0: 'batch', 2: 'anchors'} # shape(1, 116, 8400)
dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'} # shape(1,32,160,160)
elif isinstance(self.model, DetectionModel):
dynamic['output'] = {0: 'batch'} # shape(1, 84, 8400)
```
4. 导出 onnx 模型,在 ultralytics-main 新建导出文件 `export.py` 内容如下(可能会由于 torch 版本问题导出失败, 具体可参考 [#6144](https://github.com/ultralytics/ultralytics/issues/6144))
```python
from ultralytics import RTDETR
model = RTDETR("rtdetr-l.pt")
success = model.export(format="onnx", dynamic=True, simplify=True)
```
```shell
cd ultralytics-main
python export.py
```
5. engine 生成
- **方案一**:替换 tensorRT_Pro-YOLOv8 中的 onnxparser 解析器,具体可参考文章:[RT-DETR推理详解及部署实现](https://blog.csdn.net/qq_40672115/article/details/134356250)
- **方案二**:利用 **trtexec** 工具生成 engine
```shell
cp ultralytics/yolov8s.onnx tensorRT_Pro-YOLOv8/workspace
cd tensorRT_Pro-YOLOv8/workspace
bash build.sh
```
6. 执行
```shell
make rtdetr -j64
```
## 接口介绍
编译接口
```cpp
TRT::compile(
mode, // FP32、FP16、INT8
test_batch_size, // max batch size
onnx_file, // source
model_file, // save to
{}, // redefine the input shape
int8process, // the recall function for calibration
"inference", // the dir where the image data is used for calibration
"" // the dir where the data generated from calibration is saved(a.k.a where to load the calibration data.)
);
```
* tensorRT_Pro 原编译接口, 支持 FP32、FP16、INT8 编译
* 模型的编译工作也可以通过 `trtexec` 工具完成
推理接口
```cpp
// 创建推理引擎在 0 号显卡上
auto engine = YoloPose::create_infer(
engine_file, // engine file
deviceid, // gpu id
0.25f, // confidence threshold
0.45f, // nms threshold
YoloPose::NMSMethod::FastGPU, // NMS method, fast GPU / CPU
1024, // max objects
false // preprocess use multi stream
);
// 加载图像
auto image = cv::imread("inference/car.jpg");
// 推理并获取结果
auto boxes = engine->commit(image).get() // 得到的是 vector
```
## 参考
- [https://github.com/shouxieai/tensorRT_Pro](https://github.com/shouxieai/tensorRT_Pro)
- [https://github.com/shouxieai/infer](https://github.com/shouxieai/infer)