ARTICLE DETAIL

资讯详情

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

OmDet开放词汇检测模型ONNX与TensorRT部署实战

OmDet开放词汇检测模型ONNX与TensorRT部署实战 做目标检测落地的朋友应该能感受到这两年开放词汇检测open-vocabulary detection的热度。OmDet 就是其中一个绕不开的名字它能在推理阶段用自然语言文本定义类别不需要预先锁定类别集合。但模型效果再好最终要跑在服务端 GPU 或者 Jetson 边缘设备上就绕不开 ONNX 和 TensorRT 这两道工序。我最近把一个基于 OmDet 的检测流程从 PyTorch 原始权重一路拉到 ONNX Runtime 和 TensorRT 上中间踩了不少坑也积累了一套相对稳定的操作流程这篇就完整记录一下这条部署链路给正在做 DETR 系模型落地的人一个参考。这篇文章不是纯理论也不是那种只给几个命令就完事的教程。我会从为什么选 ONNX/TensorRT 两步走讲起再拆解导出 ONNX 时最容易翻车的细节然后给出 TensorRT 的 FP16/INT8 构建与量化方法最后整理一份常见问题速查表。整个过程会穿插我实际测试的数据和踩过的坑你可以直接照着做也可以拿去做你自己的检测模型部署。1. 部署前的整体设计为什么选择 ONNX/TensorRT 两步走1.1 OmDet 结构拆解OmDet 本质上是一个开放词汇检测器模型结构上和 DETR 系列一脉相承视觉 backbone 提取特征Transformer encoder-decoder 做跨层交互再配合一个文本编码器把类别描述编码成 embedding最后在解码端计算目标框和文本 embedding 的匹配得分。结构上它和传统 YOLO 这类单阶段检测器最大的区别在于模型输出的不是固定类别 ID 的置信度而是一组“类别无关的 proposal 框”加上一组与文本语义空间的相似度得分。这样设计的好处非常直观推理时你只需替换文本编码器里的类别描述就能检测出全新的目标类别而不需要重新训练模型。但这也带来一个部署层面的麻烦。传统 YOLO 的类别数是锁死的输出维度固定导出成 ONNX 非常顺利。而 OmDet 这类开放词汇模型等于把模型拆成了视觉模块和文本模块两部分文本模块的输入是一个动态的句子序列输出类别向量的长度又会随类别数变化。如果直接整模型导出ONNX 里就会出现大量动态 shape 节点后续 TensorRT 构建 engine 时很容易报错。1.2 两条部署路线的选型思路我在这个项目里一开始其实犹豫过到底是用 LibTorch 直接加载 PyTorch 权重做 C 推理还是走 ONNX/TensorRT 这条链路。LibTorch 的好处是原生支持 PyTorch 模型不需要处理算子映射问题改造成本最低。但坏处也很明显推理性能完全取决于 PyTorch 的 eager 执行模式很多算子融合和 kernel 优化都做不了显存占用也偏高。如果只是内部验证或并发量不高的场景LibTorch 够用但一旦要面对线上流量或者要跑到 Jetson 这类边缘设备上性能完全不够看。ONNX 作为中间表示最大的价值在于它把模型从 PyTorch 的运行时里解放出来。ONNX Runtime 能做图优化和算子融合还能切换 CPU/GPU 后端非常适合做跨平台验证TensorRT 则是在 NVIDIA GPU 上做最终加速的杀手锏它会对整个计算图做层融合、kernel 自动调优、精度校准尤其是对 Transformer 这种算子密集的模型优化空间非常大。我的建议很明确不要一上来就追求 TensorRT先把 ONNX 这一层跑通确认导出后的模型输出和 PyTorch 原模型能对齐再去做 TensorRT。这两步虽然麻烦但每一步都能暴露出不同的问题比直接一步跳到 TensorRT 踩一堆莫名其妙的编译错误要可控得多。1.3 每个环节的验收标准为了不让转换过程变成“能跑就行”我在动手之前给自己定了几个硬指标转换链路PyTorch 是基线ONNX Runtime 要能跑通且输出张量和 PyTorch 原始输出的数值误差在可接受范围内FP32 下单元素绝对误差小于 1e-3。精度保持TensorRT FP16 相对 ONNX FP32 的平均精度损失不超过 2%INT8 不超过 5%。性能目标512x512 输入在 A10 显卡上TensorRT FP16 的端到端延迟要比 PyTorch CUDA FP32 快 2 倍以上。这些指标不复杂但在转换过程中非常重要。没有验收标准你很容易在一个已经“精度崩了”的 engine 上继续调参浪费时间。2. ONNX 导出最容易翻车的环节2.1 导出前的“模型瘦身”OmDet 导出 ONNX 时我做的第一件事不是写导出脚本而是先把模型拆开。前面说过OmDet 有一个独立的文本编码器正常推理时它负责把类别描述编码成语义向量。这个模块如果在 ONNX 图里保留会出现几个问题一是动态的文本序列长度会影响 batch 维度的推导二是文本编码器本身也是 Transformer 结构会显著增大 ONNX 文件体积拖慢 TensorRT 的构建时间。实际操作上我提前把线上要用的类别列表固定下来比如需要检测 person、car、dog、cat 这四类那我就在导出前先用文本编码器把这几个类别的文本描述编码成固定维度的向量保存成一个 .npy 文件。真正导出 ONNX 时模型只包含 visual backbone 加 Transformer encoder-decoder文本 embedding 作为外部输入传进去。这样整个 ONNX 推理时的输入就只有图像和一组固定的类别向量所有 shape 都是确定的后续 TensorRT 构建非常省心。这里有个细节类别顺序一定要固定好不能今天检测 person、car明天就换成 car、person。因为文本编码器每跑一次得到的向量都代表一个特定类别的语义顺序一旦变了模型输出的匹配得分对应的类别就会错位。我在项目里直接写死了一个 categories.json 文件类别列表、文本 prompt 模板、向量文件都从这一个文件读取从根上避免了这个坑。2.2 dynamic shape 与 opset 的正确打开方式OmDet 导出 ONNX 的脚本我贴一个核心版本大家可以参考。import torch import onnx from omdet_model import OmDetModel model OmDetModel.from_pretrained(omdet-turbo-r50).eval().cuda() model model.half() # 如果后续目标是用TensorRT FP16可直接用half导出 # 固定输入尺寸512x512 dummy_images torch.randn(1, 3, 512, 512).cuda().half() dummy_text_emb torch.randn(1, 4, 512).cuda().half() # 4个类别 torch.onnx.export( model, (dummy_images, dummy_text_emb), omdet.onnx, input_names[images, text_emb], output_names[boxes, scores], opset_version17, dynamic_axes{ images: {0: batch, 2: height, 3: width}, text_emb: {0: batch, 1: num_classes}, boxes: {0: batch, 1: num_queries}, scores: {0: batch, 1: num_queries}, }, do_constant_foldingTrue, )有几个参数值得单独说明。opset_version 我选了 17不是越高越好。opset 版本太高TensorRT 可能还没完全支持对应的算子太低一些新的算子如 aten::scaled_dot_product_attention 之类的映射会出问题。从我的经验看TensorRT 8.6 搭配 opset 17 是最稳的组合之一。dynamic_axes 这里图片的 height 和 width 我开了动态这样后续可以灵活调整分辨率做性能和精度的平衡。但你要有心理准备OmDet 这类 DETR 模型内部有 attention 模块height 和 width 一旦变成动态值ONNX 图里就会产生大量的动态 shape 节点TensorRT 构建时需要用 --minShapes、--optShapes、--maxShapes 分别指定。如果不想折腾第一版直接固定成 512x512 导出也是可以的跑通以后再考虑动态。2.3 后处理到底该放哪里关于 DETR 系模型的后处理网上说法很多有人把 top-k 挑选、阈值过滤、NMS 全部写成 ONNX 算子放进图里有人全部留在图外。我的建议非常明确第一版全部放图外。原因是DETR 系模型本身是基于集合预测set prediction训练的训练时用匈牙利匹配直接监督每个 query 的输出推理时通常每个 query 对应一个目标框不太会出现大量重复框。所以很多实现根本不需要 NMS或者说 NMS 的作用远没有 YOLO 里那么关键。放在图外用 Python 或 C 写一个简单的阈值过滤加 top-k 选择几十行代码就能搞定。等你验证完整个转换链路确认 ONNX 输出完全正确之后再考虑把后处理算子放进 TensorRT 图里。到时候可以用 TensorRT 自带的 EfficientNMS plugin或者自己写 plugin性能和可维护性再权衡。但这些都是后话第一版的核心目标是把模型本身跑通不要给自己增加同时排查模型问题和 plugin 问题的双重负担。3. TensorRT 加速FP16 与 INT8 的取舍3.1 环境准备与版本TensorRT 的环境坑非常多我在这里先说两个最容易踩的。服务器端我用的 TensorRT 8.6.1CUDA 11.8PyTorch 2.0。注意 TensorRT 的版本一定要和 CUDA 版本匹配否则 trtexec 或 Python API 构建 engine 时会直接报找不到 cudart 的错。如果你用的是 TensorRT 10.x对应的 CUDA 版本要求又不同网上最常见的坑就是用旧版 CUDA 容器去编译新版 TensorRT 的 plugin然后各种符号找不到。Jetson 这边更典型。Jetson AGX Orin 自带的 JetPack 5.1.1 默认自带 TensorRT 8.5.3但很多在服务器上用 8.6 构建好的 ONNX拿到 Orin 上构建 engine 会因为算子版本问题失败。一个很常见的做法是“降 TensorRT 版本”把 JetPack 里自带的 TensorRT 降级到 8.4.3或者在 Docker 里用 NGC 的 PyTorch 容器里面带了对应版本的 TensorRT。我自己测试下来Orin 上用 TensorRT 8.5 构建 OmDet 的 FP16 engine 偶尔会碰到 Deformable Attention 算子崩掉降到 8.4.3 反而稳定。提示构建 TensorRT engine 的机器和最终部署的机器显卡型号、显存、TensorRT 版本都一致是最好的。跨显卡构建的 engine 不一定能加载成功尤其是从服务器构建然后拷到 Jetson 上基本都会失败。3.2 构建 FP16 engineONNX 导出来没问题之后TensorRT 构建 FP16 engine 是最快的一步。命令行用 trtexec 就能完成。/usr/src/tensorrt/bin/trtexec \ --onnxomdet.onnx \ --saveEngineomdet_fp16.engine \ --fp16 \ --minShapesimages:1x3x512x512 \ --optShapesimages:1x3x512x512 \ --maxShapesimages:4x3x512x512 \ --minShapestext_emb:1x4x512 \ --optShapestext_emb:1x4x512 \ --maxShapestext_emb:4x4x512注意几点。如果 ONNX 里开了动态 shape这里必须同时指定 min、opt、max 三个 shape少一个都会报错。optShapes 是 TensorRT 做 kernel 自动调优用的参考尺寸建议用你线上最常用的输入尺寸这样推理时性能最优。FP16 这块大部分模型开 FP16 精度损失可以忽略不计但 OmDet 里的 LayerNorm 对低精度比较敏感。我的做法是构建完 engine 后用 polygraphy 检查每一层的输出精度。polygraphy run omdet.onnx \ --trt --fp16 \ --onnxrt --fp32 \ --atol 0.01 --rtol 0.01 \ --compare这条命令会把 ONNX Runtime 的 FP32 输出和 TensorRT 的 FP16 输出做逐层对比如果误差超过 0.01它会定位到具体是哪一层出错非常方便。3.3 INT8 量化与精度校准INT8 是 OmDet 真正吃性能提升的点但也是最容易把精度搞砸的环节。TensorRT 的 INT8 量化属于训练后量化PTQ需要校准集来统计每层激活值的分布然后确定合理的量化范围。校准集的选择是精髓。我一开始图省事直接用训练集的图片做校准结果在线上测试集上精度崩了 8%。后来发现原因是训练集里大部分图片都是干净的目标大图而线上场景有很多小目标、遮挡、光线变化大的图片导致校准时的激活值分布和线上推理时差异巨大。正确做法是用贴近线上真实分布的图片做校准数量不用太多500 到 1000 张就够但一定要多样化。另外 OmDet 这种有 Transformer 结构的模型直接全图 INT8 量化会有掉点风险建议用 polygraphy 做逐层敏感度分析把特别敏感的层通常是 LayerNorm 和最后的匹配头跳过量化维持 FP16。polygraphy run omdet.onnx \ --trt --int8 --calibration-cacheomdet_int8.cache \ --fp16 --atol 0.05 --rtol 0.05 \ --save-engineomdet_int8.engine校准缓存文件会记录每一层的量化尺度和范围构建 engine 时可以直接复用不需要每次重新校准。我一般习惯把校准缓存文件保留在代码仓库里方便复现。4. 实测结果与调优4.1 性能对比与精度对齐为了让大家有个直观的参考我把自己在一台 A10 上的实测数据整理成了表格。这个数据只代表我手上的环境和版本不同显卡、不同 TensorRT 版本会有差异但趋势是一致的。部署形态精度模式延迟512x512输入相对 PyTorch CUDA FP32 的加速比PyTorch CUDAFP3227.4 ms1.0xONNX Runtime GPUFP3218.9 ms1.45xONNX Runtime GPUFP1614.2 ms1.93xTensorRTFP168.6 ms3.19xTensorRTINT86.8 ms4.03x从数据上很明显能看到TensorRT FP16 相比 ONNX Runtime GPU FP32 将近 2.2 倍的优势这主要得益于 TensorRT 的层融合和 kernel 调优。INT8 相比 FP16 又快了约 20%但这个收益取决于模型里算子的类型占比我这里偏 Transformer 结构INT8 的提升没有纯 CNN 模型那么夸张。精度方面实际测试下来 FP16 模式下 mAP 相比 FP32 掉得很少基本在 0.5 个点以内视觉上几乎看不出来INT8 模式大约掉了 1.2 个点在可接受范围。如果对精度要求极其严格建议停在对齐效果更好的 FP16 模式INT8 只用在性能优先的边缘设备上。4.2 线上部署的显存与并发调优跑通 TensorRT engine 之后真正上线前还要考虑显存和并发。首先TensorRT engine 加载进显存时占了固定大小的显存不管输入 batch 是 1 还是 4这个基础占用量都是一样的。所以如果你的服务里同时加载了多个模型显存规划要提前算好。我在项目里用了一个简单粗暴的策略模型空闲 3 分钟无请求就卸载 engine释放显存有请求时再加载。因为 OmDet 的 engine 加载时间在 1 秒以内这个策略在实际运营中体验还不错。其次batch size 的选择要和线上流量匹配。TensorRT 引擎在构建时指定的 maxShapes 决定了它支持的最大 batch但实际推理时 batch 是动态的。如果你线上 qps 不高强制凑 batch 反而会增加首包延迟不如保持 batch1 简单应对如果 qps 高且聚合量大可以把多个请求攒到一个 batch 里做推理吞吐能提升不少。这个没有定论得拿真实流量压测以后再做决定。关于 GPU 资源测算我建议部署前先用 nvidia-smi 和 Nsight Compute 跑一轮 profiling看 engine 的显存占用、GPU util 和显存带宽。这个数据能帮你规划一台 GPU 能扛多少路并发不需要上线以后盲目扩容。5. 常见问题与排查技巧实录5.1 典型问题的速查表这一节我把项目过程中实际遇到的问题整理成表排查时可以快速对号入座。问题现象可能原因解决方案导出 ONNX 报 “Unsupported ONNX opset”opset 版本过高或过低改用 opset 17并在导出后检查 onnx.checker导出时报 NotImplementedError模型里有自定义 CUDA 算子如 Deformable Attention用 ONNX 友好算子重写模型结构或用 torch.onnx.is_in_onnx_export 分支处理TensorRT 构建报 “Assertion failed: inputs.size() 3”动态 shape 没配 min/opt/max检查 trtexec 参数或 Builder 配置中的 profileFP16 精度掉点明显LayerNorm 或 attention 层对低精度敏感用 polygraphy 逐层检查对敏感层设置 FP32 保留精度INT8 校准后模型精度崩了校准集分布和线上数据差异大换用贴近线上分布的校准集数量控制在 500-1000 张在 Jetson Orin 上构建 engine 崩掉自带 TensorRT 版本和模型算子不兼容将 TensorRT 降到 8.4.x或改用 NGC 容器构建engine 在另一台机器上加载失败engine 和 GPU/TensorRT 版本强绑定同一型号 GPU 上构建并部署不要跨显卡迁移推理结果类别错位文本 embedding 顺序和推理时类别列表不一致把类别列表、文本模板、embedding 向量统一从配置文件读取5.2 几个容易忽略的坑第一个坑是动态 shape 下的显存爆炸。我在 TensorRT 里把 height/width 开成动态后测试 1080p 输入时显存直接占用飙升比固定 512x512 高了将近 3 倍。原因是 TensorRT 为了支持动态 shape要为每个可能的尺寸分配 buffer 空间。解决方式是在 maxShapes 里设置你实际能接受的最大分辨率不要为了“灵活”而盲目调大。第二个坑是文本 embedding 的 batch 对齐。我的模型导出的文本 embedding 输入是 [batch, num_classes, dim]当 batch 为 1 时没问题但一旦 batch 变成 4文本 embedding 的 batch 维度必须和图像 batch 一致。如果文本类别相同直接 expand 就行如果每个 batch 的类别不同就要分别编码再拼起来。这个细节在 Python 里处理还好到了 C 端容易忽略建议在服务层就固定好文本。第三个坑是 LayerNorm 的 FP16 精度波动。OmDet 的解码器里有大量 LayerNormFP16 下它的均值和方差计算可能因为精度不足而产生微小偏移。这个问题在单张图片上不明显但在大量连续帧的视频推理场景里累计误差会导致检测结果偶发抖动。最稳妥的解决办法是在 TensorRT 构建时对 LayerNorm 所在子图设置 FP32 精度虽然会牺牲一点速度但换来的是稳定的效果。第四个坑是校准集的预处理必须和推理保持一致。我在做 INT8 量化时一开始校准图片没有做和线上一样的 resize 和归一化导致校准统计的激活值分布严重失真量化后精度惨不忍睹。校准集的预处理要和线上推理的预处理完全一致包括 padding、缩放方式、归一化参数一个都不能改。收尾的几句体会最后说点项目之外的实在话。OmDet 这类开放词汇模型从 PyTorch 到 ONNX 再到 TensorRT整个链路跑通以后收益是非常实在的。但实际上真正花时间的地方不是性能和精度调优而是“让模型能被标准工具链接受”。如果你的模型结构里有很多自定义算子比如 Deformable Attention建议一上来就用 ONNX 友好的方式去重写或替换不要等导出一半再回头补课。还有一个经验是如果你做 TensorRT 主要是为了跑 Jetson 或边缘设备建议提前在目标设备上验证模型结构不要在服务器上把整个链路调通后才发现边缘设备构建失败那会非常被动。先在目标设备上导出一个小模型做 smoke test确认环
返回列表