CenterNet跨平台部署实战：从ONNX到TensorRT/RKNN的后处理与避坑

简介CenterNet 部署版资源包面向需要将目标检测模型移植到多种推理平台的开发者覆盖 ONNX、TensorRT、RKNN 以及地平线工具链解决模型转换与后端推理的适配问题。资源围绕 CenterNet 的中心点热图预测与后处理流程提供手写的后处理推理逻辑并附上作者在单目三维目标检测预研阶段整理的实现思路适合算法工程师或嵌入式部署人员直接参考。压缩包共 34 个文件主要以 Python 脚本、ONNX 模型、Shell 脚本、TensorRT 与 RKNN 推理文件、YAML 配置、README 说明和测试图片组成整体大小约 268.61MB。目录按不同平台划分脚本与模型一一对应可快速迁移到实际项目。资源已有两百人浏览学习下载后可同时获取各平台的转换与推理脚本、测试图片与输出结果对照图便于验证部署效果省去自行梳理部署流程的时间。1. CenterNet 部署版本真正要花时间的不是转换是后处理我拆这个 CenterNet 部署包的时候印象最深的不是模型本身而是作者那句“先把后处理手撸一遍”。很多做部署的人卡在这模型能从 PyTorch 转成 onnxTensorRT、rknn 也能跑通但一解码就翻车。这份资源把同一套 CenterNet 分别落了 onnx、TensorRT、RKNN 和 Horizon 四条移植路径每个平台都带推理脚本、测试图和输出样例。它最适合两类人一类想在边缘设备上跑 CenterNet 检测另一类是准备做单目 3D 目标检测、想先用 2D 检测把 heatmap 解码流程彻底弄明白的。后面我要讲的是后处理怎么从零写、导出与转换脚本怎么改以及不同平台下的取舍。2. 先搞懂 CenterNet 的三个输出头热图、偏移量、尺寸2.1 网络输出形态决定了部署代码怎么写CenterNet 不做 anchor 匹配也不用 RPN它把目标检测拆成三个并行的回归任务。以常见的 ResNet/DLA 主干为例输入 512×512 图像经过下采样和上采样恢复分辨率后最终输出尺寸是原图的 1/4也就是 128×128。三个输出头共用这一份特征各自负责一类信息heatmap 负责回答“这里有目标中心吗”offset 负责把中心点的整数坐标修正到小数精度size 负责输出目标的宽高。从部署角度看这三个输出头就是三块连续内存没有任何附加结构。不同平台、不同 SDK 对它们的处理方式都一样拿到矩阵做阈值过滤做坐标换算。搞清楚这一点后续移植才不会慌。输出头维度含义解码时怎么用heatmapC×128×128每个通道对应一个类别像素值为中心点概率未过 sigmoid 的 logitssigmoid 后取 topk 中心点offset2×128×128中心点亚像素偏移通道 0 为 x、通道 1 为 y加到整数中心坐标上再乘 stridesize2×128×128目标宽高通道 0 为宽、通道 1 为高乘回 stride 得到原图尺寸这里有个容易误解的点很多人在别处见过“CenterNet 输出 4 个值”的说法那是目标检测加旋转框或 3D 检测的变体。这个部署包面对的是 2D 检测场景只有这三组输出。单目 3D 检测如果要做通常是在 size head 旁边再扩展 depth、rotation 等预测分支但核心的 center 解码逻辑不变。所以先把这三组输出理清后面加分支只是多接几个“头”的事。2.2 手写 decode从 heatmap 到四个坐标的完整代码我把这个项目里的 decode 逻辑重写成一份不依赖 torch 后处理库、只用 numpy 的实现。它在 onnxruntime、TensorRT 和 RKNN 的推理结果上都能直接复用。import numpy as np from scipy.ndimage import maximum_filter def ctdet_decode(heat, offset, wh, topk100, stride4, threshold0.3): # heat: (num_classes, H, W), 未经过 sigmoid 的 logits # offset: (2, H, W), 通道 0 为 x 方向亚像素偏移, 通道 1 为 y 方向 # wh: (2, H, W), 通道 0 为宽度, 通道 1 为高度, 训练时按 stride 归一化 num_classes, H, W heat.shape # 热图先做 sigmoid, 再做 3x3 maxpool, 这就是 CenterNet 自带的 NMS heat 1.0 / (1.0 np.exp(-heat)) heat maximum_filter(heat, size3, modeconstant) # 把类别和空间位置合并成一个大列表, 取全局 topk, 这是官方 C 版实现的做法 scores_flat heat.reshape(-1) topk_indices np.argsort(scores_flat)[::-1][:topk] topk_scores scores_flat[topk_indices] cls_ids topk_indices // (H * W) spa_idx topk_indices % (H * W) ys spa_idx // W xs spa_idx % W offset offset.reshape(2, -1) wh wh.reshape(2, -1) # offset 和 wh 与类别无关, 只取对应空间位置 reg_x offset[0, spa_idx] reg_y offset[1, spa_idx] w wh[0, spa_idx] * stride h wh[1, spa_idx] * stride # 中心点坐标加上亚像素偏移, 再还原到原图尺度 cx (xs reg_x) * stride cy (ys reg_y) * stride x1 cx - w / 2 y1 cy - h / 2 x2 cx w / 2 y2 cy h / 2 keep topk_scores threshold boxes np.stack([x1[keep], y1[keep], x2[keep], y2[keep]], axis-1) scores topk_scores[keep] labels cls_ids[keep] return boxes, scores, labels这段代码要拆开看。heat进入函数时是 logits所以第一件事是 sigmoid 转成概率maximum_filter用 3×3 窗口做局部最大值抑制这一步替代了常规目标检测里的通用 NMS。取 topk 时我把所有类别的热图展开成一个一维数组再全局排序。这样做的好处是同一空间位置不会因为多个类别都响应而输出多个重复框官方实现也是这个思路。offset和wh的索引要特别注意它们的通道数固定为 2与类别数无关所以用spa_idx从展平后的矩阵里取对应的偏移和尺寸。wh在训练时通常除以 stride 做归一化解码时必须乘回stride4否则框会整体缩小四倍。threshold0.3是经验值实际工程里如果小目标漏检多可以降到 0.2但代价是假的中心点变多。2.3 为什么不建议直接复用官方 Python 后处理官方 CenterNet 仓库里的后处理在 PyTorch 环境下没有问题但它依赖torch.topk、torch.max_pool2d这些算子。部署到 ONNX Runtime、TensorRT 或 RKNN 后输入是三组纯内存数据没有 autograd也没有 tensor你必须自己处理 batch、通道、索引这些维度。很多移植项目死在半路不是因为模型转换失败而是因为把 torch 版后处理原样拷过去发现torch.topk在 RKNN 的 C API 里根本没有对应实现。注意这个项目里 onnx 版本只导出三个输出头不导出 decode。也就是说无论你最后跑在哪个平台后处理都要自己实现一份。这是刻意的设计避免把平台不支持的算子塞进模型图里。我一般会建议工程团队把这段 numpy decode 先放在 PC 端跑通用它对比 onnxruntime 的结果确认无误后再按平台语言重写一遍。这样每个平台的移植都只是“翻译”逻辑而不是重新设计逻辑。3. 导出 ONNX 与 onnxruntime 验证每一步都要确认输出形状3.1 torch.onnx.export 的参数取舍opset、动态 batch、输出命名这个资源包里已经有一份导好的centernet.onnx但如果你要从自己训练的权重重新导有几个参数值得按下面这样给。用 PyTorch 导出时最关键的三个决定是 opset 版本、是否开动态轴、输出命名。import torch # model 是训练完的 CenterNet, 这里需要切到 eval 并固定 batch 为 1 model.eval() dummy_input torch.randn(1, 3, 512, 512) torch.onnx.export( model, dummy_input, centernet.onnx, opset_version12, input_names[images], output_names[heatmap, offset, size], dynamic_axes{ images: {0: batch}, heatmap: {0: batch}, offset: {0: batch}, size: {0: batch}, }, )opset 我建议用 12而不是官方默认生成的更高版本。TensorRT 7 系列和部分 RKNN 工具链对高版本 opset 支持不完整opset 12 在兼容性和算子覆盖度上最稳。dynamic_axes把 batch 维设成动态这样同一份 onnx 既能跑 batch1 的调试也能在 TensorRT 里配置多 batch。三个输出必须同时声明动态 batch否则导出时会报维度不一致。热词里有人问“pytorch 转 onnx 之后怎么运行”这里顺手点破onnx 不是可执行文件它只是一张计算图。你要么用 onnxruntime 直接跑要么转成 TensorRT engine、RKNN 模型或 Horizon hbm。这个项目目录里给了onnxrun相关的 demo 脚本就是走 onnxruntime 验证的思路。3.2 onnxruntime 跑通预处理、输入输出、结果核对import cv2 import json import numpy as np import onnxruntime as ort mean np.array([0.408, 0.447, 0.470], dtypenp.float32) std np.array([0.289, 0.274, 0.278], dtypenp.float32) def preprocess(image, input_size512): # 保持长宽比缩放, 剩余部分用 0 填充 h, w image.shape[:2] scale min(input_size / h, input_size / w) nh, nw int(h * scale), int(w * scale) resized cv2.resize(image, (nw, nh)) canvas np.zeros((input_size, input_size, 3), dtypenp.float32) canvas[:nh, :nw] resized / 255.0 # onnxruntime 默认输入为 NCHW, 这里需要 BGR-RGB, 再减均值除方差 canvas canvas[:, :, ::-1].transpose(2, 0, 1) canvas (canvas - mean[:, None, None]) / std[:, None, None] return canvas.astype(np.float32)[None, ...] session ort.InferenceSession(centernet.onnx, providers[CPUExecutionProvider]) image cv2.imread(test.png) input_tensor preprocess(image) outputs session.run( [heatmap, offset, size], {images: input_tensor}, ) heatmap, offset, wh [o[0] for o in outputs] # 取 batch0 boxes, scores, labels ctdet_decode(heatmap, offset, wh) np.save(test_onnx_result.npy, boxes)预处理这一段是跨平台移植最容易出分歧的地方。/255.0之后做BGR-RGB再减mean、除std这是训练时的标准配置。mean和std的值并不是写死的要以训练代码里的transforms.Normalize为准。我见过有人把 mean/std 直接用 ImageNet 默认值导致 RKNN 上精度崩盘原因就是训练时用的是 CenterNet 官方这套统计值。session.run里的输出名必须和导出时的output_names完全一致如果导出时只写了[output1, output2]这里就要跟着改。跑完后如果test_onnx_result.jpg能画对框说明 onnx 图没有丢算子如果全图只有零散的点优先怀疑预处理其次是 heatmap 是否重复做了 sigmoid。4. TensorRT、RKNN 和 Horizon 移植实战平台差异全踩一遍4.1 TensorRT从 onnx 到 engine 的两种途径资源包里的onnx2trt_rt7.py是基于 TensorRT 7 API 的转换脚本用一个 python 循环遍历 onnx 节点转成 engine。这个脚本对 TRT 7 有效但如果是 TensorRT 10.x直接用官方trtexec更省事。热词里反复出现“tensorrt 版本如果是 10.x 是否支持 gtx1070”说明不少人在老 GPU 上栽过。GTX 1070 属于 Pascal 架构TRT 10 对它的支持不如 TRT 7 完整尤其是 FP16 和动态 shape 组合时直接 build 会报算子不支持。# TRT 10.x 下推荐用 trtexec 转换, 开 FP16 前先确认目标 GPU 架构 trtexec --onnxcenternet.onnx \ --saveEnginecenternet.trt \ --minShapesimages:1x3x512x512 \ --optShapesimages:1x3x512x512 \ --maxShapesimages:8x3x512x512 \ --fp16这个命令行里minShapes/optShapes/maxShapes三个参数构成了动态 batch 的区间。如果部署时只跑固定 batch1可以直接去掉这三个参数性能往往更好。--fp16在 Pascal 架构上收益有限而且容易让 heatmap 的峰值被削平我的经验是精度优先时先不开。如果用 python 来加载 engine核心调用长这样import tensorrt as trt logger trt.Logger(trt.Logger.WARNING) runtime trt.Runtime(logger) with open(centernet.trt, rb) as f: engine runtime.deserialize_cuda_engine(f.read()) context engine.create_execution_context() context.set_input_shape(images, (1, 3, 512, 512)) # 之后把输入拷到 GPU, 调用 context.execute_v2(bindings) # 三个输出的 binding 名字要和 onnx 导出时的 output_names 一致set_input_shape必须在你第一次执行推理前调用否则 TRT 会用默认的 profile 形状去分配内存动态 batch 下会报维度不匹配。绑定的显存 buffer 要用cuda.mem_alloc一次性分配输入和输出都是连续内存不能按通道分开拷。4.2 RKNN量化校准集和 dataset.txt 的坑RKNN 这条路我建议重点关注量化部分。onnx2rknn_demo_ZQ.py这个脚本名字里的 ZQ 指向瑞芯微的零拷贝方案它在 toolkit 2 上的一般流程是加载 onnx指定 target_platform设置量化数据集然后 build 和导出 rknn 文件。from rknn.api import RKNN rknn RKNN() rknn.config( mean_values[[0.408*255, 0.447*255, 0.470*255]], std_values[[0.289*255, 0.274*255, 0.278*255]], target_platformrk3588, quantized_dtypew8a8, ) rknn.load_onnx(modelcenternet.onnx, input_size_list[[3, 512, 512]]) rknn.build(do_quantizationTrue, datasetdataset.txt) rknn.export_rknn(centernet.rknn)这里mean_values和std_values单位是 0-255 尺度所以要乘 255跟 onnxruntime 里的归一化写法不一样。dataset.txt每行放一张图片的绝对路径至少准备 20 张覆盖不同场景的图。量化时工具链会从这里采样统计激活分布校准集如果全是测试集里的相似场景换到户外光照下小目标会成片消失。rknn 推理时输入 shape 在 toolkit 2 里统一为 NCHW不需要手动转 NHWC。拿到输出后decode 逻辑和 ONNX 完全一致唯一要确认的是输出张量的维度顺序尤其是 heatmap有些板端 SDK 会返回[C, H, W]有些会加一层 batch。4.3 Horizonmapper 目录和工具链边界这个包里centernet_horizon/mapper是地平线工具链的编译入口。地平线的路线和 RKNN 类似也是把 onnx 编译成自己的模型格式再在板端 runtime 加载。它的工具链版本迭代比 RKNN 更激进不同版本对 opset 的支持差别很大。opset 12 导出的这个模型在部分版本 mapper 上能直接过在另一些版本会报 Unsupported op。我的处理方式比较保守先看 mapper 目录里有没有别人留的配置没有就按地平线官方文档把 onnx 重新过一遍hb_mapper跑通后固定工具链版本不要再升级。Horizon 的板端后处理也是自己写 decode因为它同样不依赖 torch 算子。这里没有更多实际踩坑记录因为地平线工具链每个版本差异太大我不建议在没有板子的时候盲写。4.4 三个平台的关键差异对照平台转换入口动态 batch精度策略后处理注意点onnxruntime直接跑 centernet.onnx支持FP32预处理和 numpy decodeTensorRTtrtexec 或 onnx2trt 脚本支持需 profile默认 FP32FP16 可选输出 binding 顺序RKNNonnx2rknn_demo_ZQ.py一般固定w8a8 量化dataset.txt 校准集质量Horizonmapper 工具链固定量化工具链版本锁定这张表的结论是代码逻辑可以复用但每个平台都要重新过一遍预处理、输入 shape 和量化配置。任何一步对不上输出的框就是偏的、漏的甚至完全没有框。5. CenterNet 部署避坑记录五条反复出现的翻车现场5.1 topk 排序不一致导致框错位现象同一张 test.png在 PC 端 onnxruntime 解码出的框位置正确换成 C 重写后框全部向左上角偏移几个像素置信度还正常。原因我用np.argsort默认的快速排序而仿写的 C 代码用了std::partial_sort两者对相同 score 的排序稳定性不同。topk 索引一旦错位后面取 offset 和 wh 时取到的是另一个位置的值中心点和尺寸自然对不上。解决不要依赖排序稳定性。先按分数取 topk 索引再把索引存储为 long long 类型整个传递链路都保持整数索引如果 C 侧的partial_sort结果仍不稳定改成std::sort对 pairfloat, int 排序确保分数相同时按索引排序。从那以后我排序一律显式带上索引。5.2 TensorRT 10.x 在 Pascal 老 GPU 上 build 失败现象GTX 1070 上用 TensorRT 10.x 跑 trtexecFP16 开启后报Unsupported op: ConvFP32 能 build但速度比预期慢很多。原因Pascal 架构的 GPU 对 INT8/FP16 卷积的支持和 Turing、Ampere 不一样TRT 10 在高版本里把部分旧架构的 kernel 路径裁掉了。热词检索里也有人在问 10.x 是否仍支持 GTX 1070结论是能用但要避开 FP16 和动态 shape 组合。解决固定 TensorRT 7.2 或 8.2 版本重新转换如果版本不能换就放弃--fp16用 FP32 engine。CenterNet 这种单阶段检测器对 FP16 的精度损失本来就敏感heatmap 的峰值置信度容易被削掉所以牺牲一点速度换正确框是值得的。5.3 RKNN 量化后小目标全丢现象FP32 下检测正常w8a8 量化后大目标还在小目标几乎全漏尤其距离远的目标。原因量化校准集dataset.txt里 20 张图全部来自同一个园区场景光照和尺度单一。小目标在热图上峰值低量化后落到低比特区间直接低于 0.3 阈值。解决扩充校准集加入夜间、逆光、远距离目标图片数量增加到 40 张以上。如果还不行把quantized_dtype改成混合量化让 heatmap 输出保持 FP16只量化主干卷积。混合量化在 RKNN 里常见做法是do_quantizationFalse跑一遍再看哪些层对精度影响最大手动标记为不量化。5.4 预处理顺序不一致导致整张图输出乱框现象同一份 onnxonnxruntime 正常转成 rknn 后输出的框完全错乱甚至出现负坐标。原因rknn.config 里mean_values填的是 0-255 尺度而我 PC 端 onnxruntime 用的是归一化到 0-1 后减均值除方差。两边预处理不一致模型输入分布完全变了。解决rknn.config 里按mean[0.408*255, 0.447*255, 0.470*255]std[0.289*255, 0.274*255, 0.278*255]设置确保 BGR 到 RGB 的通道反转发生在 resize 之后而不是之前。这个坑出现频率最高所以我每次换平台第一件事就是打印预处理后的像素值范围而不是先看输出框。5.5 重复检测框和通用 NMS 的关系现象两个重叠目标类别相同decode 后输出了两个中心点但框严重重叠需要再做一次 NMS 过滤。原因3×3 maxpool 只能压掉同一目标内部的局部重复响应当两个同类目标靠得很近且 heatmap 峰值连成一片时topk 会同时选中两个相邻中心点。解决不要迷信“CenterNet 不需要 NMS”。它只是不需要常规的提案网络但后处理阶段加一个轻量 NMS 并不多余。常见做法是按类别做 IOU0.5 的抑制每类最多保留 100 个框。加了这层过滤后重复框数量会明显下降mAP 通常还能涨一点。6. 最后一个技巧用三平台输出 diff 验收移植正确性移植完成后最怕的不是没有框而是“看起来差不多但差一点”。我的习惯是固定一张 test.png让 onnxruntime、TensorRT、RKNN 三个平台各输出一份解码后的 boxes然后做最大绝对误差对比。这个验收步骤花不了十分钟但能挡住绝大多数移植事故。import numpy as np # ref: onnxruntime 解码后的 boxes # target: TensorRT 或 RKNN 解码后的 boxes # boxes 形状为 (N, 4), 列顺序为 x1, y1, x2, y2 def compare_boxes(ref, target, coord_atol1.0, score_atol0.05): if ref.shape[0] ! target.shape[0]: print(f框数量不一致: ref{ref.shape[0]}, target{target.shape[0]}) return False diff np.abs(ref - target) max_coord_diff diff[:, :4].max() max_score_diff diff[:, 4].max() if diff.shape[1] 4 else 0 print(f坐标最大误差: {max_coord_diff:.3f} px) print(f分数最大误差: {max_score_diff:.4f}) return max_coord_diff coord_atol and max_score_diff score_atol坐标误差小于 1 个像素、置信度误差小于 0.05我就认为这个平台移植接过了。如果误差在 2-3 像素之间先去查stride有没有乘错如果误差是整体偏移一个固定值查 offset 分支取值的坐标索引如果误差出现在后半段再考虑 topk 排序差异。这三个方向基本覆盖了我在 CenterNet 移植里遇到过的所有问题。从那以后我每次在项目里加一个新的部署平台都强制走一遍这个流程先跑 onnxruntime 解码再转换平台引擎再固定图片做 diff。这套三步走帮我挡掉过好多次“模型代码看着没问题”的假象也让我逐渐养成了后处理手写的习惯。CenterNet 这类中心点检测方法真正理解了解码逻辑换任何平台都只是把这套 numpy 翻译成对应语言的活儿。希望帮到你。本文还有配套的精品资源点击获取