ARTICLE DETAIL

资讯详情

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

TensorRT部署YOLOv5实例分割:从ONNX导出到engine后处理实战

TensorRT部署YOLOv5实例分割:从ONNX导出到engine后处理实战 简介基于TensorRT的YOLOv5实例分割部署源码包面向有C/CUDA基础、希望深入学习模型推理加速的开发者也适合计算机、电子、数学等专业学生作为课程设计或毕业设计的参考资料。压缩包共180个文件整体约19.64MB以编译目标文件、Makefile与Shell脚本、JSON与Proto配置、CUDA Kernel源文件为主并包含预处理算子、HSigmoid/HSwish激活函数、H264码流转换、Base64图像传输等模块目录结构清晰能对照完整工程理解TensorRT部署链路。目前已有1591人学习下载。通过该源码可学习TensorRT C API调用、自定义CUDA算子注册与集成、多模块协同编译等关键技巧同时可借鉴其中基于mongoose的嵌入式服务与图像上传接口设计快速迁移到实际项目中。资料作为“参考资料”而非定制需求要求使用者具备相应基础能够自行调试并扩展功能适合有独立排错能力的学习者。1. 基于TensorRT部署YOLOv5实例分割engine文件只是部署链路的一半把一个训练好的YOLOv5实例分割模型搬到TensorRT比纯检测模型多出来的工作量几乎全在mask这条输出线上。检测头返回的不只是框和类别还有一个包含32个mask系数的张量外加一张stride 4的proto特征图后处理必须把两者做矩阵乘再sigmoid才能还原出实例掩码。正因为多了这层很多网上流传的TensorRT实例分割源码包真正难复现的地方不在engine构建而在输出张量名、动态shape、FP16精度这几个容易对不齐的点上。这篇文章按模型导出、引擎构建、后处理、工程落地四个阶段讲参数和命令可直接抄适合要自己训练数据集、再部署到Jetson Orin或PC端GPU的读者。2. YOLOv5实例分割模型导出输出张量形状与TensorRT版本匹配动手构建engine之前先把ONNX文件的两个输出搞清楚。YOLOv5 v7.0的segment分支沿用yolov5网络结构的backbone和PAN neck只在检测头旁边多了一条proto分支这条分支的输出在TensorRT里是一个独立绑定的张量名字、形状、dtype都要与解析代码一一对应。2.1 检测头与proto分支的输出约定以640×640输入为例名为output的输出形状是[1, 25200, 41nc32]。25200的来历是P3/P4/P5三个尺度的网格数之和80×80加40×40加20×20等于8400每格预测3个anchor得到25200组预测。每组前4位是xywh框第5位是objectness接着nc位是类别分数最后32位是mask系数。另一个输出proto形状固定为[1, 32, 160, 160]对应stride 4的mask原型图。推理阶段用检测头保留的mask系数与proto做矩阵乘结果经过sigmoid就得到每个实例在160×160网格上的粗掩码。实例分割的后处理就是从这一个张量出发裁出目标区域再放大回原图。源码里这两个输出通常保持默认名output和proto。如果你在训练时改过模型yaml里的depth_multiple或width_multiple输出张量的通道数会跟着变解析代码里的硬编码数字都要随之修改这是最容易埋雷的地方。2.2 segment/export.py导出ONNX的命令与参数YOLOv5官方仓库里实例分割相关的训练和导出脚本在segment子目录下不能用根目录的export.py那只会导出检测模型。python segment/export.py --weights yolov5s-seg.pt \ --include onnx --opset 12 --simplify --dynamic--include指定导出格式为onnx--opset 12在TensorRT 8.x上兼容性最好--simplify会调用onnx-simplifier清理导出时残留的冗余算子--dynamic让batch维度保持动态。各参数的作用见下表。参数推荐值作用--opset12TRT 8.5及以上稳定支持opset 17需要TRT 8.6--simplify开启消除冗余节点减小engine体积--dynamic开启让batch可动态建议高宽固定为640--includeonnx只生成onnx文件提示导出后立即检查输出名和形状确认images这个输入名保留。TensorRT解析时按名字绑定张量拼错一个字母都要查半天。2.3 Jetson Orin上的TensorRT版本选型与降版本TensorRT的engine文件与构建时的TRT版本强绑定跨版本直接加载会报invalid engine或版本不匹配。Jetson Orin上TRT版本由JetPack决定常见组合关系如下表。平台JetPack/CUDA版本自带TensorRT建议ONNX opsetJetson OrinJetPack 5.1.2 / CUDA 11.48.5.212Jetson OrinJetPack 6.0 / CUDA 12.28.6.312 或 17PC端Ampere/AdaCUDA 12.x10.x17网上搜orin降tensorrt版本多数情况发生在JetPack重刷之后系统里的TRT版本和旧engine构建时不匹配重建又遇到CUDA小版本不一致。降版本的做法是用pip指定版本号重装装完执行python -c import tensorrt; print(tensorrt.version)验证。yolov5环境配置阶段就把tensorrt版本写进requirements锁定能避免团队协作时每个人构建出不同行为的engine。2.4 自己训练数据集后的输出形状检查如果用yolov5训练自己的数据集改了data.yaml里的nc后output最后一维变成37nc也就是41nc32。类别顺序以训练时的data.yaml为准部署端解码时类别索引必须与之一致否则会出现框框对、类别全错的诡异现象第一反应往往是数据加载问题实际是索引没对齐。超参数文件data/hyps/hyp.scratch-low.yaml里的增强配置不会改变输出形状但会改变最终权重因此模型重训后必须重新执行一次导出。导出一分钟的事别省。import onnx m onnx.load(yolov5s-seg.onnx) for out in m.graph.output: shape [d.dim_value for d in out.type.tensor_type.shape.dim] print(out.name, shape)这段脚本在导出后立刻确认输出名与形状。动态batch维的dim_value是0看到0不用慌只要名字是output和proto、最后一维等于37nc就可以进入构建步骤。3. 构建TensorRT引擎trtexec与Python API两条路径拿到验证过的ONNX后下一步是让TensorRT做图优化和层融合落成engine文件。常见做法是先用trtexec命令行把流程跑通再到代码里用Python API做同样的事两条路径产出的engine等价参数含义也一致。3.1 trtexec最小命令与min/opt/max语义/usr/src/tensorrt/bin/trtexec \ --onnxyolov5s-seg.onnx \ --saveEngineyolov5s-seg_fp16.engine \ --fp16 \ --minShapesimages:1x3x640x640 \ --optShapesimages:1x3x640x640 \ --maxShapesimages:4x3x640x640三个Shapes参数只针对输入张量imagesproto输出由TRT根据网络自动推导不需要手动写profile。min/opt/max分别代表推理时可接受的最小batch、性能优化基准batch和显存上限batch建议值如下表。参数建议影响--minShapes1低于它的batch直接报错--optShapes实际最高频batchTRT选kernel和tiling的依据--maxShapes4~8视显存超过上限会触发重新profile或失败如果业务确认只跑单batch我一般会把高宽直接固定在640连动态batch都不开。这样proto固定为[1, 32, 160, 160]后处理里的缩放因子恒定整个链路的可复现性高很多。实例分割比检测更值得这么做因为proto的空间维度跟随输入分辨率变化动态高宽会让后处理代码多出一堆分支。3.2 用Python API在程序内构建engineimport tensorrt as trt def build_engine(onnx_path, save_path, max_batch4): 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(onnx_path, rb) as f: assert parser.parse(f.read()), ONNX解析失败 config builder.create_builder_config() if builder.platform_has_fast_fp16: config.set_flag(trt.BuilderFlag.FP16) profile builder.create_optimization_profile() profile.set_shape(images, (1, 3, 640, 640), (1, 3, 640, 640), (max_batch, 3, 640, 640)) config.add_optimization_profile(profile) serialized builder.build_serialized_network(network, config) if serialized is None: raise RuntimeError(engine构建失败检查ONNX与TRT版本) with open(save_path, wb) as f: f.write(serialized)create_network必须加EXPLICIT_BATCH标志因为导出的ONNX带动态batch不加会解析失败。platform_has_fast_fp16用于确认当前GPU能跑FP16桌面卡和Jetson基本都是True。build_serialized_network失败时经常不抛异常而是返回None显式判一下更稳妥。workspace内存默认按显存分配在Jetson这类小显存设备上建议显式限制config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 2 30)2GB的临时缓冲区对yolov5s-seg级模型足够也避免把整个显存吃满。3.3 FP16与INT8的精度取舍FP16对实例分割mask的精度损失通常在可接受范围实测常见mAP掉点不到一个点边缘轮廓偶尔出现锯齿。INT8想保住mask质量需要准备500到1000张有代表性的校准图跑EntropyCalibrator2而且分割这种细节敏感任务对校准集分布特别挑剔人物占比高的业务必须把人物图放进去否则背景mask会成片漏检。两种精度的取舍经验如下表。精度是否需要校准集端到端加速经验值mask掉点风险FP16否相对FP32约1.2~1.8倍低个别边缘锯齿INT8是500张以上约2.5~3倍高需逐类验证我的建议是第一步只开FP16先把整条流水线跑通、精度对齐验证满足指标后再考虑INT8。另外注意后处理里矩阵乘的dtypenumpy用float32足够别把mask系数和proto乘出来后又转成float16累加那点精度损失在crop后会被放大mask边界会布满噪点。3.4 构建期常见报错与定位ONNX解析失败报Could not parse ONNX model先查opsetopset 12的模型在TRT 8.5以下版本会解析不过版本没问题就回导出步骤重新导一次。报Cannot find input: images输入名不是images用2.4节的形状脚本看实际名按实际名改profile。报显存不足把maxShapes降下来或给workspace设上限Jetson上优先用set_memory_pool_limit。engine构建成功但推理结果全是0多半是输入预处理或输出buffer形状不匹配先回到单batch跑一遍trtexec看输出是否正常。把这几类错误在源码工程里做成预处理检查直接抛中文异常比让下游调用方看一份英文TRT日志高效得多。4. YOLOv5实例分割后处理mask解码、NMS与精度对齐引擎推理返回的两个原始张量不能直接当作可视化结果。output里25200组预测要先筛置信度再过NMSproto要和保留下的mask系数做矩阵乘。这一章给的numpy实现直接对标yolov5仓库里的process_mask逻辑。4.1 从output张量筛选目标框与置信度import numpy as np def decode_detection(output, nc, conf_thres0.25): pred output[0] # [25200, 41nc32] cls_score pred[:, 5:5 nc] cls_id cls_score.argmax(1) conf pred[:, 4] * cls_score[np.arange(len(pred)), cls_id] keep conf conf_thres boxes_xywh pred[keep, :4].copy() scores conf[keep] coeffs pred[keep, 5 nc:5 nc 32] # mask系数 # xywh转xyxyYOLOv5的框是中心点加宽高 boxes boxes_xywh.copy() boxes[:, 0] boxes_xywh[:, 0] - boxes_xywh[:, 2] / 2 boxes[:, 1] boxes_xywh[:, 1] - boxes_xywh[:, 3] / 2 boxes[:, 2] boxes_xywh[:, 0] boxes_xywh[:, 2] / 2 boxes[:, 3] boxes_xywh[:, 1] boxes_xywh[:, 3] / 2 return boxes, scores, cls_id[keep], coeffs导出后的ONNX里output张量已经带解码后的像素坐标和sigmoid过的objectness与类别分数这里直接相乘得到综合置信度不需要再做一次激活。先过滤掉绝大多数背景候选把25200筛到几百个再做NMS。mask系数的索引起点是5加nc单独截出来留给下一步解码。4.2 用proto与mask系数解码实例掩码def decode_masks(protos, coeffs, boxes, img_size640): c protos.shape[0] masks (coeffs protos.reshape(c, -1)) # [N, 160*160] masks 1.0 / (1.0 np.exp(-masks)) # sigmoid masks masks.reshape(-1, 160, 160) # 每个实例一张160x160 ratio 160 / img_size for i, (x1, y1, x2, y2) in enumerate(boxes): x1, y1, x2, y2 int(x1 * ratio), int(y1 * ratio), int(x2 * ratio), int(y2 * ratio) x1, y1 max(x1, 0), max(y1, 0) x2, y2 min(x2, 160), min(y2, 160) crop masks[i, y1:y2, x1:x2].copy() masks[i] 0 masks[i, y1:y2, x1:x2] crop return masks这段对应yolov5里的scaled_boxes加crop_mask。关键点是缩放比例必须用proto尺寸除以原图尺寸输入640时是四分之一输入1280时变成八分之一。矩阵乘方向别写反coeffs在左、proto展平在右形状是[N, 32]乘[32, 25600]。crop后对每个实例做cv2.resize到原图尺寸再按0.5阈值转成二值mask就是最终分割结果。注意如果原图经过letterbox处理box和mask坐标都在letterbox后的画布上映射回原图时要先减去pad偏移再除以缩放比这一步漏掉会导致整幅图的mask整体错位。4.3 NMS实现与阈值参数表NMS直接用numpy写就行不必为此引入torch。按分数降序依次保留当前最高分框扔掉与它IoU超过阈值的其余框循环到候选为空def nms(boxes, scores, iou_thres0.45): order scores.argsort()[::-1] keep [] while order.size 0: i order[0] keep.append(i) ious box_iou(boxes[i], boxes[order[1:]]) order order[1:][ious iou_thres] return np.array(keep)yolov5默认按类别分别做NMS上面这段是class-agnostic的实现直接整批跑会把不同类别的重叠框误删要复现默认行为就按cls_id分组后逐组调用nms。检测框和mask共用同一组筛选结果NMS在mask解码之前做能省掉大量无效矩阵乘。阈值参数沿用训练侧默认值就好参数默认值调节方向conf_thres0.25调低提升召回代价是NMS耗时上升iou_thres0.45调高抑制重叠目标密集场景调整mask阈值0.5调低mask边缘更膨大按可视化需求微调4.4 与PyTorch结果做tensor级对齐部署完第一件事是对齐。同一张图分别走yolov5的torch推理和TensorRT流水线对比框的IoU和mask的IoU。对齐脚本里最关键的是预处理完全一致letterbox的缩放、填充颜色、归一化方式都不能有微小差异很多说TensorRT精度不行的结论最后查出来都是预处理没复制到位。def mask_iou(a, b): inter np.logical_and(a, b).sum() union np.logical_or(a, b).sum() return inter / max(union, 1)我一般以box IoU大于0.9、mask IoU大于0.85作为通过线。不一致时先查预处理再查NMS顺序最后才怀疑FP16。顺序颠倒会把半小时的排错拖成一下午。5. 源码工程落地显存管理、动态batch与精度回归技巧engine有了、后处理通了最后是把几块拼进真正能长期跑的源码工程里。这里讲三个最影响生产的落地细节都是只影响稳定性、不影响demo正确性的点。5.1 engine生命周期与输出buffer预分配engine加载一次后整个进程复用context按线程各持一个避免多线程推理时上下文互踩。输出buffer按maxBatch预分配推理循环里只做host与device之间的拷贝不要每次推理都cudaMalloc频繁申请会让显存碎片化跑一晚上之后出现莫名的分配失败。TensorRT要求buffer地址512字节对齐Python侧用cupy分配默认满足C侧cudaMalloc也满足别拿普通vector 的data直接传。推理时两个输出tensor的地址固定只在batch变化时整体换一套buffer这个约定能省掉每次查询张量名的开销。5.2 optShapes要贴近真实batch很多工程把optShapes设成maxShapes认为这样上限最高实际会让单路推理变慢。TRT按optShapes选择kernel和向量化宽度业务固定1路输入时opt设8意味着按8路的tiling策略执行单batch反而多算了无效数据。常见做法是optShapes等于生产环境出现频率最高的batchminShapes设1兜底maxShapes设成显存允许的上限应对突发。batch切换时记得调用context.set_input_shape并且确认输出buffer容量大于新batch。这个坑在压测脚本里最常见batch从1切到4程序不报错结果tensor被截断mask错得莫名其妙。5.3 把精度回归脚本变成CI门禁最后一个技巧是把4.4节的对齐逻辑固化成回归脚本放进源码仓库的scripts目录每次换权重、升TensorRT版本、改后处理参数后自动跑一遍for img, gt_mask in zip(valid_images, gt_masks): boxes, masks trt_infer(img) # TRT整条流水线 ref_boxes, ref_masks torch_infer(img) # 原PyTorch流水线 assert mask_iou(masks[0], ref_masks[0]) 0.85 assert box_iou(boxes[0], ref_boxes[0]) 0.9torch侧固定随机种子只跑eval模式数据取验证集里框数量差异大的100张图能同时覆盖单目标、密集行人、大目标占满画面三类场景。回归不通过时先看预处理差异再看阈值参数最后才动FP16。脚本失败直接让CI标红比靠肉眼对比两张可视化图可靠得多也避免上次还能跑、这次效果变了这类问题在项目里反复出现。本文还有配套的精品资源点击获取
返回列表