- 人工智能
- 深度学习
- 计算机视觉
【免费下载链接】tensorrtx
Implementation of popular deep learning networks with TensorRT network definition API
导读
本文以 yolo11/readme.md 为主体,系统讲解在 tensorrtx 开源仓库中如何将 Ultralytics YOLO11(n/s/m/l/x 全尺寸)部署到 NVIDIA TensorRT-10 平台,覆盖检测(det)、分类(cls)、分割(seg)、关键点(pose)、旋转框(obb)五大任务的完整链路:PyTorch 权重转 .wts → CMake 编译 → 引擎序列化 → C++ 推理 → Python API 推理 → INT8 量化校准。读者按本文操作可得到可复现、可运行的 YOLO11 TensorRT 推理程序,并理解其网络定义与插件实现的底层原理。
1. 项目概述:YOLO11 在 tensorrtx 中的定位
tensorrtx 是一个使用 TensorRT network definition API 手工搭建流行深度学习网络的仓库(项目描述:"Implementation of popular deep learning networks with TensorRT network definition API")。YOLO11 是其中的一个独立子项目,位于 yolo11/,其 readme 明确指出:
Yolo11 model supports TensorRT-10.
与依赖 ONNX 导出的部署方式不同,本子项目完全使用 TensorRT C++ API(INetworkDefinition、addConvolutionNd、addPoolingNd、addPluginV2等)逐层重建网络,因此对 TensorRT 版本高度敏感,当前实现锁定 TensorRT 10 系列。
1.1 支持矩阵
readme 中列出的支持范围如下:
| 任务 | 精度支持 | API 支持 | 入口程序 |
|---|---|---|---|
| YOLO11-det 目标检测 | FP32 / FP16 / INT8 | Python / C++ | yolo11_det.cpp |
| YOLO11-cls 图像分类 | FP32 / FP16 / INT8 | Python / C++ | yolo11_cls.cpp |
| YOLO11-seg 实例分割 | FP32 / FP16 / INT8 | Python / C++ | yolo11_seg.cpp |
| YOLO11-pose 姿态估计 | FP32 / FP16 / INT8 | Python / C++ | yolo11_pose.cpp |
| YOLO11-obb 旋转目标检测 | FP32 / FP16 / INT8 | Python / C++ | yolo11_obb.cpp |
从源码结构看,五个任务的 C++ 入口共用 src/ 与 include/ 下的网络定义代码(model.cpp分别提供buildEngineYolo11Det/Cls/Seg/Pose/Obb五个构建函数),差异集中在检测头(Head)的输出组织方式上,这为一次编译、多任务复用打下了基础。
1.2 目录结构
yolo11/ ├── CMakeLists.txt # 编译脚本,产出 5 个可执行文件 + libmyplugins.so ├── gen_wts.py # PyTorch .pt → .wts 权重转换脚本 ├── yolo11_det.cpp / _cls / _seg / _pose / _obb # C++ 入口 ├── yolo11_*_trt.py # Python API 推理示例 ├── include/ # 头文件:config.h / block.h / model.h / calibrator.h / preprocess.h / postprocess.h 等 ├── src/ # block.cpp / model.cpp / calibrator.cpp / preprocess.cu / postprocess.cpp / postprocess.cu └── plugin/ # yololayer.cu / yololayer.h(检测头自定义插件)2. 环境要求
readme 给出了作者实测通过的软件版本组合:
| 组件 | 版本 |
|---|---|
| CUDA | 12.9 |
| cuDNN | 9.10.2 |
| TensorRT | 10.10.0.31 |
| OpenCV | 4.8.0 |
| ultralytics | 8.3.238 |
| Python | 3.12 |
需要说明的是:这是 readme 作者声明的工作环境,并非严格的下限要求。实际部署时只要 TensorRT 10.x 主版本一致、CUDA/cuDNN 与所用 TensorRT 构建版本匹配,即可按相同流程操作。
从 CMakeLists.txt 还可以看到两个部署细节:
- 编译期默认设置
CMAKE_CUDA_ARCHITECTURES 75 80 86 89 90,即覆盖 Turing(75)、Ampere(80/86)、Ada Lovelace(89)、Hopper(90)等主流架构;请按实际 GPU 计算能力修改该列表,否则可能编译出无法在当前显卡加载的 kernel。 - 构建脚本对 aarch64(嵌入式平台,如 Jetson)与 x86 分别处理:
CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64"时链接/usr/local/cuda/targets/aarch64-linux,否则链接/usr/local/cuda与 TensorRT 安装目录。TensorRT 头文件/库路径在脚本中写死为/workspace/shared/TensorRT-10.10.0.31,与本机实际路径不一致时必须手动修改。
3. 核心配置项解析(include/config.h)
所有可调参数集中在 include/config.h,readme 也明确指出"Other configs please check src/config.h"。该文件是调整精度、输入尺寸、类别数、阈值的唯一入口。
3.1 精度控制宏
// #define USE_FP32 #define USE_FP16 // #define USE_INT8三个宏三选一:注释掉 FP32、打开 FP16 即默认以半精度构建;打开 INT8 则走量化流程(见第 8 节)。对应实现中,src/model.cpp 在构建引擎时根据宏设置BuilderFlag::kFP16或BuilderFlag::kINT8。
3.2 常用配置参数
| 宏 | 默认值 | 含义 |
|---|---|---|
kInputTensorName/kOutputTensorName | "images"/"output" | 输入/输出张量名 |
kProtoTensorName | "proto" | 分割任务的原型掩码张量名 |
kNumClass | 80 | 检测/分割类别数(COCO) |
kPoseNumClass | 1 | 姿态任务类别数(只有 person 一类) |
kNumberOfPoints | 17 | 姿态关键点数量 |
kObbNumClass | 15 | OBB 旋转框类别数 |
kObbNe | 1 | OBB 额外输出参数个数(角度) |
kBatchSize | 1 | 批大小 |
kGpuId | 0 | 使用的 GPU 编号 |
kInputH/kInputW | 640 / 640 | 检测、分割、姿态输入尺寸 |
kObbInputH/kObbInputW | 1024 / 1024 | OBB 输入尺寸 |
kNmsThresh | 0.45 | NMS IoU 阈值 |
kConfThresh | 0.5 | 置信度阈值 |
kConfThreshKeypoints | 0.5 | 关键点置信度阈值 |
kMaxInputImageSize | 3000×3000 | 预处理缓存上限 |
kMaxNumOutputBbox | 1000 | 单图最大输出框数 |
kInputQuantizationFolder | ./coco_calib | INT8 校准图片目录 |
kClsNumClass | 1000 | 分类任务类别数(ImageNet) |
kClsInputH/kClsInputW | 224 / 224 | 分类任务输入尺寸 |
需要自定义数据集(例如改为 20 类)时,只需修改kNumClass并重新构建引擎;OBB 与分类任务的输入尺寸与检测不同,已在配置中独立控制。
4. 权重转换:PyTorch .pt → .wts
tensorrtx 的通用约定是:先用 Python 脚本把 PyTorch 权重序列化为自定义的.wts文本格式,C++ 侧再通过loadWeights(见 src/block.cpp)按参数名读取并构建nvinfer1::Weights。
4.1 转换脚本参数
gen_wts.py 提供三个命令行参数:
| 参数 | 必填 | 说明 |
|---|---|---|
-w/--weights | 是 | 输入 .pt 权重路径 |
-o/--output | 否 | 输出 .wts 路径,缺省为权重同名文件 |
-t/--type | 否(默认detect) | 任务类型,可选detect / cls / seg / pose / obb |
4.2 转换原理
脚本内部执行如下关键步骤(gen_wts.py第 40-56 行):
torch.load(pt_file, map_location='cpu')['model'].float()加载模型并强制转为 FP32;- 对
detect/seg/pose/obb任务,读取检测头锚点model.model[-1].anchors,计算anchor_grid = anchors * stride后将其从模型中移除(delattr),因为在 TensorRT 侧由插件统一处理锚框; - 遍历
model.state_dict(),第一行写入参数总个数,之后逐行写出参数名 参数个数及每个浮点数的 IEEE-754 单精度十六进制表示(struct.pack('>f', float(vv)).hex())。
.wts是明文文本文件,行数与参数个数一一对应;C++ 侧 loadWeights 按同一格式解析为std::map<std::string, Weights>。
4.3 生成五类权重
readme 给出的完整命令流程为(以 nano 为例,其余尺寸同理):
# 1. 准备 ultralytics 训练代码与预训练权重(下载 ultralytics v8.3.238 源码包, # 并从官方 Assets 发布页分别获取以下 .pt 文件) # yolo11n.pt (检测) # yolo11n-cls.pt (分类) # yolo11n-seg.pt (分割) # yolo11n-pose.pt (姿态) # yolo11n-obb.pt (旋转框) # 2. 将转换脚本拷贝到 ultralytics 目录并依次转换 cp [PATH-TO-TENSORRTX]/yolo11/gen_wts.py . python gen_wts.py -w yolo11n.pt -o yolo11n.wts -t detect python gen_wts.py -w yolo11n-cls.pt -o yolo11n-cls.wts -t cls python gen_wts.py -w yolo11n-seg.pt -o yolo11n-seg.wts -t seg python gen_wts.py -w yolo11n-pose.pt -o yolo11n-pose.wts -t pose python gen_wts.py -w yolo11n-obb.pt -o yolo11n-obb.wts -t obb每个-t类型必须与 .pt 权重自身任务一致,否则检测头结构对不上,后续构建引擎时会因权重缺失而断言失败。转换后每个任务得到一个.wts文件。
5. 编译构建
readme 的构建步骤:
cd [PATH-TO-TENSORRTX]/yolo11 mkdir build cd build cmake .. make编译成功后会生成:
- 五个可执行文件:
yolo11_det、yolo11_cls、yolo11_seg、yolo11_pose、yolo11_obb; - 动态库
libmyplugins.so:由 plugin/yololayer.cu 编译而来,包含检测头自定义插件YoloLayer_TRT,供引擎序列化与 Python 推理加载使用。
从 CMakeLists.txt 可见,myplugins仅链接nvinfer与cudart,各可执行文件再额外链接myplugins与 OpenCV。编译前请确认:
- CUDA 架构列表与实际显卡匹配;
- TensorRT 头文件/库路径正确(脚本默认
/workspace/shared/TensorRT-10.10.0.31); - OpenCV 可被
find_package(OpenCV)找到。
6. 命令行参数约定与推理流程
五个可执行文件遵循统一的-s/-d参数协议,以 yolo11_det.cpp 中的parse_args(第 122-167 行)为例:
./yolo11_det -s [.wts] [.engine] [n/s/m/l/x] # 序列化:由 .wts 构建并保存 engine ./yolo11_det -d [.engine] ../images [c/g] # 反序列化:加载 engine 并批量推理-s:构建引擎并写入 plan 文件。serialize_engine内部创建IBuilder,调用buildEngineYolo11Det后把IHostMemory数据写入二进制文件。-d:从 plan 文件反序列化引擎,逐批读入图片执行推理,结果图保存在 build 目录下(输出文件名为_原图名)。[c/g]表示后处理位置:c表示在 CPU 上做解码 + NMS(batch_nms);g表示在 GPU 上通过cuda_decode+cuda_nms(见 src/postprocess.cu)完成解码与 NMS,再由batch_process回读结果。注意g模式当前不支持kBatchSize > 1(代码中会直接退出)。
6.1 子模型 n/s/m/l/x 的网络超参
第五个参数[n/s/m/l/x]决定网络宽度与深度。yolo11_det.cpp中硬编码了与 Ultralytics YOLO11 各尺寸一致的缩放系数:
| 子模型 | depthgd | widthgw | max_channels | 内部type |
|---|---|---|---|---|
| n | 0.50 | 0.25 | 1024 | n |
| s | 0.50 | 0.50 | 1024 | s |
| m | 0.50 | 1.00 | 512 | m |
| l | 1.00 | 1.00 | 512 | l |
| x | 1.00 | 1.50 | 512 | x |
这些系数在 src/model.cpp 中通过get_width(x, gw, max_channels)(通道数取 8 的倍数向上取整)与get_depth(x, gd)(深度四舍五入且结果不小于 1)换算成实际网络层数/通道数,从而在不改动代码的情况下复用同一套网络定义构建不同尺寸的模型。
6.2 检测任务(Detection)
cp [PATH-TO-ultralytics]/yolo11n.wts . # 构建并序列化 TensorRT 引擎 ./yolo11_det -s yolo11n.wts yolo11n.engine n # 运行推理(c = CPU 后处理,g = GPU 后处理) ./yolo11_det -d yolo11n.engine ../images c # 结果保存在 build 目录推理主循环(yolo11_det.cpp第 225-253 行)按kBatchSize分批次cv::imread读图,经cuda_batch_preprocess(src/preprocess.cu)完成缩放填充到 640×640,然后enqueueV3异步推理,最后draw_bbox绘制并cv::imwrite输出。
6.3 分类任务(Classification)
分类使用独立的 224×224 输入,输出经全局平均池化 + 全连接(model.10,见buildEngineYolo11Cls)得到 1000 类得分:
cp [PATH-TO-ultralytics]/yolo11n-cls.wts . # 构建并序列化 TensorRT 引擎 ./yolo11_cls -s yolo11n-cls.wts yolo11n-cls.engine n # 下载 ImageNet 类别标签文件(imagenet_classes.txt) wget <imagenet_classes.txt 下载地址> -O imagenet_classes.txt # 运行推理 ./yolo11_cls -d yolo11n-cls.engine ../images分类程序读入标签文件后,在推理结果中按 top-1/top-5 输出类别名称与置信度。
6.4 分割任务(Segmentation)
分割任务的引擎输出除了检测框外,还包含 32 维掩码系数(Detection.mask[32],见 include/types.h),推理时与proto特征相乘还原掩码:
cp [PATH-TO-ultralytics]/yolo11n-seg.wts . # 构建并序列化 TensorRT 引擎 ./yolo11_seg -s yolo11n-seg.wts yolo11n-seg.engine n # 下载 COCO 类别标签文件 wget -O coco.txt <coco-labels.txt 下载地址> # 运行推理(第三个参数 c 为后处理位置,第四个参数为标签文件) ./yolo11_seg -d yolo11n-seg.engine ../images c coco.txt6.5 姿态任务(Pose)
姿态输出包含 17 个关键点的(x, y, conf)三元组,配置中kNumberOfPoints = 17、kConfThreshKeypoints = 0.5:
cp [PATH-TO-ultralytics]/yolo11n-pose.wts . # 构建并序列化 TensorRT 引擎 ./yolo11_pose -s yolo11n-pose.wts yolo11n-pose.engine n # 运行推理 ./yolo11_pose -d yolo11n-pose.engine ../images6.6 旋转框任务(OBB)
OBB 是特殊任务:输入尺寸为 1024×1024(kObbInputH/W),类别数 15(kObbNumClass),且Detection结构体带angle字段用于角度回归:
cp [PATH-TO-ultralytics]/yolo11n-obb.wts . # 构建并序列化 TensorRT 引擎 ./yolo11_obb -s yolo11n-obb.wts yolo11n-obb.engine n # 下载示例图片并放入 images 目录 wget -O P0015.png <P0015.png 下载地址> mv P0015.png ../images # 运行推理 ./yolo11_obb -d yolo11n-obb.engine ../images7. 网络定义的源码级原理
7.1 基础算子封装(src/block.cpp)
src/block.cpp 用 TensorRT API 手工实现了 YOLO11 的全部基础模块:
- convBnSiLU(第 74-94 行):Conv2d + BatchNorm + SiLU。其中 BatchNorm 被数学等价地折叠为
IScaleLayer(addBatchNorm2d,第 40-72 行),将gamma/beta/mean/var融合成 scale/shift 两个向量,避免显式计算 BN 算子; - C3K2(第 239-275 行):对应 YOLO11 主干中的 C3k2 模块,通过
addSlice把cv1输出按通道一分为二,一侧串联多个bottleneck(或 C3k),再与另一侧addConcatenation汇合,最后cv2卷积; - C2PSA(第 380-415 行):YOLO11 引入的 PSA(位置敏感自注意力)模块。从源码可以推断,其核心是
PSABlock(第 357-378 行)内的Attention(第 293-355 行):qkv卷积 → reshape 多头 → 矩阵乘注意力图 → Softmax → 与v加权 → 叠加可学习的相对位置编码pe(Depthwise Conv)→proj卷积,再串接两层 FFN,每个子层都带 shortcut 残差; - SPPF(第 113-136 行):三个串联的最大池化(5×5、stride=1、padding=2)后与原始特征
addConcatenation,再 1×1 卷积; - DFL(第 138-157 行):分布聚焦损失对应的解码分支,将 4×16 的分布通过 Softmax 与 1×1 卷积加权得到 4 个回归量;
- DWConv(第 417-437 行):深度可分离卷积(
conv->setNbGroups(ch)),用于 PSA 中的编码分支。
7.2 检测头插件(plugin/yololayer.cu)
三个尺度的特征图经拼接后接入自定义插件YoloLayer_TRT(addYoLoLayer,block.cpp 第 159-218 行)。插件通过PluginField传入combinedInfo数组,包含类别数、关键点数、输入尺寸、kMaxNumOutputBbox以及 seg/pose/obb 标志位,再配合各尺度网格大小px_arry。插件在 GPU 上一次性完成 anchor 解码、多尺度合并与阈值过滤,输出统一为[kMaxNumOutputBbox]长度的Detection结构(见 include/types.h),CPU 后处理只需做最终的 NMS。
7.3 尺寸换算与推理缓冲
src/model.cpp 中get_width/get_depth负责把 n/s/m/l/x 的 gd/gw 映射为实际通道与深度;C++ 入口则按kOutputSize = kMaxNumOutputBbox * sizeof(Detection)/sizeof(float) + 1分配输出缓冲,并在运行时通过engine->getTensorShape读取实际输出框数model_bboxes(yolo11_det.cpp第 206-207 行)。
8. INT8 量化
readme 给出了 INT8 量化三步流程:
- 准备校准图片:可从训练集随机挑选 1000 张以上,或直接下载作者提供的 COCO 校准集
coco_calib; - 解压到构建目录:将校准图片目录放入
yolo11/build下(路径必须与kInputQuantizationFolder一致,默认./coco_calib); - 开启量化并重新编译:
# 编辑 include/config.h # 注释 #define USE_FP16,取消注释 #define USE_INT8 makeconfig.h中默认// #define USE_INT8即关闭状态,打开后 src/model.cpp 会:
- 检查
builder->platformHasFastInt8()并断言; - 设置
BuilderFlag::kINT8; - 创建
Int8EntropyCalibrator2(src/calibrator.cpp),按kInputW/kInputH批量读取coco_calib目录中的图片完成熵校准,校准表保存为int8calib.table。
校准完成后,后续-s序列化生成的 engine 即 INT8 引擎,推理流程与 FP16 完全一致。需要强调:INT8 引擎对当前输入尺寸与校准集敏感,更换分辨率或类别数后应重新校准。
9. Python API 推理
readme 提供可选的第 3 步:在 Python 中加载 C++ 序列化好的 engine 与插件库进行推理。
# 安装 python-tensorrt、pycuda 等依赖 # 确保 build 目录中已有 yolo11n.engine 与 libmyplugins.so python yolo11_det_trt.py ./build/yolo11n.engine ./build/libmyplugins.so9.1 常见坑与修复
readme 明确记录了两个环境相关 FAQ:
- Windows 下
pycuda._driver.LogicError; - Linux 下
Segmentation fault。
修复方式相同:在 Python 文件头部显式初始化 PyCUDA 运行时:
import pycuda.autoinit import pycuda.driver as cuda这正是 yolo11_det_trt.py 第 13-14 行的写法——pycuda.autoinit会在模块导入时完成 CUDA 上下文初始化,pycuda.driver提供显存与流管理接口。脚本还通过ctypes.CDLL加载libmyplugins.so注册插件(否则反序列化 engine 时找不到YoloLayer_TRT),并定义了与 C++ 一致的CONF_THRESH = 0.5、IOU_THRESHOLD = 0.4等常量,以及YoLo11TRT类封装预处理、推理与后处理。
yolo11_seg_trt.py、yolo11_pose_trt.py、yolo11_obb_trt.py、yolo11_cls_trt.py与检测脚本结构一致,分别按各自任务解析输出张量(分割取proto+ 32 维掩码系数,姿态取 17×3 关键点,OBB 取角度)。
10. 总结与注意事项
基于 readme 与源码,将 YOLO11 在 tensorrtx 上的部署要点归纳如下:
- 版本锁定:网络使用 TensorRT 10 的显式 batch API 与
enqueueV3接口构建,不要混用 TensorRT 7/8 的旧 API 风格; - 权重格式:
.wts是 tensorrtx 自有的文本权重格式,必须与 gen_wts.py 的-t任务参数严格对应; - 配置唯一入口:精度、输入尺寸、类别数、NMS/置信度阈值、量化目录等全部在 include/config.h 修改,改后需重新
make; - 后处理两种模式:
c(CPU)通用且支持多 batch,g(GPU)性能更优但当前仅支持单 batch; - INT8 需要校准:量化前务必准备足够的校准图片,量化后的 engine 在更换输入分辨率或数据集后需要重新校准;
- Python 推理:必须同时提供 engine 与
libmyplugins.so,且要显式import pycuda.autoinit,这是 Windows 与 Linux 上常见异常的根因。
若需查阅同一部署框架下的其他模型,可参见仓库根目录 README.md 与其他子项目(如 yolov8/、yolov10/、yolov13/)的 readme,它们遵循相同的.wts→ engine → 推理的整体流程。
- 人工智能
- 深度学习
- 计算机视觉
【免费下载链接】tensorrtx
Implementation of popular deep learning networks with TensorRT network definition API
相关推荐
tensorrtx 中 YOLOv10 的 TensorRT-10 部署指南:从 .pt 权重到 FP32/FP16/INT8 引擎
tensorrtx 中 YOLOv10 的 TensorRT 10 部署指南:从 .pt 权重到 FP32/FP16/INT8 引擎 导读 本文基于 tenso
人工智能深度学习计算机视觉tensorrtx 中 YOLOv8 的 TensorRT-10 完整部署指南:从 .pt 权重到 FP32/FP16/INT8 推理
tensorrtx 中 YOLOv8 的 TensorRT 10 完整部署指南:从 .pt 权重到 FP32/FP16/INT8 推理 导读 本文围绕 tens
人工智能深度学习计算机视觉TensorRT 部署 YOLOv3-Tiny 完整指南:wts 权重转换、引擎构建与推理(tensorrtx 实战)
TensorRT 部署 YOLOv3 Tiny 完整指南:wts 权重转换、引擎构建与推理(tensorrtx 实战) YOLOv3 Tiny 是 YOLOv3
人工智能深度学习计算机视觉
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考