tensorRT_Pro-Vision

简介

tensorRT_Pro-Vision 是一个基于 TensorRT C++ API 的高性能计算机视觉推理框架，支持 20+ 主流视觉模型的一键部署。该仓库从 shouxieai/tensorRT_Pro fork 而来，经过大幅重构和扩展。

目前已支持 检测、分类、实例分割、语义分割、姿态估计、旋转框检测 (OBB)、深度估计、文字识别 (OCR)、车道线检测、多目标跟踪 等多项视觉任务的端到端 GPU 推理。🚀🚀🚀

核心特性：

• 基于 TensorRT C++ API，支持 TensorRT 8.x / 10.x 版本自适应编译（通过 NV_TENSORRT_MAJOR 宏自动适配 API 差异）
• 使用 TensorRT 官方 ONNX 解析器 (libnvonnxparser.so)，无需 protobuf 依赖
• 预处理、后处理、NMS 全部 GPU 加速（CUDA Kernel 实现）
• 支持 FP32 / FP16 / INT8 三种推理精度，内置 Entropy 和 MinMax 校准器
• 支持动态 Batch，最大 batch size 可在编译时配置
• 代码结构清晰，各模型任务独立模块化，易于扩展

✨ What's New

• 2026/6/26 重构
  • TensorRT 10.x 全面支持：通过 NV_TENSORRT_MAJOR 宏实现 TRT 8.x / 10.x 双版本自适应编译，API 差异由 #if NV_TENSORRT_MAJOR >= 10 条件编译自动处理
  • 移除 vendored ONNX 解析器：不再维护自定义的 onnx-tensorrt 解析器（~15,000 行代码），统一使用 TensorRT 官方 libnvonnxparser.so，大幅简化代码库
  • 移除 protobuf 依赖：不再需要单独安装 protobuf 3.11.4，构建更加简单
  • 修复 getMaxBatchSize() 警告：显式 batch 模式下使用 getProfileDimensions / getProfileShape 正确获取最大 batch size
  • 仓库重命名：tensorRT_Pro-YOLOv8 → tensorRT_Pro-Vision，体现多任务视觉推理的定位
  • 全新 Banner：9 宫格任务全景图，直观展示所有支持的视觉任务

📜 完整历史更新记录（2023-2026）请查看 v1.0.0 README，其中包含 CSDN 文章同步讲解链接。

环境配置

该项目依赖于 CUDA、cuDNN、TensorRT、OpenCV 库，请在 Makefile 或 CMakeLists.txt 中手动指定路径配置。

依赖	TensorRT 8.x 推荐	TensorRT 10.x 推荐
CUDA	>= 10.2	>= 12.0
cuDNN	>= 8.x	>= 9.x
TensorRT	>= 8.4	>= 10.0
OpenCV	>= 4.x	>= 4.x

克隆项目：

git clone https://github.com/Melody-Zhou/tensorRT_Pro-Vision.git

1. 修改 Makefile 中的库文件路径：

# ===== TensorRT 8.x =====
lean_tensor_rt := /opt/TensorRT-8.6.1.6
lean_cudnn     := /home/zhouwenguang/lean/cudnn-8.5.0.96
lean_cuda      := /usr/local/cuda-11.4
lean_opencv    := /home/zhouwenguang/lean/opencv-4.6.0

# ===== 或者 TensorRT 10.x =====
# lean_tensor_rt := /home/zhouwenguang/lean/TensorRT-10.16.1
# lean_cudnn     := /home/zhouwenguang/lean/cudnn-9.18.0
# lean_cuda      := /usr/local/cuda-12.8
# lean_opencv    := /home/zhouwenguang/lean/opencv-4.6.0

2. 编译：

make -j$(nproc)

1. 修改 CMakeLists.txt 中的库文件路径

2. 编译：

mkdir build && cd build
cmake .. && make -j$(nproc)

各项任务支持

1. 下载 YOLOv3

git clone https://github.com/ultralytics/yolov3.git

2. 修改代码, 保证动态 batch

# ========== 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 模型

cd yolov3
python export.py --weights=yolov3.pt --dynamic --simplify --include=onnx --opset=11

4. 复制模型并执行

cp yolov3/yolov3.onnx tensorRT_Pro-Vision/workspace
cd tensorRT_Pro-Vision

# 修改代码在 src/application/app_yolo.cpp: app_yolo 函数中, 使用 V3 的方式即可运行
# test(Yolo::Type::V3, TRT::Mode::FP32, "yolov3");

make yolo -j64

1. 下载 YOLOX

git clone https://github.com/Megvii-BaseDetection/YOLOX.git

2. 导出 onnx 模型

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. 复制模型并执行

cp YOLOX/yolox_s.onnx tensorRT_Pro-Vision/workspace
cd tensorRT_Pro-Vision

# 修改代码在 src/application/app_yolo.cpp: app_yolo 函数中, 使用 X 的方式即可运行
# test(Yolo::Type::X, TRT::Mode::FP32, "yolox_s");

make yolo -j64

1. 下载 YOLOv5

git clone https://github.com/ultralytics/yolov5.git

2. 修改代码, 保证动态 batch

# ========== 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 模型

cd yolov5
python export.py --weights=yolov5s.pt --dynamic --simplify --include=onnx --opset=11

4. 复制模型并执行

cp yolov5/yolov5s.onnx tensorRT_Pro-Vision/workspace
cd tensorRT_Pro-Vision

# 修改代码在 src/application/app_yolo.cpp: app_yolo 函数中, 使用 V5 的方式即可运行
# test(Yolo::Type::V5, TRT::Mode::FP32, "yolov5s");

make yolo -j64

1. 下载 YOLOv6

git clone https://github.com/meituan/YOLOv6.git

2. 修改代码, 保证动态 batch，并去除 anchor 维度

# ========== 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)

# 根据不同的 head 去除 anchor 维度
# ========== effidehead_distill_ns.py ==========
# YOLOv6/yolov6/models/heads/effidehead_distill_ns.py第141行
# return torch.cat(
#     [
#         pred_bboxes,
#         torch.ones((b, pred_bboxes.shape[1], 1), device=pred_bboxes.device, dtype=pred_bboxes.dtype),
#         cls_score_list
#     ],
#     axis=-1)
# 修改为：
return torch.cat(
    [
        pred_bboxes,
        cls_score_list
    ],
    axis=-1)

# ========== effidehead_fuseab.py ==========
# YOLOv6/yolov6/models/heads/effidehead_fuseab.py第191行
# return torch.cat(
#     [
#         pred_bboxes,
#         torch.ones((b, pred_bboxes.shape[1], 1), device=pred_bboxes.device, dtype=pred_bboxes.dtype),
#         cls_score_list
#     ],
#     axis=-1)
# 修改为：
return torch.cat(
    [
        pred_bboxes,
        cls_score_list
    ],
    axis=-1)

# ========== effidehead_lite.py ==========
# YOLOv6/yolov6/models/heads/effidehead_lite.py第123行
# return torch.cat(
#     [
#         pred_bboxes,
#         torch.ones((b, pred_bboxes.shape[1], 1), device=pred_bboxes.device, dtype=pred_bboxes.dtype),
#         cls_score_list
#     ],
#     axis=-1)
# 修改为：
return torch.cat(
    [
        pred_bboxes,
        cls_score_list
    ],
    axis=-1)

3. 导出 onnx 模型

cd YOLOv6
python deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 640 --dynamic-batch --simplify

4. 复制模型并执行

cp YOLOv6/yolov6s.onnx tensorRT_Pro-Vision/workspace
cd tensorRT_Pro-Vision

# 修改代码在 src/application/app_yolo.cpp: app_yolo 函数中, 使用 V6 的方式即可运行
# test(Yolo::Type::V6, TRT::Mode::FP32, "yolov6s");

make yolo -j64

1. 下载 YOLOv7

git clone https://github.com/WongKinYiu/yolov7.git 

2. 导出 onnx 模型

python export.py --dynamic-batch --grid --simplify --weights=yolov7.pt

3. 复制模型并执行

cp yolov7/yolov7.onnx tensorRT_Pro-Vision/workspace
cd tensorRT_Pro-Vision

# 修改代码在 src/application/app_yolo.cpp: app_yolo 函数中, 使用 V7 的方式即可运行
# test(Yolo::Type::V7, TRT::Mode::FP32, "yolov7");

make yolo -j64

1. 下载 YOLOv8

git clone https://github.com/ultralytics/ultralytics.git

2. 修改代码, 保证动态 batch

# ========== 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 内容如下：

# ========== export.py ==========
from ultralytics import YOLO

model = YOLO("yolov8s.pt")

success = model.export(format="onnx", dynamic=True, simplify=True)

cd ultralytics-main
python export.py

4. 复制模型并执行

cp ultralytics/yolov8s.onnx tensorRT_Pro-Vision/workspace
cd tensorRT_Pro-Vision
make yolo -j64

1. 下载 YOLOv8

git clone https://github.com/ultralytics/ultralytics.git

2. 修改代码, 保证动态 batch

# ========== exporter.py ==========

# ultralytics/engine/exporter.py第323行
# output_names = ['output0', 'output1'] if isinstance(self.mo