1. 为什么要专门聊SMOKE,以及它给部署挖了哪些坑
如果你做过一段时间自动驾驶或者边缘端3D感知,大概率会遇到这样的场景:算法组给出来的模型在PyTorch里跑得好好的,map还很好看,但推到你手上要做TensorRT加速时,各种问题就炸了。要么某些算子不支持,要么精度掉得没法看,要么动态batch调了半天还是性能拉胯。这时候你手头的模型如果是SMOKE,那恭喜你,你踩到的是“看起来简单、做起来处处是门道”的典型代表。
SMOKE是单目3D目标检测里非常有代表性的方案。它不依赖点云,只靠一张RGB图像去预测目标的3D框参数,包括中心点投影位置、深度、长宽高、航向角。它继承了CenterNet的角点检测思路,把3D检测做成了“关键点检测 + 属性回归”的单阶段任务。模型结构不复杂,训练代码也好找,但真正让它有工程挑战性的地方在于四件事:一是输出头冗余,二是反卷积和子像素卷积带来的算子兼容性,三是深度回归的尺度极其敏感,四是后处理里3D信息解算的细节容易被忽略。这四个问题,恰恰是TensorRT部署环节中最常见的拦路虎。
这篇文章我不会像教科书那样从“什么是深度学习”开始讲,而是以一个已经把SMOKE端到端在TensorRT上跑通的开发者视角,把整个链路拆开给你看。内容包括ONNX导出、TensorRT构建、精度调优、INT8量化、推理封装后处理,以及一张我在踩坑过程中整理出来的问题排查速查表。适合手里已经有一个SMOKE模型准备加速的工程师,也适合想理解“单个3D检测算法从训练到落地的完整过程”的算法同学。目标只有一个:让你照着这条路径走完,能拿到一个可以上实时推理链路的、稳定不掉点的SMOKE部署产物。
2. 算法层面先把账算清楚:SMOKE的输出到底在回归什么
2.1 单目3D检测的坐标体系和难度来源
部署模型之前,我不建议直接上手转换,先回到算法本身把输入输出挖明白。SMOKE的输入就是一张普通的三通道RGB图,分辨率常见的是384乘1280。注意它是单目的,不借助深度图,所以网络必须靠视觉特征去“猜”深度。这里有个绕不开的视觉常识:一张平面图像里,近处的车大、远处的车小,层叠遮挡又严重,算法要从这种2D表象中恢复出3D属性,注定是条件受限的。
SMOKE把问题定义成:每个物体的3D框中心投影到2D图像上,成为一个中心关键点。网络负责找出这些关键点的高斯热图,同时在关键点位置上回归三个东西:深度、3D尺寸(长宽高)、航向角(观察视角alpha)。关键点热图的输出通道数等于类别数,一般就是3到8类;深度输出是一个通道;尺寸输出三个通道;航向角输出两个通道,因为把角度通过正弦余弦编码成两个连续值。到部署阶段,这些都必须从引擎里原样取出来,不能经过任何不可导或者定制的中间层。
2.2 关键点热图、深度、尺寸、航向角的张量规划
张量规划这一步我看着很多人翻车,因为PyTorch训练代码里会把所有预测cat成一个大的输出张量,方便loss计算。SMOKE的常见做法是最后输出类似5到6个通道的头,具体排列是热图、深度、尺寸、转角。你把模型用ONNX导出之后,输出的结构就是你训练的head排列方式,这不代表TensorRT引擎也得这么编排。
我建议在导出前就把输出头拆散,不要让网络自己cat,或者至少在onnx里做成多个独立输出。原因有两个。第一,TensorRT对多输出的支持非常成熟,每个输出可以独立设置精度类型和显存优化,比如热图保持FP16、深度回归保持FP32精度就比全部混在一起更容易调优。第二,在后处理阶段,你不需要一次性把所有数据拷贝回CPU,可以单独把深度和航向角相关的张量留在GPU上做解码,减少无谓的拷贝开销。
2.3 部署里必须处理的归一化和坐标变换
这是SMOKE最容易掉点的根源之一。训练时图像的均值和方差归一化通常发生在网络输入前,数据加载器里做的。部署推理时,如果你直接用TensorRT,可别指望引擎里会内置一套归一化算子替你处理,绝大多数情况下你得在预处理阶段手动做。也就是先resize到384乘1280,转成NCHW布局,再按通道减均值、除方差。这个顺序不能反,前面做resize再归一化,和先归一化再resize,视觉上数值完全不同,映射到深度估计上差异更大。
后处理阶段还有一次坐标变换:网络输出的中心点位置是在输入图像坐标系下的像素坐标,不是原图坐标。你要根据resize的缩放倍数和可能存在的padding偏移把它映射回原图。这个映射如果在代码里写错了,出来的3D框投影到图像上看起来会整体平移几十个像素,但某些场景下检测精度指标却不会掉太多,极具欺骗性。所以项目落地时,我建议凡是涉及坐标变换的代码,必须在可视化工具里做一次“3D框投影回原图”校验,不要只看mAP数字。
3. 部署链路设计:PyTorch到TensorRT到底要走哪几步
3.1 整体路线拆解与每步交付物
我的落地路线是这样的:
- 第一步,PyTorch中拿到训练好的SMOKE权重,冻结batch norm,走一遍全流程得到基准输出。
- 第二步,剥离后处理,只保留网络主干到输出头,用torch.onnx.export导出ONNX文件。
- 第三步,用onnx-simplifier做算子化简,然后逐层对比PyTorch和ONNX的中间输出误差。
- 第四步,用TensorRT自带工具构建engine,先用FP32建一版,确保数值对齐后,再往FP16和INT8推进。
- 第五步,写推理封装层,负责预处理、拷入显存、推理、解析输出、后处理解码出3D框。
- 第六步,做性能测试和精度回归,重点盯深度误差和关键点召回率。
这条链路表面上看跟普通2D检测没有太大区别,但每一步都有细节坑。比如PyTorch里的某些算子,在ONNX导出后会变成多个基础算子组合,或者反卷积变成转置卷积的时候行为不一致。正因如此,我强烈建议在导ONNX那一步就做一个输出张量的逐值比对工具,别等TensorRT构建完了再来排错。
3.2 动态Shape与固定Shape怎么选
对于SMOKE这类基于CenterNet的单目检测器,我的建议是能在固定分辨率上跑的,就尽量不要用动态Shape。原因很直接:动态Shape在TensorRT里会对context做重配置,某些插件或者校准时算法会退化成保守路径,性能损失可能到30%到50%。这个损失在3D检测这种已经很吃算力的任务上,是你不愿意付的代价。
当然你不能完全拍脑袋固定一个分辨率,要结合你自己部署摄像头的实际分辨率和检测距离需求来定。比如你接的是1920乘1080的图像,图像中心区域是主要检测区,那可以原图直接缩放到384乘1280,纵向等比缩放、横向裁剪或者padding。注意这里的裁剪和padding策略,在后处理时要反算回原图的偏移。我个人更倾向固定Shape后配合一个包围盒Pad层,因为它能避免图像拉伸造成的变形,这对于关键点热图的中心定位精度至关重要。
3.3 构建引擎时的工作区、BatchSize与版本匹配
TensorRT构建时有一堆参数,但真正影响SMOKE部署效果的是三个:最大BatchSize、工作空间大小、与PyTorch版本匹配的算子兼容性。
最大BatchSize不是越大越好。我一直跟团队里同学强调,BatchSize设为2通常是最稳妥的起步值。因为推理时如果要用左右摄像头两路同时进引擎,那BatchSize至少得是2。但如果你BatchSize设到8甚至更大,构建engine过程中的优化策略会为了保留大Batch资源而牺牲部分计算融合,实时性反而可能下降。工作空间如果给得偏小,某些算子比如反卷积的融合就是不执行,推理耗时直接上一个台阶。保守起见给到默认的1GB到2GB即可,除非嵌入式设备显存受限才去往下压。
版本匹配这件事成也它败也它。TensorRT不同版本的ONNX解析器,对同一个节点的支持程度可以不一样。我曾经遇到过一个在旧版本里能顺利导入的ONNX,升级到新版本后某个of操作解析失败。所以别盲目跟风升级环境,锁定你自己的Conda虚拟环境、CUDA、TensorRT为一套组合,并且跑通后不要轻易动。落盘到具体项目上,这套组合甚至应该写进文档里随代码仓库一起维护。
4. ONNX导出环节的精细操作,决定了后续调优是否省力
4.1 从模型剥离后处理,导出纯网络结构
SMOKE的开源实现里,模型forward和后处理通常混在一起:模型输出热图后,会直接做一个3乘3的MaxPool拿到局部峰值,然后从峰值位置取深度、尺寸、朝向。这种风格对训练友好,对部署极其致命。TensorRT不会帮你完成这些Sampling坐标点、索引输出张量的操作,遇到这类节点多半直接报错。
根除方法只有一个——在继承原模型类的时候,重写forward函数,让它的返回值变成四个原始张量(热图、深度、尺寸、转角),其余后处理逻辑全部用NumPy或者PyTorch在外部实现。导出ONNX时用这个只含主干和任务头的模型来导出,后处理代码与部署端共享一套NumPy逻辑,避免训练代码和部署代码维护两套标准。
我在一个项目里帮一位算法同学排查过,他导出的ONNX包含了peak提取和topk操作,结果TensorRT构建阶段直接报不支持。我当时建议他改模型输出,把提取峰值当后处理做,总共只改了几行代码,问题就解决了。这不算技巧,是设计习惯问题,但从结果看,很多人就是会在这一步磨掉几天时间。
4.2 反卷积与像素重排的转换陷阱
SMOKE的Decoder部分,为了恢复高分辨率热图,会用到转置卷积。TensorRT对转置卷积的支持是有条件的,在FP16下偶尔会出现输出的数值变化比普通卷积明显得多。这跟转置卷积本质上是“插入空洞再卷积”有关,数值分布上容易出现较大的极值。
我的建议是,ONNX导出时优先检查一下转置卷积的groups参数和输出padding设置。SMOKE类模型一般用的是2倍或4倍上采样,输出padding通常是1,kernel size 3或4。如果这组参数在TensorRT里解析成错误的结果,最直接的排查办法是导出一小段只包含这个反卷积的ONNX,用trtexec单独构建,再对比PyTorch对应输出。用最小化复现的方式很快就能定位是通用解析问题还是模型本身结构问题。
另外一点是Pixel Shuffle这类重排算子,SMOKE不用,但如果你改版过网络插入了它,就要特别注意。TensorRT的底层能够支持pixel shuffle,但某些版本下生成的kernel效率不高,或者要求输入通道数必须被固定整除。这类算子的转换陷阱,我建议统一做法:能把上采样交给双线性插值解决的,就别用Pixel Shuffle;必须用的,做好性能benchmark,对比核验。
4.3 ONNX Simplifier与精度保守策略
很多工程朋友导出ONNX后喜欢直接上onnx-simplifier,然后发现精度对不上了,就怀疑是简化掉了什么关键逻辑。其实onnx-simplifier只是做常量折叠和等价替换,正常不会改变输出。但如果你模型里带有一些比较冷门的自定义函数,simplifier的pattern匹配可能误判,导致输出变化。
我的经验是,先在原始ONNX上做一轮量化校准和TensorRT构建,确认无误之后再做simplify,并且比对simplify前后所有输出张量的余弦相似度,阈值设到0.999以上。低于这个值,就不要用simplified版本。工程上最忌讳的是同时改多个变量,所以“先验证旧文件,再切换到新文件”对你排查精度问题是最高效的路径。
5. TensorRT构建与量化,性能曲线是怎么一点点挤出来的
5.1 FP32基准:先对齐数值再谈加速
我接管一个新模型时,基准做法永远是先用FP32构建一版引擎,跑一套与PyTorch相同输入的向量,对比输出张量的误差。FP32下数值误差应该非常小,大约在1e-4量级,如果这个量级对不上,那通常不是TensorRT的问题,而是ONNX导出环节就丢了信息。常见原因包括torch.onnx.export时没有设置opset_version足够新,或者模型里有原地操作被静默略过了。
我用的对比工具其实很简单,不需要上visdom或者tensorboard,就是NumPy算一下绝对误差和相对误差,然后按照输出通道分开打印。热图通道的误差允许稍微大一点,因为后处理要经过argmax这类抑制操作;但深度和航向角通道上的误差要严格控制,因为它们是连续回归值,直接决定最终检测框的空间位置。
5.2 FP16与INT8量化的取舍,以及校准那点事
FP16对SMOKE整体来说是收益明显的,运算吞吐基本能提升60%以上,但代价是深度通道的精度可能劣化。这类通道输出的是一个绝对值,通常跨度很大,近距离几米到远距离80米以上,FP16的尾数只有10位,在大数值上容易产生明显的量化步进。所以如果你发现FP16下深度误差明显变大,一个低成本的解决方案是:单独把深度输出改为FP32精度,其他输出保留FP16。TensorRT允许按张量设置精度,这个能力在3D检测场景下极为实用。
INT8量化则是另一个级别的挑战。SMOKE的INT8量化,校准数据选什么,几乎决定了最终mAP损失有多大。我建议校准集里不仅要有正常的城市道路场景,还要包含极端光照、雨雾、金灰色背景等边界样本,因为3D检测对深度的敏感度远高于2D分类,而深度估计对图像纹理细节的依赖非常强,纹理信息一旦被量化掉,深度预测立刻飘。实际操作时我习惯分两步做:先用100张左右校准图跑一轮,查看每一层的激活值分布;然后针对动态范围特别大的层单独设阈值,而不是依赖默认的最小化信息熵策略。
5.3 多Stream与DLA:你该关注的性能选项
如果你用的是带DLA的嵌入式计算平台,可以让关键点热图的计算走DLA,深度与尺寸头仍走GPU。这种异构分配在SMOKE这种多输出模型身上效果明显,因为热图分支通常是卷积链中最宽最深的,DLA可以消化掉这部分负载,而深度分支对数值敏感度较高,留在GPU上更稳。
不过DLA并不是万能的。TRT中启用DLA后,某些算子会不兼容,比如自定义插件多半要CPU辅助或者纯GPU执行,这时开销反而更高。我在实际项目里,只有在BatchSize固定、输入尺寸固定且不需要INT8的动态校准下才启用DLA。调试阶段关掉它,因为DLA的数值结果跟GPU有差异,你分不清到底是量化问题还是DLA计算差异。先把GPU上的结果调到数值可接受,再往DLA上迁移,是稳妥的策略。
6. 推理封装与后处理实现,距离一个能上车的引擎还差三步
6.1 显存管理与输入拷贝的双Buffer技巧
很多人把TensorRT推理封装理解为“load engine feed tensor get tensor”,太天真了。SMOKE的输入是图像,输出是多组张量。在同一条推理链路里,预处理后的图像要拷进GPU,这步涉及到CPU到GPU的PCIe传输。如果你一步一步串行处理:先等图像解码完成,再拷贝图像,再推理,那你CPU和GPU大量时间是空转的。
我会在工程上做一个双层缓冲,一个CPU buffer负责完成图像解码和归一化,另一个GPU buffer负责接收并执行推理。当第N帧图像在GPU推理时,CPU已经开始准备第N+1帧的预处理了。这么改动对单帧时延没有影响,但对吞吐量和流水线稳定性帮助极大。如果你是把推理节点包成一个ROS节点或者中间件服务,这种设计直接决定了你的节点CPU占用率能否保持平稳。
6.2 解码3D框:从热图到空间位置的完整脉络
后处理的流程看起来长,其实核心就四步。第一步,在热图上找峰值点,也就是局部最大值,并且通过边界剔除和阈值过滤留下目标类别。第二步,根据峰值点的坐标去索引深度、尺寸和朝向,得到该物体的原始属性向量。第三步,利用相机内参矩阵将2D中心点反投影到3D空间,结合深度恢复出物体中心的相机坐标。第四步,利用尺寸和朝向构造3D框的八个角点。
这里有一个我踩过的具体坑是深度解码公式。SMOKE的深度输出通常不是真值本身,而是通过一个先验均值加指数或者sigmoid变换得到。不同版本的代码可能细节不同,比如有的是depth = 1 / (sigmoid(depth_pred) * scale),有的是exp(depth_pred)再加上均值。你在写部署代码时必须严格对着训练代码里的decode函数,别自己凭记忆写一个差不多的公式进去。数值上差一点,3D框投影到正视图看起来还行,你从俯视图去看,框的纵向位置会整体偏移,车尾和车头的位置全错。
6.3 可视化验证:把3D框投影回2D图像里检查
我是强烈建议你把这套部署代码的可视化模块做完的,甚至在项目验收里把它列为一项交付物。因为3D检测部署全部调完,你只看mAP曲线有一些校准误差很难定位,但看可视化的鸟瞰图或者原图投影,很多问题一目了然。
我在真实项目里遇到过这样一个问题:某一辆处在图像边缘的车,给定一个深度的预测,网络认为它在3D空间里和另一辆车重叠了20%。这个结果如果在纯指标评估里,只会表现为一个小误差项,但如果把3D框画在俯视图里,几乎能看到物体框之间的物理重叠。通过可视化,我可以迅速追查是某个类别深度估计偏差大,还是整个摄像头的外参标定出了偏差。这个排查路径比单调对比日志高效得多。
7. 常见问题与排查技巧实录
很多问题单独出现时很难一次定位,我把实际项目里反复遇到的几类情况整理成一个速查表,方便你直接对照。
| 现象 | 直接原因 | 排查建议 |
|---|---|---|
| FP16引擎下深度整体偏大或偏小 | 深度输出头的数值范围太大,FP16尾数不足 | 单独给深度头指定FP32精度 |
| 模型转换后某些类别召回明显下降 | 热图输出被量化,峰值得分偏低 | 检查校准集是否覆盖多样性;对热图层单独设动态范围 |
| 输入对齐后检测框整体偏移 | resize与归一化顺序导致图像数值不对 | 重建预处理步骤,与PyTorch逐像素比对输入到网络的值 |
| 固定Shape下性能不佳 | 未绑定多个输入流或设定的工作空间不够 | 调整工作空间到2GB,确认实际batch与构建引擎一致 |
| ONNX里提示不支持某算子 | 后处理被混入模型图里 | 舍弃推理阶段的自定义节点,把该算子的对齐任务交给后处理 |
| 多路摄像头交替推理帧率高但卡发CTK | 反复调用context分配资源 | 推理循环启动前一次性分配context|上下文缓冲绑定,循环内不重复初始化 |
| INT8量化后深度误差集中在远距离 | 远距离样本在校准集里占比太低 | 增加远距离小目标样本,按距离分桶校准 |
| 原始模型有mish等激活函数 | TensorRT某些版本直接用GELU近似 | 尽量用ReLU或SiLU替换;或用插件实现精确版本 |
7.1 精度掉点时的排查顺序,应该先做什么
如果发现TensorRT引擎的精度显著低于PyTorch,请记住一个顺序:先查预处理差异,再查ONNX中间层输出,不要一上来就怀疑量化。很多所谓的“量化掉点”,实际上是你把图像的归一化顺序弄错了,或者是某个resize策略跟训练时不一致。先拿同一张测试图分别跑一遍PyTorch和TensorRT,把最终输出张量保存下来,计算每个通道的余弦相似度。定位到某个通道差异大,再看它前面对的输入层输出是否一致。逐层追下去,通常到不了第5层就能找到根本问题。
7.2 性能不达标的排查思路
性能问题比精度问题难查,因为影响因子太多。我一般会先看构建engine时的日志,把每一层的耗时打印出来,看看有没有诡异的超大buffer或者某个算子编译出了回退路径。然后量化到推理循环里,打印清楚每一帧中预处理耗时多少,GPU推理耗时多少,后处理耗时多少。我见过最典型的案例是后处理写得极烂,用了双层循环在CPU上挨个峰值点解码,这份耗时甚至超过GPU推理耗时。SMOKE的热图峰值点通常只有几十个到两三百个,完全可以把解码改成全向量化操作,一次NumPy矩阵运算完成。
7.3 动态Shape报错时的应急方案
如果业务上被迫要用动态Shape,而你的引擎构建过程报错,别慌。先缩小动态范围,比如把高度范围从384到384,宽度从1280到1920,而不是动不动就设置成可能从64到4096的大动态范围。动态范围越大,TensorRT构建时的候选实现越少,性能也就随之下降。SMOKE这类模型输出特征层级固定,输入的动态范围如果不能完全避免,就控制在一个接近实际使用的小区间内。同时要特别关注,预处理器统一padding时的偏移量填充值是不是会影响数理统计,有的填充值搭配不好直接改变多个关键点的响应强度。
8. 一点个人体会
我没有办法给你一个保证“照着做一定不掉点”的万能配方,但如果你把这条路走完一遍,尤其是把精度对齐和量化校准当成独立的验收环节去做,SMOKE在TensorRT上的提速效果是实打实的。我自己在最开始做这种单目3D检测部署时,也不太理解为什么要把反卷积的kernel配置记得那么死,直到一次FP16推理的实测把深度误差拉到十几米我才反应过来。后来我对团队里所有人的要求是:任何算法在部署前,先把训练时的decode代码看懂,并且保证有一份双视校验工具,没有校验工具就不开优化。
从精度对齐到性能榨干,是一个反复迭代的过程,但每一步的判断逻辑是相通的:不要同时改太多变量,任何一次改动都只影响一个因素,并且务必保留对比基准。希望这篇文能让你在接触SMOKE部署时少走一些弯路,少熬几个通宵。