ARTICLE DETAIL

资讯详情

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

TensorRT加速YOLOv5与DeepSORT:从ONNX到engine的实时多目标跟踪部署

TensorRT加速YOLOv5与DeepSORT:从ONNX到engine的实时多目标跟踪部署 简介基于TensorRT部署YOLOv5与DeepSORT实现行人检测跟踪的完整工程包面向需要在NVIDIA Jetson Xavier等嵌入式平台或X86架构上进行实时多目标跟踪的算法工程师与开发者。方案将YOLOv5检测网络与DeepSORT跟踪算法结合通过TensorRTX完成模型转换与引擎生成可在低算力设备上获得更优推理性能。资源共35个文件压缩包约42.95MB主体为26个Python脚本涵盖检测器、跟踪器及demo_trt.py运行入口另含yolov5s.engine预生成引擎文件、libmyplugins.so自定义插件、环境依赖requirements.txt、配置yaml以及测试视频基本实现开箱即用。目前已有399人学习浏览便于快速验证算法效果或在此基础上二次开发。文件中还包含README与LICENSE目录结构清晰适合部署入门与项目移植参考。1. 先搞清TensorRT部署YOLOv5和DeepSORT的边界把TensorRT、YOLOv5、DeepSORT这一串名词放在一个工程里通常是在解决一个很具体的问题监控画面里的行人检测框连不成一条连续轨迹而纯用PyTorch跑推理又到不了边缘设备的实时要求。行人检测跟踪的常规做法是YOLOv5负责定位行人DeepSORT负责给每个检测框分配并维护稳定IDTensorRT负责把YOLOv5的卷积部分和ReID特征提取网络压到几十毫秒内。这个方案适合正在把自己的数据集训练出来的YOLOv5权重推上板端的人也适合给视频分析后端设计推理服务的工程师。要动手前先理解边界DeepSORT的卡尔曼滤波和匈牙利匹配仍然在CPU上执行TensorRT加速的只是检测网络和外观特征网络。这个区别直接决定了整个项目的代码结构。2. TensorRT engine生成YOLOv5从pt到onnx再到engine2.1 为什么中间必须过ONNX不能从pt直接到engineTensorRT不认PyTorch的pt文件这是刚接触部署时最常见的认知落差。pt文件里保存的是网络结构和权重但它的表达方式和TensorRT需要的计算图、kernel选择机制是两套东西。ONNX在这个链路里的角色是中间表示它把YOLOv5的推理图固定下来TensorRT再把这个图解析成自己的engine文件。常见做法是用YOLOv5官方仓库里自带的export.py导出ONNX因为导出过程会去掉训练分支、anchor decode和一些只在训练阶段有效的算子只保留推理所需的计算图。如果自己写脚本转换很容易把detect头的三个输出导成三个分支后面在TensorRT里还需要额外拼接。用原生export.py导出的模型输入名一般是images输出名随代码版本不同而变化这个差异要在解析输出前先打印确认。2.2 用export.py导出ONNX参数怎么给假设已经用YOLOv5训练过自己的数据集权重在runs/train/exp/weights/best.pt导出命令最常见的写法是cd yolov5 python export.py \ --weights runs/train/exp/weights/best.pt \ --include onnx \ --img-size 640 640 \ --batch-size 1 \ --opset 17 \ --simplify参数里有几个值得说明。--img-size要和训练时保持一致如果训练用的1280这里改成640会导致框的位置和置信度分布明显漂移--opset建议用17或18TensorRT 8.x对这些版本的支持比较成熟--simplify会调起onnxsim做常量折叠能减少一部分冗余节点。训练阶段调过的yolov5超参数不会影响导出流程只影响权重质量所以这一步不用回头动训练配置。有没有训练自己的数据集对导出操作本身没有影响区别只在类别名。行人检测和车牌识别、水果识别这类项目在导出环节走的是同一套命令后续差异都发生在后处理的类别过滤上。2.3 用trtexec生成FP16 engine顺带设置动态shape拿到ONNX之后最省事的做法是用trtexec直接构建engine不需要先写Python代码trtexec --onnxyolov5s.onnx \ --saveEngineyolov5s_fp16.engine \ --fp16 \ --minShapesimages:1x3x640x640 \ --optShapesimages:1x3x640x640 \ --maxShapesimages:1x3x640x640 \ --memPoolSize2048这里把三个shape写成了同一个值等效于固定输入分辨率。后面如果要多分辨率推理就把optShapes换成实际最常用的尺寸让TensorRT针对那个尺寸优化kernel。trtexec里经常用到的参数见下表参数作用常见取值--fp16允许使用FP16精度kernel无参数--int8开启INT8量化需要准备校准集无参数--calib指定INT8校准数据集图片目录或文件--minShapes/--optShapes/--maxShapes声明动态输入范围与推理尺寸匹配--memPoolSize构建期显存池上限旧版本叫--workspace1024~4096--builderOptimizationLevel控制层融合力度1~5默认3--verbose输出完整构建日志无参数构建成功后trtexec最后几行会输出Total Host Memory、Total Device Memory以及Average on 10 runs这类耗时数据可以直接看出engine在当前设备上的基础延迟。提示动态shape的optShapes一定要覆盖实际推理中出现最多的尺寸否则后续API里set_input_shape的取值偏离optShapes太多推理时会触发一次额外的显存重排表现为明显卡顿。2.3.1 构建失败时先看什么trtexec报错时日志很长前几十行多半是层信息和warning真正有效的内容在ERROR关键字附近。常见情况是报Invalid plugin或Type not supported。前者说明ONNX里残留了YOLOv5的特殊算子回2.2检查--simplify是否真正执行成功后者说明算子版本不被当前TensorRT支持优先考虑提高--opset重新导出而不是换更低的opset去迁就。3. DeepSORT模块的工程化改造ReID特征和匹配逻辑要分开看3.1 卡尔曼滤波和级联匹配在DeepSORT中的地位DeepSORT本身不是神经网络而是一个多目标跟踪算法框架。它用卡尔曼滤波预测每个track在当前帧的边界框位置再结合外观特征和运动特征计算检测框与track之间的代价矩阵最后用匈牙利算法做关联。YOLOv5的输出在这里只是输入每个检测框需要转换成cx, cy, aspect ratio, height四个观测值再加上速度分量组成8维状态向量。经典实现里级联匹配的马氏距离门控阈值取的是卡方分布自由度4下的0.95分位数9.4877。这个数值不是拍脑袋定的它是判断“当前观测和预测位置偏差是否仍在合理范围”的统计依据。实际工程中我们不会去改这个值而是通过调整其他参数来改善跟踪质量。参数常见值作用n_init3连续命中3帧才把track从tentative转confirmedmax_age30track在未匹配状态下保留30帧超时删除max_cosine_distance0.2外观特征余弦距离超过该值直接不参与匹配gating_threshold9.4877马氏距离门控阈值nn_budget100每个track缓存的外观特征数量上限这几个参数是联动的。max_age给大能扛住短时遮挡但track长期匹配不上时计算量会涨n_init给大能过滤掉误检但ID确认变慢。工程上我一般先固定n_init和max_age只调max_cosine_distance直到ID切换次数有明显变化再回头动前两个。3.2 ReID特征提取是否也走TensorRTDeepSORT中负责外观特征的是一个小型ReID网络输入通常是64x128或128x256的检测框裁剪图输出128或512维特征。它的参数量远小于YOLOv5但每帧需要对所有已确认track和检测框都做一次特征提取。如果完全用PyTorch在CPU上跑这部分会成为明显的短板在Jetson这类设备上我一般会把ReID网络也转成engine和YOLOv5共用同一套TensorRT上下文。也有一种折中做法ReID网络保留ONNX Runtime毕竟它结构简单转成engine节省的时间并不显著。但这个方案在显存管理上多维护了一套Runtime内存紧张的板子上并不划算。我的建议是经常跑的设备ReID也进TensorRT一次性实验项目可以用ONNX Runtime顶着。3.3 用Python API加载TensorRT engine的通用模板无论加载YOLOv5还是ReID的engine代码都是同一套我习惯先封装成一个小类import tensorrt as trt class TRTEngine: def __init__(self, engine_path): logger trt.Logger(trt.Logger.WARNING) runtime trt.Runtime(logger) with open(engine_path, rb) as f: self.engine runtime.deserialize_cuda_engine(f.read()) self.context self.engine.create_execution_context() self.inputs, self.outputs [], [] for i in range(self.engine.num_io_tensors): name self.engine.get_tensor_name(i) mode self.engine.get_tensor_mode(name) if mode trt.TensorIOMode.INPUT: self.context.set_input_shape(name, self.engine.get_tensor_shape(name)) self.inputs.append(name) else: self.outputs.append(name) def infer(self, input_tensor): # 固定shape下set_input_shape只调用一次运行时直接execute_v2 self.context.execute_v2(bindings[input_tensor.data_ptr()] self.output_bindings) return self.outputs这里有几个关键点。deserialize_cuda_engine接收的是engine文件反序列化出的字节流部署时只要这个字节流就够不需要重新构建num_io_tensors和get_tensor_mode是TensorRT 8.x之后的接口老代码里常用的num_bindings和get_binding_name虽然还能用但不建议再写。set_input_shape对固定shape引擎调用也不会报错所以统一加上。真正执行时用execute_v2bindings列表长度必须等于输入输出总数且全部是设备端地址。4. YOLOv5和DeepSORT的串联标准跟踪时序与可运行主循环4.1 先预测还是先检测跟踪更新的标准时序把两个模型放在同一个进程里最容易乱的是更新顺序。DeepSORT论文和参考实现里的标准时序是先对已有track做卡尔曼预测再用当前帧检测结果做关联最后统一更新状态。常见错误是检测结果出来就立刻调tracker.update完全跳过predict。跳过后track的位置少了一步先验目标持续运动时会慢半拍近距离行人的ID切换概率明显上升。完整时序可以用下面这段注释形式表达输入一帧图像 YOLOv5前处理 # letterbox 归一化 YOLOv5推理 # TensorRT执行 YOLOv5后处理 # 置信度过滤 NMS DeepSORT predict # 卡尔曼滤波预测每个track当前帧状态 DeepSORT update # 级联匹配 IOU匹配 特征匹配 状态更新 画出track框 # 读取track_id和bbox并绘制4.2 Python主循环一个可以直接改的最小实现按上面的时序主循环大概长这样import cv2 import numpy as np det_engine TRTEngine(yolov5s_fp16.engine) reid_engine TRTEngine(reid.engine) tracker DeepSort(modelreid_engine, max_age30, n_init3) cap cv2.VideoCapture(pedestrian.mp4) while True: ok, frame cap.read() if not ok: break blob, ratio, (dw, dh) letterbox(frame, (640, 640)) blob blob.astype(np.float32) / 255.0 dets det_engine.infer(np.expand_dims(blob, 0)) boxes, scores, class_ids postprocess(dets, conf_thres0.4, iou_thres0.5) person_mask class_ids 0 # COCO里0是person换数据集时改这里 xywhs xyxy2xywh(boxes[person_mask]) confs scores[person_mask] tracker.predict() tracks tracker.update(xywhs, confs, frame) for x1, y1, x2, y2, track_id in tracks: cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) cv2.putText(frame, str(track_id), (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2) cv2.imshow(tracking, frame) if cv2.waitKey(1) 0xFF ord(q): break这段代码的关键点在几处。letterbox返回的ratio和(dw, dh)必须保存下来把检测框映射回原图时要除以ratio、减去padding漏掉这一步画出来的框会整体偏向右下角。postprocess读取YOLOv5导出的ONNX输出时shape通常是1x25200x8585对应4个框坐标加1个objectness再加80个类别置信度新版代码也可能输出成3个分支所以第一次运行时打印输出shape再写后处理。tracker.predict()的位置必须在update之前不同repo的DeepSort封装接口名称不一定相同对接时先打印tracker的成员方法确认一下比对着文档猜更稳妥。4.3 前后处理的开销也不要忽视整个管线里TensorRT推理可能只占15ms但Python侧的letterbox、NMS和类型转换经常要吃掉3到5ms多路视频流时这个数字还要乘上。有人只盯着engine的推理耗时整体帧率却上不去多半就卡在这段。如果目标是把端到端延迟压到20ms以内常见做法是把letterbox改成GPU上的缩放和归一化NMS换成TensorRT自带的efficientNMS插件让输出张量直接给出最终框和分数。后处理还有一个和数据集相关的坑用自己数据集训练YOLOv5时行人这一类的class_id不一定是0。车牌识别、水果识别这类项目跟踪目标根本不在COCO的80类里所以postprocess里写死“class_ids 0”之前先打印一次检测结果的类别索引确认它对应的是什么再决定过滤条件。5. 性能调优、显存管理和TensorRT版本兼容的实战技巧5.1 FP16和INT8的取舍YOLOv5转engine时FP16是性价比最高的选项绝大多数GPU上推理延迟能比FP32降一半左右对行人检测这类任务几乎看不出精度差异。INT8收益更大但代价是校准。校准集的选择比量化算法本身更影响结果至少准备200张和真实场景光照接近的图片分辨率不需要和训练一致但内容分布必须接近线上数据。行人属于小目标比大目标更容易在INT8下漏检所以常规路径是先上FP16观察漏检率再决定是否值得上INT8。精度相对延迟显存占用常见问题FP321.0x高延迟高很少用于部署FP16约0.5x-0.7x中精度损失很小INT8约0.3x-0.5x低小目标漏检校准集要求高表格里是相对值不要直接拿别处的数字当自己机器的结果。以trtexec输出的Average on 10 runs为准而且同一台设备上不同TensorRT版本跑出的数值可能差百分之十几。5.2 多路视频流下怎么组织TensorRT推理多路摄像头场景里最常见的两种做法是batch推理和多context并行。batch推理是把N路帧拼成一个Nx3x640x640的张量做一次推理吞吐高但某一路输入晚到会拖慢整体多context是同一个engine创建多个context各自绑一路流代码简单延迟也更稳定。我一般优先选多context因为行人检测跟踪更看重延迟稳定性而不是纯吞吐。# 常见做法为每路流创建独立context共享同一个engine contexts [det_engine.engine.create_execution_context() for _ in range(stream_count)]显存不够时再切到batch方案。切换时要注意TensorRT的batch维度并不总是等于输入shape的第一维有些插件对固定维度的优化比动态维度激进批量模式下建议回trtexec重新构建一次专用engine。5.3 Jetson Orin上降TensorRT版本或换版本时要注意的坑板端部署最常遇到的是版本问题。Jetson上的TensorRT不是独立安装的它和JetPack的CUDA、cuDNN、libnvinfer_plugin.so绑在一起。直接pip install一个其他渠道的tensorrt版本运行时大概率报符号找不到或者版本mismatch。有人问Orin要不要降TensorRT版本我的建议是先确认当前版本再决定dpkg -l | grep -i nvinfer如果必须换版本更可控的方案是用TensorRT官方容器在容器里装指定版本engine的构建和推理都放在容器内完成。另一个选择是用NVIDIA SDK Manager重刷对应JetPack版本而不是单独替换系统里的几个库文件。另一个很常见的现场是把x86机器上构建好的engine拷到Orin上直接跑报错信息五花八门绝大多数原因是host端TensorRT版本比板端高engine里的kernel和板端不兼容。这种情况不用纠结版本差异在目标设备上重新构建一次engine比做任何兼容层都省事。换TensorRT版本后同一个ONNX构建出的engine大小和推理速度有变化是正常的不代表模型变了。新版本上重新构建并做一轮精度对比比猜测原因更高效。6. TensorRT部署后的验证平均FPS和ID切换定位技巧6.1 用平均FPS和ID切换次数验收部署完别只看画面“好像挺流畅”。两个最直接的指标是完整管线的平均FPS以及每100帧发生的ID切换次数。平均FPS用一段固定长度视频测import time start time.time() for i, frame in enumerate(video_frames): run_pipeline(frame) # 替换成实际主循环 if i 100: avg_fps 100 / (time.time() - start) print(favg fps: {avg_fps:.1f})用100帧的平均值而不是瞬时帧率能过滤掉前几帧初始化engine和CUDA context带来的抖动。FPS低于预期时先回看trtexec构建时输出的device memory和平均耗时确认engine本身没问题再回头查预处理和后处理。6.2 定位ID切换发生的位置ID切换是DeepSORT部署里最影响体验的问题但不需要靠肉眼整段盯视频。把每帧的track_id和框中心点存下来按id分组检查同一id的轨迹点时间戳是否连续from collections import defaultdict traj defaultdict(list) for frame_idx, tracks in enumerate(all_tracks): for x1, y1, x2, y2, track_id in tracks: traj[track_id].append((frame_idx, (x1 x2) / 2, (y1 y2) / 2)) for tid, pts in traj.items(): gaps [pts[i][0] - pts[i - 1][0] for i in range(1, len(pts))] if max(gaps, default0) 3: print(tid, max(gaps), 帧没有连续出现)轨迹点时间间隔超过3帧说明这个ID中断过大概率是在中间被另一个ID接管了。顺着打印出的frame_idx跳到视频对应帧检查那一帧的检测框是否被NMS滤掉、ReID特征是否因为裁剪尺寸不对而失真很快就能定位原因。拿到gap分布后再去调max_age或n_init依据就具体了。本文还有配套的精品资源点击获取
返回列表