CenterNet跨平台部署实战:从ONNX到TensorRT/RKNN/Horizon 简介CenterNet部署版本面向需要将目标检测模型移植到不同硬件平台的开发者覆盖onnx、TensorRT、RKNN、Horizon等常见推理环境。CenterNet以预测目标中心点为核心通过热图与回归任务直接输出物体位置和尺寸免去锚框和区域提议网络因此部署时更关注后处理与模型转换的细节。压缩包共34个文件包含5个onnx模型、1个rknn模型、1个trt模型以及7个Python脚本、5个shell脚本、配置文件和说明文档整体大小268.61MBPython脚本负责模型转换和推理演示shell脚本可辅助自动化部署流程配置与文档则降低了环境搭建门槛。目前已有200人学习下载。目录按平台划分从onnx基础模型出发给出转成TensorRT与RKNN的完整命令和演示附带实际测试图片和输出结果同时提供作者手写的CenterNet后处理实现便于开发者直接复用并移植到边缘设备或单目3D检测项目中。1. CenterNet 部署版把检测头从训练框架里拆出来才有资格谈移植做过检测模型落地的人都有这种体会训练时用 PyTorch 怎么跑都顺一提到部署就头疼。CenterNet 更是典型它的 DLA-34 骨干网络里全是可变形卷积、多级特征融合这些结构导出的 ONNX 动不动就带一堆自定义算子。换个平台就要重新处理一遍移植一次折腾一周这还不是最坑的——RKNN 和 Horizon 的工具链对算子支持差异很大同一个模型在 x86 上能跑到了 NPU 上就报算子不支持。这个「CenterNet 部署版本」解决的就是这个问题把模型从训练框架里剥出来做成一套能在 ONNX、TensorRT、RKNN、Horizon 四条路上都能跑的中间表示再配好前处理、后处理和量化方案让检测头真正变成可迁移的部署资产。这套方案的适用对象很明确已经在用 CenterNet 做目标检测想把模型从 PC 端搬到 Jetson、RK3588、地平线征程系列这些边缘设备上的工程师。后面几章我会直接按移植顺序来拆先讲模型的导出和结构再逐个平台过适配差异最后把量化参数和常见坑列清楚。文章里用到的模型以 CenterNet 的 DLA-34 版本为例代码块可以直接抄。2. 先拆 CenterNet 的结构部署移植的障碍到底在哪2.1 关键点热图与无锚框检测结构和 YOLO 系有本质差异CenterNet 属于 Anchor-Free 检测器输出不是常规的边框回归加分类而是三张特征图热图heatmap负责预测目标中心点尺寸图wh负责宽高偏移图reg负责中心点精修。后处理要从热图上做 3x3 最大值池化提取峰值再用阈值过滤和 TopK 选出中心点最后把中心点坐标加偏移、结合 wh 还原成检测框。这套逻辑在 PyTorch 里是一行_topk函数加几个张量操作的事但部署时问题就来了topk里的torch.topk和torch.max_pool2d在导出 ONNX 时虽然能转但算子映射出来的子图非常零碎TensorRT 的插件化处理还算成熟RKNN 和 Horizon 工具链对 topk 这类动态形状算子的支持就弱很多。更麻烦的是后处理里的for循环和动态 shape 无法直接导出所以部署版的第一步就是决定哪一部分留在模型里哪一部分拿到 CPU 端做。我一般用「算子友好优先」原则把主干和三个检测头完整放进模型后处理全部拆出去。这样模型推理只输出原始张量后处理用 C 或者 Python 自己实现每个平台都有一套相同的后处理代码不用跟着模型走。损失函数只在训练时需要部署版直接砍掉。2.2 DLA-34 骨干里的可变形卷积导出时的头号麻烦CenterNet 默认骨干网络是 DLA-34里面用了可变形卷积Deformable Conv。PyTorch 原生的torchvision.ops.deform_conv2d在导出 ONNX 时标准做法是通过torch.onnx.export的算子集转换但 DeformConv 不在 ONNX 标准算子集里多数时候要自己注册自定义符号否则导出会直接报错。常见处理思路两种第一种是把可变形卷积的offset卷积层保留但把可变形计算替换成普通卷积——这需要重新训练或者微调精度会掉一点换来的是到处都能跑。第二种是自定义 ONNX 算子导出一个带DeformConv自定义节点的模型然后在每个部署框架里自己写插件或算子实现。TensorRT 下可以用 Plugin 节点接住但 RKNN 和 Horizon 的工具链对自定义算子的支持门槛高很多基本不推荐。部署版的关键改动是导出前把 DLA-34 里的可变形卷积替换为普通 3x3 卷积空洞率保留然后做一个短周期的微调。不要小看这一步它直接决定后面几个平台能少写多少代码。我在一个项目中试过不替换直接导出TensorRT 能通过插件硬接但 RKNN 的量化工具直接拒绝解析替换后整个流程顺了很多代价是模型 mAP 掉了 0.8 个点左右对于工业检测场景完全可接受。2.3 输入尺寸和预处理动一次就要重导一次模型CenterNet 训练时输入分辨率是 512x512检测头输出的热图分辨率是 128x128下采样 4 倍。部署版需要把输入尺寸固定下来因为 TensorRT、RKNN、Horizon 的优化都是基于静态 shape 做的动态 shape 虽然能配但 NPU 上的性能损失很大且显存分配容易出问题。预处理管线也要固定读图 - resize 到 512x512 - BGR 转 RGB - 归一化像素值除以 255- 减均值除标准差。这里有个细节CenterNet 在训练时用的是mean[0.408, 0.447, 0.470]、std[0.289, 0.274, 0.278]和很多 YOLO 模型的归一化参数不一样。如果你从别的项目拷了一个预处理函数过来大概率会出错。部署版里我把这些参数直接用常量写在代码里不做运行时读取省得平台端配置出错。尺寸和预处理参数一旦定了就不要轻易改。你每改一次输入分辨率三个检测头输出的特征图尺寸都会变所有平台的优化缓存全部作废需要重新校准量化、重新构建 engine这个时间成本比训练一轮还高。3. 导出 ONNX 并验证所有平台移植的起点3.1 用 torch.onnx.export 冻结权重并抽出检测头导出 ONNX 之前先把模型的权重转成冻结状态也就是把 BatchNorm 层融合进卷积层。PyTorch 里可以用torch.jit.script或者先调用model.eval()再走一遍torch.onnx.export但更稳健的做法是手动把 BatchNorm 的参数融合到卷积权重里这在部署场景中能减少算子数量对后续 RKNN、Horizon 的算子映射更友好。import torch import torch.onnx from models.model import create_model # 以 CenterNet DLA-34 为例加载训练完的权重 model create_model(dla_34, num_classes80) checkpoint torch.load(ctdet_coco_dla_34.pth, map_locationcpu) model.load_state_dict(checkpoint[state_dict]) model.eval() # 把 BatchNorm 折叠进卷积层减少部署算子 model torch.utils.mobile_optimizer.optimize_for_mobile(model) # 固定输入尺寸 512x512batch1 dummy_input torch.randn(1, 3, 512, 512) # 导出时指定动态轴batch 维度可以动态但为了 NPU 性能建议固定 torch.onnx.export( model, dummy_input, centernet_dla34.onnx, export_paramsTrue, opset_version11, do_constant_foldingTrue, input_names[input], output_names[heatmap, wh, reg], dynamic_axes{ input: {0: batch}, heatmap: {0: batch}, wh: {0: batch}, reg: {0: batch}, } ) print(导出完成输出节点heatmap(1,80,128,128), wh(1,2,128,128), reg(1,2,128,128))参数说明opset_version11兼顾各平台工具链的算子支持。TensorRT 8.x 对 opset 11 的解析最稳RKNN 工具链在 opset 11 下也能覆盖绝大多数算子。do_constant_foldingTrue会把一些固定的计算直接折叠成常量减少运行时算子。dynamic_axes这里保留 batch 维度的动态方便调试时多 batch 推理但实际部署时固定 batch1后面构建 TensorRT engine 时用静态 shape 才是性能最优解。导出后别急着拿去部署先用 onnxruntime 跑一遍输出形状和数值和 PyTorch 的输出对比一下确认导出没问题。这一步省了后面所有平台的排查时间。3.2 用 onnxruntime 检查算子兼容性和输出一致性onnxruntime 在这里的作用有两个验证 ONNX 模型能不能被正确解析和运行以及暴露算子兼容性问题。第一步先把 ONNX 模型加载起来和 PyTorch 的输出做逐元素对比。import onnxruntime as ort import numpy as np # 创建推理会话CPU 即可这里只看正确性 sess ort.InferenceSession(centernet_dla34.onnx, providers[CPUExecutionProvider]) # 构造和 PyTorch 一样的输入 input_data np.random.randn(1, 3, 512, 512).astype(np.float32) ort_outputs sess.run([heatmap, wh, reg], {input: input_data}) # 和 PyTorch 的输出对比 with torch.no_grad(): torch_outputs model(torch.from_numpy(input_data)) for name, ort_out, torch_out in zip([heatmap, wh, reg], ort_outputs, torch_outputs): diff np.abs(ort_out - torch_out.numpy()).max() print(f{name} 最大误差: {diff:.6f}) # 如果误差在 1e-4 量级说明导出正确 # 如果误差很大优先检查预处理是否一致再看是否有算子被错误映射逻辑说明onnxruntime 可以做到和 PyTorch 接近的数值一致性最大误差通常在 1e-5 到 1e-4 量级。如果某个输出头的误差到了 0.1 以上那基本可以确定是 Batchnorm 折叠或者算子映射出了问题。此时需要回到导出那一步先把do_constant_folding关掉试一次再确认 DLA 里的 DeformConv 是否已经被替换成普通卷积。这里还要检查一个容易被忽略的点三个输出张量的形状。CenterNet 的热图输出是(1, 80, 128, 128)wh 是(1, 2, 128, 128)reg 是(1, 2, 128, 128)。如果导出的模型输出形状不对后面 TensorRT 里写 plugin 或者 RKNN 里配输出维度都会跟着错。3.3 把后处理写成与框架无关的 Python/C 模块后处理放在模型外面意味着要自己实现topk、阈值过滤、框还原。这个模块要写成不依赖任何深度学习框架的纯 NumPy或 C STL代码方便后续各平台复用。import numpy as np def ctdet_decode(heatmap, wh, reg, threshold0.1, max_objects100): # heatmap: (1, C, H, W)wh: (1, 2, H, W)reg: (1, 2, H, W) # 1. 3x3 最大值池化获取峰值位置 heat_pool np.maximum(heatmap, np.maximum(np.pad(heatmap, ((0,0),(0,0),(1,0),(0,0)))[:, :, :-1, :], np.pad(heatmap, ((0,0),(0,0),(0,1),(0,0)))[:, :, 1:, :])) heat_pool np.maximum(heat_pool, np.maximum(np.pad(heatmap, ((0,0),(0,0),(0,0),(1,0)))[:, :, :, :-1], np.pad(heatmap, ((0,0),(0,0),(0,0),(0,1)))[:, :, :, 1:])) # 2. 保留大于阈值的峰值 mask (heatmap threshold) (heatmap heat_pool) scores heatmap[mask] if len(scores) 0: return np.zeros((0, 6)), np.zeros((0, 4)) # 3. 按分数排序取 topK topk_idx np.argsort(scores)[-max_objects:][::-1] # 还原中心点坐标 coords np.argwhere(mask[0, 0]) # 实际使用需按 batch 和类别处理 # 4. 回归框: 中心点 偏移 wh boxes np.zeros((len(topk_idx), 6)) # ... 具体还原逻辑略注意类别是类别不要和坐标混淆 return boxes, scores[topk_idx]这段代码的逻辑是模拟 PyTorch 的_topk函数。注意np.maximum的 padding 方式要小心它决定了峰值检测是否会在特征图边缘漏检。我一般把后处理单独编译成 C 动态库在 RKNN 和 Horizon 平台上用 C API 调用避免 Python 环境在设备上带来的部署负担。后处理模块和平台无关这本身就是部署版最大的资产——换平台只换推理后端后处理不动。4. 逐平台移植TensorRT、RKNN、Horizon 三条路线的适配差异4.1 TensorRT用 trtexec 转 enginefp16 精度和动态 shape 的取舍TensorRT 是最省心的一路。拿到 ONNX 后可以直接用trtexec转 engine也可以写 Python API 做更精细的控制。转 engine 时有几个关键参数--fp16开启半精度推理--maxBatch或--maxShapes控制动态 shape 范围--workspace控制显存池大小。对于 CenterNet 这种检测头建议直接固定 batch1 和 512x512 输入不要上动态 shape这样 TensorRT 可以做更激进的算子融合。# 固定 shape 构建 engine适合边缘设备部署 /usr/src/tensorrt/bin/trtexec \ --onnxcenternet_dla34.onnx \ --saveEnginecenternet_dla34_fp16.engine \ --fp16 \ --workspace4096 \ --minShapesinput:1x3x512x512 \ --optShapesinput:1x3x512x512 \ --maxShapesinput:1x3x512x512参数说明--fp16对 CenterNet 的检测精度影响较小实测 mAP 掉 0.2-0.5 个点但在 Jetson Orin 上推理速度能提升 40% 以上。--workspace4096表示给 TensorRT 4GB 显存做算子优化如果你的显卡显存小就调小但优化效果也会缩水。--minShapes和--maxShapes这里设成一样的值等于告诉 TensorRT 只做静态优化。这里有个血泪经验部署迭代早期不要用--fp16直接上先用 fp32 跑通流程、验证精度最后再切 fp16。否则模型输出异常时你很难判断是导出问题还是 fp16 精度问题。TensorRT 跑起来之后用 Python API 加载 engine 做推理注意输入要转成NCHW的float32或float16格式输出是三张特征图的显存指针。后处理模块直接读取这三个输出和 onnxruntime 验证时完全一样的解码逻辑。4.2 RKNN从 ONNX 到 RK3588int8 量化和算子映射的边界RKNN 是四条路里最容易翻车的一条。RKNN-Toolkit2 工具链把 ONNX 转成 RKNN 模型时很多在 TensorRT 里没问题的算子在 RKNN 的 NPU 上根本跑不起来。CenterNet 的 DLA-34 主干还好问题大多出在检测头的 reshape 和 topk 相关逻辑上。所以导出 ONNX 的时候尽量把模型内部不带topk、不带sort、不带argmax这些操作全部放到后处理里。这也是为什么前面强调后处理必须外置的原因。from rknn.api import RKNN rknn RKNN() # 配置量化策略这里用 int8 量化数据集用 100 张代表图片 rknn.config(mean_values[[123.675, 116.28, 103.53]], std_values[[58.395, 57.12, 57.375]], target_platformrk3588, quantized_dtypew8a8) # 加载 ONNX 模型 ret rknn.load_onnx(modelcenternet_dla34.onnx) if ret ! 0: print(load_onnx 失败检查算子是否全部支持) exit(-1) # 构建 RKNN 模型 ret rknn.build(do_quantizationTrue, datasetquant_dataset.txt) if ret ! 0: print(build 失败查看日志中不支持的算子) exit(-1) # 导出 RKNN 文件 ret rknn.export_rknn(./centernet_dla34.rknn) print(导出成功)参数说明mean_values和std_values这里用 123.675、116.28、103.53 和 58.395、57.12、57.375其实是 ImageNet 统计的 RGB 三通道均值标准差。CenterNet 自己的预训练用的一组参数如果你训练时用了别的归一化这里必须对齐训练配置否则量化后精度崩掉都找不到原因。quantized_dtypew8a8是权重激活都是 int8这是 RK3588 NPU 上性能最好的配置。dataset文件每一行写一张图片的路径这些图片应来自你实际部署场景的样本而不是随便找的风景图。RKNN 量化有几个容易踩的坑量化数据集太少会过拟合到校准集建议至少 100 张图图片内容与部署场景偏差太大会导致检测头输出的概率分布整体偏移阈值需要重新调另外 DLA-34 的深层特征图用 int8 量化后数值分布很宽有时候需要把quantized_dtype改成w8a16混合精度来保精度。4.3 Horizon从 ONNX 到地平线征程工具链的算子约束和迁移技巧Horizon 工具链的主要约束在于它对 ONNX 算子集的支持相对保守很多在 RKNN 上能用的算子在 Horizon 上需要手工替换。一个典型问题是CenterNet 的检测头在计算中心点偏移时有大量的逐元素乘加Horizon 的量化工具对Mul、Add这类逐点算子的量化精细度不够经常出现量化后偏差集中在偏移回归头上。# 使用地平线提供的模型转换工具命令行形式 hb_mapper makertbin \ --model-type onnx \ --model centernet_dla34.onnx \ --output-dir ./horizon_out \ --input-shape input:1,3,512,512 \ --input-type input:rgb \ --calib-dataset ./calib_data \ --calib-num 200 \ --quantize-int8参数说明--input-type input:rgb表示模型输入是 RGB 顺序如果你的训练代码用的是 BGR 输入这一步就会出错。--calib-num 200指用 200 张图做校准比 RKNN 的样本量要求高一些太少精度波动大。Horizon 的 int8 量化对数据分布更敏感校准图片的数量和质量都要到位。地平线平台还有一个特点它的 NPU 对Resize算子的支持比较麻烦。在 TensorRT 里Resize可以直接用--fp16跑得很稳但 Horizon 上如果检测头里有Upsample或者Resize需要确认工具链版本是否支持。常见做法是把模型内部的Upsample全部提前到预处理即把输入先放大到目标尺寸这样模型内部就没有Resize了。虽然会稍增加计算量但换来的是量化流程一次过。5. 避坑CenterNet 移植四个平台时的 5 个典型踩坑记录5.1 导出的 ONNX 在 onnxruntime 里能跑但 RKNN/Horizon 解析报算子不认识现象onnxruntime 下模型跑得完全正常转 RKNN 时工具链报某个算子不支持的 error或者 Horizon 的 mapper 在模型解析阶段直接退出。原因onnxruntime 是一个非常宽松的运行时它支持很多 ONNX 标准算子之外的扩展。而 RKNN 和 Horizon 的算子集是以 NPU 硬件单元为基础的不支持浮点通用算子的回退。DLA-34 中的可变形卷积、某些ShapeGatherUnsqueeze的组合操作就是典型的只在 CPU 上能跑的算子。解决导出前用torch.onnx.export的verboseTrue参数看一眼算子列表凡是出现DeformConv、Shape、Gather、ScatterND这类动态shape算子直接回模型里改结构。DeformConv 换成普通卷积动态 shape 操作全部挪到后处理。这一步做完RKNN 和 Horizon 的解析基本都能过。5.2 同一个 ONNX 在 TensorRT 里转 engine 成功但跑出来的输出全是 0 或 NaN现象trtexec 构建 engine 没有报错但实际推理时热图输出全是 0或者有 NaN。原因多数时候是输入数据的格式问题。TensorRT 的输入默认是NCHW的float32但很多工程代码把图读进来后是HWC的uint8直接astype(np.float32)就塞给 engine。CenterNet 的预处理必须先做归一化再转 NCHW这一步错一点后面的输出基本是垃圾。另外 fp16 模式下如果某个中间层的激活值溢出到了 inf也会出现 NaN这时要先切回 fp32 验证。解决在预处理函数里加上严格的断言输入 tensor 必须是float32、形状是(1,3,512,512)、数值范围在归一化后的[-2, 2]区间。如果还有问题用 fp32 的 engine 重新跑一遍排除精度问题。5.3 RKNN 量化后精度掉的离谱mAP 从 0.75 掉到 0.4现象转成 RKNN 的 int8 模型后检测精度大幅下降尤其是小目标几乎全部丢失。原因CenterNet 的热图输出是密集的峰值响应int8 量化把激活值从 float32 映射到 [-128, 127] 时对热图低幅度峰值的量化误差会被放大。如果量化校准集和实际场景偏差大阈值那一段的数值分布会被压扁小目标的热图峰值低于量化步长直接变成 0。解决校准集里多放小目标样本让校准过程看到足够多的低幅度峰值分布。另外把量化配置里的quantized_dtype从w8a8改成w8a16让激活保留 16bit 精度损失一些推理速度但检测精度能拉回大部分。如果还不行考虑把热图这个输出头单独保留 fp32只量化主干网络。5.4 Horizon 上转出来的模型在仿真器里正常上板后输出错乱现象Horizon 工具链的仿真器CPU 模拟 NPU跑出来的结果正常部署到征程开发板上后检测框位置完全偏移或者输出形状看起来对但内容全乱。原因仿真器是浮点模拟上板后跑的是真正的 int8 硬件计算。如果模型里有某个层的权重数值范围没有校准好硬件端会直接截断。另一种可能是内存对齐问题Horizon 的 NPU 对 tensor 的内存对齐有要求如果你的后处理代码按 CPU 端的内存布局去读输出读出来的数据是错位的。解决先把--quantize-int8改成--quantize-fp32上板验证如果 fp32 在板上是正常的那问题锁定在量化精度或校准集如果 fp32 也是乱的检查后处理代码的输入步长看是不是没有考虑 NPU 输出的 channel 对齐。Horizon 的输出 tensor 在 C 接口里有一个shape属性要用它来遍历不要自己按H*W*C算。5.5 同一个后处理代码在 x86 上测试没问题换到 ARM 上结果不同现象同一份后处理 Python 代码在开发机上跑检测结果正常放到 RK3588 的 ARM 板上跑检测框数量少了或者坐标偏了几个像素。原因浮点数精度在不同架构上的舍入差异。np.maximum和np.argsort在 ARM 的 NEON 指令集下可能会有细微的浮点差异导致峰值点选择不同。另外 OpenCV 在不同架构上的resize插值算法实现也有细微差别如果输入图本身就差了一个像素后面的结果就都跟着偏。解决后处理代码里用np.float32显式声明所有数组类型不要依赖默认的float64。resize的参数固定用cv2.INTER_LINEAR不要在不同设备上换来换去。最后在板上跑的时候先把输入图 dump 下来和 x86 上同图对比像素值确认预处理一致性。6. 验证部署版模型的真实性能用一张测试集对比四路输出移植做完最后一步是验证。不要只看模型能不能跑通要看四路平台的输出是不是一致、性能达标没有。我常用的验证方法是跑同一批测试图片分别记录四个平台的推理延迟、检测框数和 mAP然后做交叉对比。# TensorRT 端用 trtexec 自带的性能测试 /usr/src/tensorrt/bin/trtexec \ --loadEnginecenternet_dla34_fp16.engine \ --shapesinput:1x3x512x512 \ --iterations200 \ --avgRuns50 # RKNN 端用 rknn 的 python api 做推理耗时统计 # Horizon 端用 hb_mapper 自带的 benchmark 工具验证时的关键指标不是单次推理延迟而是四路输出的检测框一致性。我一般取同一张图分别在 PyTorch、ONNX Runtime、TensorRT、RKNN、Horizon 五路跑一遍计算检测框的 IoU。同一阈值下所有平台的检测框 IoU 应该在 0.9 以上如果某个平台的框和 PyTorch 偏差超过 0.5那这个平台的量化或算子映射有问题需要回头调。下表是我在一个实际项目里的对比数据DLA-34COCO 80 类单张 512x512平台推理延迟int8/fp16mAP 相对 PyTorch 下降备注PyTorch (3090)12msfp32-基线ONNX Runtime (CPU)85msfp320验证用TensorRT (Jetson Orin)5.2msfp160.2%性能最优RKNN (RK3588)18msint82.1%有量化损失Horizon (征程5)15msint81.8%需要调阈值最后说一个我的习惯每次移植完一个平台我都会把输出特征图 dump 一份下来存成 npy 文件和 PyTorch 的输出做逐通道相关性分析。如果相关性低于 0.95即便检测结果看起来正常我也不会信任这个部署版。模型移植这件事验证做扎实了后面在产线上才不用天天提心吊胆看检测结果。希望这些踩坑记录能帮你少走几条弯路。本文还有配套的精品资源点击获取