ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

TensorRT+YOLOv8实例分割:C++/Python跨平台部署实战

TensorRT+YOLOv8实例分割:C++/Python跨平台部署实战 简介这是一套基于 TensorRT 与 YOLO 的部署工具同时支持目标检测与实例分割提供 C 和 Python 两种接口面向开发者、算法研究人员以及有监控安防、工业质检、智能交通、医疗影像辅助诊断等落地需求的团队。压缩包共 1523 个文件大小约 314MB内含 hpp/h 头文件、cpp/cu 源码、Python 脚本、CMake 构建配置、MD 文档、预编译 dll/lib/exe 及模型 engine 文件等配套文档说明与示例相对完整可在 Linux 和 Windows 上直接搭建运行。目前已有 300 余人学习下载。除核心推理引擎与样例工程外资源还收纳了构建缓存、Pyd 扩展、Dockerfile、Whl 包及多种图像/视频样例便于二次开发与调试。整体上兼顾工程部署与算法验证能帮助用户把 YOLO 能力快速集成到产品或研究环境中。1. TensorRT 部署 YOLO 实例分割先跑通 C/Python 和 Linux/Windows 这四条路一个刚做完的视频分析项目甲方给的卡是 T4要求 1080p25fps 输入同时输出检测框和实例掩码开发机是 Ubuntu交付现场是 Windows。模型选了 YOLOv8-segPyTorch 直接推理在单路下都吃紧最后把权重过一遍 TensorRT多路才撑了下来。这个标题对应的就是这类场景用 TensorRT 部署 YOLO 实例分割与目标检测同时提供 C 和 Python 两套推理方式并覆盖 Linux 和 Windows。适合正在做工业质检、安防分析、机器人视觉的团队也适合想在 GPU 上把延迟榨干的开发者。下面从输出头解析、engine 构建、后处理到上线压测按实际落地顺序走一遍。2. 推理链路设计从 PyTorch 权重到 TensorRT engine先看清输出头和 shape 的坑2.1 为什么实例分割任务比纯检测更需要 TensorRT纯目标检测在 TensorRT 里的收益主要是卷积层融合和 FP16/INT8 量化但实例分割多了一个 mask 分支PyTorch 的原始实现里mask 的生成要在 CPU 或 GPU 上做矩阵乘、sigmoid、resize、按框裁剪这部分开销可能占整帧延迟的三到四成。TensorRT 能把检测头和 mask coefficient 分支保留在 GPU 上prototype 输出也不再往返 CPU但后处理依然得自己写。常见做法是先把 YOLOv8-seg 或 YOLOv5-seg 导出成 ONNX再交给 TensorRT 构建 engine。这个过程中最容易翻车的地方不在卷积层而在输出头检测输出、mask 系数、prototype 三者的通道数和顺序决定了后续解析代码怎么写。很多人把纯检测的解析逻辑套到分割模型上结果框是对的掩码全花。TensorRT 的 engine 对用户来说是个黑匣子构建一次后序列化成文件加载后只能按 tensor name 取数据。你必须在构建前就想清楚输入输出 tensor 的名称、shape 和数据类型否则后面每一步都在猜。所以这一章先把模型的导出结果讲透再谈构建参数。2.2 ONNX 导出YOLOv8-seg 与 YOLOv5-seg 输出头的差异导出 YOLOv8-seg 用 ultralytics 自带的接口最省事常见做法是直接跑下面这段from ultralytics import YOLO model YOLO(yolov8s-seg.pt) model.export( formatonnx, imgsz640, dynamicFalse, opset12, simplifyTrue, )参数说明formatonnx指定导出格式imgsz640固定输入分辨率训练时如果是 640 就保持一致dynamicFalse这里先关掉动态维度把 batch 和宽高都固死省掉后续 profile 的麻烦opset12是为了兼容旧版 TensorRT新版用 17 也可以但低版本更稳simplifyTrue会做一轮图优化去掉一些冗余节点。导出后用onnxruntime或netron看一眼输出节点YOLOv8-seg 通常是三个输出检测头是1×84×8400其中 4 是框坐标80 是 COCO 类别数也就是常说的 coco80mask 系数是1×32×8400表示每个 anchor 对应 32 个 mask 权重prototype 是1×32×160×160是模型学习到的 32 张基础掩码图。也有导出脚本把前两个输出拼成1×116×8400解析时要看清楚。YOLOv5-seg 的情况不同它的输出是一个大节点shape 类似1×117×25200通道组成是 4 个框坐标、80 个类别、1 个目标置信度、32 个 mask 系数。25200 是三个尺度 anchor 的总数。解析逻辑比 YOLOv8 多一步乘置信度其余思路一致。2.3 固定 shape 和动态 shapebatch、分辨率与 mask 尺寸的连带影响构建 TensorRT engine 时第一个决策是固定 shape 还是动态 shape。固定 shape 最简单输入固定为1×3×640×640优化 profile 里 min、opt、max 三个值写一样构建快运行时无需处理形状变化。动态 shape 能换来灵活性比如同一份 engine 跑 640 和 1280 两种分辨率或者一个进程里 batch 从 1 切到 8。代价是每个输入 tensor 都要配一组 min/opt/max输出 tensor 的 shape 在运行时可能随输入变化。实例分割这里有个隐藏关联prototype 输出的高宽是输入高宽的四分之一输入 640 时是 160×160输入 1280 时变成 320×320。如果在代码里写死了 160切到高分辨率后掩码生成全部错位而且不会报错只会出烂图。我一般建议第一次做先固定 batch1、固定 640 分辨率把整条推理链路跑通。等后处理稳定了再考虑动态 shape 或多 batch。用表格总结一下选型差异维度固定 shape动态 shape构建难度低profile 三个值相同高需要配 min/opt/max推理代码简单直接用固定 buffer每次要查输出 shape分辨率切换不行只能重构建可以但 mask 尺寸要跟着算多 batch需要重新构建运行时指定 batch适用场景上线后参数不变的场景算法同学频繁调输入的场景固定 shape 的 engine 换分辨率只能重新构建所以在需求还没冻结的阶段动态 shape 更省心。但动态 shape 必须配合严谨的后处理后面避坑章节会专门讲一个 batch 变化导致掩码错位的案例。3. Python 先行的最小闭环engine 构建、推理与掩码后处理的完整脚本3.1 TensorRT、CUDA、cuDNN 版本搭配先把环境摸清楚TensorRT 不是独立运行的它依赖 CUDA 和 cuDNN三个版本之间是强绑定。Ubuntu 上装 TensorRT 最常见的方式是下载 deb 包用dpkg -i安装装完ldconfig刷新库缓存Windows 上通常是解压 zip 后把lib目录加到环境变量 PATH在 CMake 里指路径。Python 侧用 pip 安装tensorrt包是最省事的但要注意 pip 包版本和本机 CUDA 版本必须匹配。常见翻车现场是import tensorrt成功trt.Builder一创建就报错原因是 pip 包自带的 CUDA runtime 和驱动不兼容。遇到这种情况先跑trt.__version__确认版本再用nvidia-smi看驱动版本不要把 CUDA toolkit 装成驱动两个是不同概念。3.2 构建 engine 的 Python 脚本从 ONNX 到 plan 文件下面是一段最小构建脚本TensorRT 8.x 到 10.x 的 API 略有差异但核心流程一致import tensorrt as trt logger trt.Logger(trt.Logger.WARNING) builder trt.Builder(logger) network builder.create_network( 1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) ) parser trt.OnnxParser(network, logger) with open(yolov8s-seg.onnx, rb) as f: ok parser.parse(f.read()) if not ok: for i in range(parser.num_errors): print(parser.get_error(i)) config builder.create_builder_config() config.set_flag(trt.BuilderFlag.FP16) # 固定 shape 的 profilemin/opt/max 相同 profile builder.create_optimization_profile() input_tensor network.get_input(0) profile.set_shape(input_tensor.name, (1, 3, 640, 640), (1, 3, 640, 640), (1, 3, 640, 640)) config.add_optimization_profile(profile) plan builder.build_serialized_network(network, config) with open(yolov8s-seg.engine, wb) as f: f.write(plan)这段代码的逻辑是先建 builder 和 network用 ONNX parser 把模型读进来然后创建 builder config开启 FP16接着创建 optimization profile把输入形状三个值都设为1×3×640×640这就相当于告诉 TensorRT 只处理这一种 shape最后build_serialized_network生成序列化后的 engine写到文件里。参数说明1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)是固定写法表示使用显式 batch 模式这是 TensorRT 7 之后的要求FP16 开关在 T4 这种卡上收益明显但如果后面精度掉点严重可以先关掉对比profile.set_shape的三个参数分别是 min、opt、max固定 shape 时三者必须一致。构建时间取决于模型大小和 TensorRT 版本YOLOv8s-seg 在 T4 上一般几十秒到几分钟。如果构建中途被杀掉常见原因是显存分配不足可以尝试降低分辨率或关掉 FP16。3.3 Python 推理与掩码后处理把三个输出组装成掩码加载 engine 和推理的代码我用 pycuda 管理显存这是社区里最常见的选择import tensorrt as trt import numpy as np import pycuda.driver as cuda import pycuda.autoinit logger trt.Logger(trt.Logger.WARNING) runtime trt.Runtime(logger) with open(yolov8s-seg.engine, rb) as f: engine runtime.deserialize_cuda_engine(f.read()) context engine.create_execution_context() # 创建固定大小的显存 buffer input_vol trt.volume((1, 3, 640, 640)) output0_vol trt.volume((1, 84, 8400)) output1_vol trt.volume((1, 32, 8400)) output2_vol trt.volume((1, 32, 160, 160)) d_input cuda.mem_alloc(input_vol * 4) d_output0 cuda.mem_alloc(output0_vol * 4) d_output1 cuda.mem_alloc(output1_vol * 4) d_output2 cuda.mem_alloc(output2_vol * 4) context.set_tensor_address(images, int(d_input)) context.set_tensor_address(output0, int(d_output0)) context.set_tensor_address(output1, int(d_output1)) context.set_tensor_address(output2, int(d_output2)) # 假设 input_np 已经完成 letterbox 和归一化 cuda.memcpy_htod(d_input, input_np.astype(np.float32).ravel()) context.execute_v2([int(d_input), int(d_output0), int(d_output1), int(d_output2)]) # 把结果拷回内存 out0 np.empty(output0_vol, dtypenp.float32) out1 np.empty(output1_vol, dtypenp.float32) out2 np.empty(output2_vol, dtypenp.float32) cuda.memcpy_dtoh(out0, d_output0) cuda.memcpy_dtoh(out1, d_output1) cuda.memcpy_dtoh(out2, d_output2)这段代码里有几个容易踩的地方TensorRT 10.x 默认用set_tensor_address加execute_v2但更老的版本用的是context.bindings或context.set_binding_shape。我用的是新版 API如果你手头是 8.2 以下版本编译期会报错到时换成旧式 binding 索引即可。拿到三个输出后检测部分的解析和 YOLO 推理相同对out0做置信度过滤和 NMS选出保留的框。然后对每个保留框取out1里对应的 32 个 mask 系数和out2里的 prototype 做矩阵乘再经过 sigmoid得到掩码def generate_masks(det_boxes, mask_coeffs, prototypes, img_hw): # det_boxes: (N, 4) 原始尺度下的框 # mask_coeffs: (N, 32) # prototypes: (1, 32, 160, 160) proto prototypes.reshape(32, -1) # (32, 25600) masks mask_coeffs proto # (N, 25600) masks 1.0 / (1.0 np.exp(-masks)) masks masks.reshape(-1, 160, 160) # 缩放到原图尺度并按 box 裁剪 for i, box in enumerate(det_boxes): x1, y1, x2, y2 box mask masks[i] # 用 cv2.resize 把 160x160 放大到原图然后截取 box 区域 # 最后做阈值 0.5 二值化 return masks核心是mask_coeffs proto这一步矩阵乘的结果就是每个像素属于当前实例的概率。参数说明prototypes的 160×160 对应输入 640×640 的四分之一下采样如果输入分辨率变了要先确认这个尺寸sigmoid 前的数值范围没有归一化直接1/(1exp(-x))即可裁剪时要把 mask 先放大到原图尺寸再截取否则直接截 160×160 会偏小。3.4 Python 侧测性能先 warmup再统计端到端延迟Python 推理稳定后性能测试也要在 Python 侧做因为后续 C 复刻时可以用这个数字做基准。测延迟时最忌讳上来就循环计时GPU 第一次推理会做 kernel 加载和显存初始化这部分时间不能算进真实延迟。import time import numpy as np # warmup 20 次 for _ in range(20): run_inference(input_np) latencies [] for _ in range(500): t0 time.perf_counter() dets, masks run_inference(input_np) latencies.append((time.perf_counter() - t0) * 1000) latencies.sort() print(avg%.2fms p50%.2fms p99%.2fms % ( np.mean(latencies), latencies[len(latencies)//2], latencies[int(len(latencies)*0.99)]))这里的run_inference包含预处理、H2D copy、execute、D2H copy 和后处理。真正上线时后处理可能放到 CPU 上多线程做所以这个数字是端到端参考最终以 C 数据为准。4. 换 C 上生产内存管理、多路并发与 Linux/Windows 双平台构建4.1 为什么 Python 验证完之后要换成 CPython 方案的开发效率高适合算法同事快速验证正确性但生产环境常驻服务的时候C 的优势非常明显没有 GIL 锁多路并发可控显存分配和释放完全由自己掌握部署现场不用装 Python 环境和一堆 pip 包。另一个现实原因是很多 Windows 工控机上没有 Python 运行时C 编译成单一 exe 加几个 dll 就能跑。不要误解为 Python 不能上生产很多中低并发场景 Python 够用。但如果你要同时跑四路以上视频流C 在 CPU 后处理和显存池复用上自由度更高这是实践里的普遍结论。4.2 C 最小推理骨架IRuntime、ICudaEngine 与 IExecutionContext下面是一个最小 C 推理类用 TensorRT 新版 API 写的#include NvInfer.h #include cuda_runtime_api.h #include fstream #include vector #include iostream using namespace nvinfer1; class YoloTensorRT { public: bool loadEngine(const std::string path) { std::ifstream f(path, std::ios::binary); if (!f.good()) return false; std::vectorchar data((std::istreambuf_iteratorchar(f)), std::istreambuf_iteratorchar()); runtime_ createInferRuntime(logger_); engine_ runtime_-deserializeCudaEngine(data.data(), data.size()); context_ engine_-createExecutionContext(); // 遍历所有 IO tensor分配显存 for (int i 0; i engine_-getNbIOTensors(); i) { const char* name engine_-getIOTensorName(i); auto dims engine_-getTensorShape(name); size_t vol 1; for (int j 0; j dims.nbDims; j) vol * dims.d[j]; void* ptr nullptr; cudaMalloc(ptr, vol * sizeof(float)); tensors_[name] ptr; context_-setTensorAddress(name, ptr); } return true; } void infer(float* host_input, void* stream) { cudaMemcpyAsync(tensors_[images], host_input, input_size_ * sizeof(float), cudaMemcpyHostToDevice, stream); context_-enqueueV3(stream); // D2H copy 由调用方按需执行 } private: ILegacyLogger logger_; IRuntime* runtime_ nullptr; ICudaEngine* engine_ nullptr; IExecutionContext* context_ nullptr; std::mapstd::string, void* tensors_; size_t input_size_ 3 * 640 * 640; };这段代码要做的事和 Python 版本一一对应createInferRuntime创建运行时deserializeCudaEngine反序列化 enginecreateExecutionContext创建执行上下文然后为每个 IO tensor 分配显存并调用setTensorAddress把地址绑定进去。enqueueV3是异步推理必须传入 CUDA stream不能传默认 stream 后立刻读取结果。参数说明输入尺寸这里写死了3*640*640如果动态 shape 需要从 engine 的 tensor shape 读取数据类型默认 float如果模型里有 INT8 输出要改sizeof对应类型setTensorAddress在 TensorRT 10.x 是标准用法9.x 开始也支持8.x 则用setBindingAddress注意版本差异。4.3 多路并发CUDA Stream 与显存池配置多路视频流并发时常见做法不是开多个 context而是多个 context 共用同一个 engine每个 context 绑定一个 CUDA stream。每个 stream 独立提交异步任务GPU 自己调度。四路输入的显存占用线性增长T4 16GB 上跑四路 640 分辨率 YOLOv8s-seg 通常不是显存先爆而是算力先饱和。显存分配策略上不要在每帧推理时cudaMalloc那是性能毒药。正确做法是启动时一次性分配循环复用。TensorRT 8.5 之后也支持cudaMallocAsync配合显存池但需要 driver 支持老卡上不一定稳定我一般还是手动预分配。多路并发还有一种思路是把多路拼成一个 batch 推理比如四路各取一帧拼成4×3×640×640。这样 GPU 利用率更高但需要所有路帧率同步否则会等最慢的那路。实际项目里我倾向于每路独立 context 加异步推理逻辑简单延迟也更均衡。4.4 Linux 与 Windows 的 CMake 构建差异C 工程用 CMake 是跨平台最省心的方式。下面是两份常见配置cmake_minimum_required(VERSION 3.18) project(yolo_trt LANGUAGES CXX CUDA) find_package(CUDAToolkit REQUIRED) # Linux 常见路径/opt/TensorRT # Windows 常见路径C:/TensorRT set(TENSORRT_ROOT /opt/TensorRT CACHE PATH TensorRT root) find_library(NVINFER_LIB nvinfer PATHS ${TENSORRT_ROOT}/lib) find_library(NVPARSER_LIB nvonnxparser PATHS ${TENSORRT_ROOT}/lib) add_executable(yolo_trt main.cpp) target_include_directories(yolo_trt PRIVATE ${TENSORRT_ROOT}/include) target_link_libraries(yolo_trt PRIVATE ${NVINFER_LIB} ${NVPARSER_LIB} CUDA::cudart) if(WIN32) target_compile_definitions(yolo_trt PRIVATE NOMINMAX) endif(WIN32)Linux 上 TensorRT 的库文件是libnvinfer.soWindows 上是nvinfer.libfind_library会自动匹配Windows 下运行时需要nvinfer.dll在 PATH 里否则 exe 起不来。NOMINMAX宏是 Windows 的经典坑不定义的话min/max会被 windows.h 里的宏替换模板代码直接编译失败。Linux 上如果链接报undefined reference优先检查 ldd 是否找到所有依赖常见问题是libcudnn.so版本不匹配Windows 上用 VSCode 配置 C/C 环境时CMake 工具链指向的编译器必须和 CUDA 版本配套VSCode 里两个不同的编译器会导致 CUDA 头文件识别混乱。5. 部署避坑指南engine 迁移、动态 shape、显存泄漏与精度掉点的 5 个排查记录5.1 现象engine 在 Linux 上构建、Windows 上加载直接崩溃一个项目里开发用 Ubuntu交付是 Windows 工控机。把 Linux 上构建好的.engine文件拷贝到 WindowsdeserializeCudaEngine返回空指针甚至直接段错误。原因是 TensorRT 的 engine 文件不是可移植文件它绑定 GPU 架构、CUDA 版本、TensorRT 版本和驱动能力。T4 的 engine 放到 RTX 3090 上都可能加载失败更别说跨操作系统。解决方法是把「构建 engine」这一步放到目标机器上执行或者交付时带上 ONNX 和构建工具在部署现场首次启动时自动构建。还有一种思路是用 Docker 固定环境但 Windows 工控机往往没有 Docker最后还是现场构建最稳。5.2 现象动态 shape 下 batch1 正常batch4 掩码错位代码里配了动态 shapemin(1,3,640,640)opt(4,3,640,640)max(8,3,640,640)。单路推理一切正常改 batch4 后检测框位置没问题掩码整体偏移有些掩码跑到别的物体上。原因是 prototype 输出 tensor 的 shape 是动态的batch4 时输出是4×32×160×160但后处理代码里把out2当成1×32×160×160来 reshape数据布局完全错位。TensorRT 在动态 shape 下不会帮你调整输出 buffer 的解析所有 shape 变化都要在推理后重新读取。解决方法是推理后调context-getTensorShape(output2)用返回值里的维度去算 volume再 reshape。不要在代码里写死任何160之类的数字。这个坑用 Python 复现更容易发现因为 NumPy reshape 超出预期时大概率报错C 里直接数组越界表现是随机内存错误更隐蔽。5.3 现象跑久了显存上涨几十小时后 OOM服务刚启动时显存占用 3GB跑一天后涨到 10GB再跑半天直接 out of memory。用nvidia-smi看是进程的显存持续增长不是波动。常见原因是推理循环里反复调用cudaMalloc或createExecutionContextTensorRT 某些版本在execute_v2失败或 buffer 地址变化时不会立刻释放旧显存泄漏点是自己的代码。解决办法是启动阶段把所有 buffer 分配好推理阶段只做 copy 和 enqueue绝不在循环里分配任何 GPU 资源。还有一类隐藏问题cudaMemcpy如果使用异步版本但没有同步 stream显存不会泄漏但会导致后续显存碎片化表现为分配越来越慢配合cudaMemPoolSetAttribute可以缓解。5.4 现象FP16 精度掉点不大但掩码边缘明显变差检测框和置信度在 FP16 下几乎没有变化但掩码的边缘出现锯齿和断点某些小物体掩码直接缺失。FP16 的精度损失对分类和框回归影响有限因为这两个任务对数值精度不敏感但 mask 分支的 prototype 是逐像素回归边缘区域的值本身就接近 0 到 1 的阈值边界半精度浮点的舍入误差会把边缘像素推到错误一侧。解决方法是先对比 FP32 engine 和 FP16 engine 的掩码输出确认掉点幅度如果影响验收可以对 mask 分支设置更高的精度TensorRT 的 per-layer precision 控制可以只让 mask 相关层保持 FP32但操作复杂。更实用的做法是后处理时对掩码做一次高斯模糊或形态学开闭运算边缘质量能提升不少代价很小。5.5 现象INT8 量化后分割输出明显退化比 FP16 还差为了压榨 T4 的算力上了 INT8 量化检测框还能用但掩码大面积失效部分实例掩码整块消失量化校准集的验证精度比 FP16 低了十几个点。原因是 INT8 量化对激活值的动态范围非常敏感mask 分支输出的分布和检测分支差异很大校准集如果只按检测场景采样prototype 部分的量化 scale 就没有代表性。解决方法是重新构造校准集故意混入不同掩码形状、不同物体尺寸的样本让 mask 分支的动态范围覆盖更全。经验值是校准集至少 500 张以上包含小物体占比高的场景量化后 mAP 掉点控制在 3 个点以内才算合格。6. 验证与压测技巧把帧率、显存和精度测成能汇报的数字6.1 压测脚本P99 延迟比平均帧率更能说明问题上线前我会先跑一段压测1000 帧推理丢掉前 50 帧 warmup统计剩余 950 帧的延迟分布。重点看 P99因为多路并发下偶发的长尾延迟会被前端放大成卡顿平均帧率 50fps 不代表每一帧都稳定。压测代码优先用 Python 写逻辑简单import time lat [] for i in range(1000): if i 50: infer_once() continue t0 time.perf_counter() infer_once() lat.append((time.perf_counter() - t0) * 1000) lat.sort() print(favg{sum(lat)/len(lat):.1f}ms p99{lat[int(len(lat)*0.99)]:.1f}ms)6.2 T4 上 1080p25fps 640 分辨率的路数估算估算路数时我习惯先测单路延迟 T再按帧间隔倒推。1080p25fps 的每帧间隔是 40ms如果单路端到端延迟 8ms理论上可以把 5 路塞进 40ms 内但 GPU 多路并发时延迟会非线性上升所以实际按 60%~70% 使用率规划标称 5 路就按 3 到 4 路设计留出余量给 CPU 后处理和视频解码。6.3 我的习惯构建与推理分离engine 按版本归档现在的教训是每次改模型权重都重新构建 engine不改文件名结果跑了几个月后没人知道线上 engine 对应哪个版本。现在我把 ONNX、engine、后处理代码三样东西打包成一个带日期和模型名字的目录构建脚本和推理脚本分开engine 永远由脚本生成不手工拷贝。希望这条习惯和前面的踩坑记录能帮到你少走弯路毕竟 TensorRT 这层黑匣子翻过一次车就知道后悔药不好找。本文还有配套的精品资源点击获取
返回列表