ARTICLE DETAIL

资讯详情

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

ONNX 动态 Shape 全链路:导出、推理与跨平台部署避坑

ONNX 动态 Shape 全链路:导出、推理与跨平台部署避坑 模型在 PyTorch 里跑得好好的一转成 ONNX 就翻脸——喂进去一张 1080p 的图直接甩你一行Got: (1, 3, 1080, 1920), Expected: (1, 3, 640, 640)。这种场面我遇到过太多次了尤其是做车牌识别、人物抠图这类输入尺寸天生就千变万化的场景。ONNX 的动态输入和动态输出本质上就是给张量的每一个维度两个身份选项要么钉死成一个常数要么写成一个符号名让它在每次推理时跟着你喂进去的数据自己变。这篇东西不打算讲 ONNX 的基础语法那些文档里都有。我要聊的是从导出、验证、推理到跨平台部署这一整条链路上动态 shape 到底会在哪些地方给你使绊子以及每个绊子背后真正的成因。如果你正卡在同一个模型换张图就报错、转 TensorRT 后只能跑一个尺寸、量化完精度崩了这类问题上这篇的排查路径应该能省你不少时间。文中给的代码都是可以直接抄去改的我也会把参数为什么要这么填讲清楚。1. 动态维度在 ONNX 图里到底是怎么表达的1.1 一个维度只有三种状态别把它想复杂ONNX 里描述形状用的是TensorShapeProto里面是一串Dimension。每个Dimension只有三种可能的状态理解这三种状态后面所有的报错你都能自己定位。第一种是定值维dim_value被赋了一个具体的整数比如dim_value 1。这种维度在推理时不许变你喂进去的形状对不上就直接抛异常。第二种是符号维dim_param被赋了一个字符串比如dim_param batch或者dim_param height。这个字符串本身没有数值含义它只是一个占位标签。关键在于同一个标签在同一个图里多次出现时ONNX Runtime 会认为它们是同一个值。这一点非常重要也是很多诡异的 shape 报错的根源——你把输入高度标成了height而图内某个中间张量恰好也用了同一个标签ORT 就会尝试去满足这个约束结果可能跟你预期完全不一样。第三种是完全空既没有dim_value也没有dim_param。这表示这里有一个维度但大小未知。onnx.shape_inference在推不出来的时候就会留下这种空维度。提示不少人以为只要把输入标成符号维就万事大吉了其实图内部的中间张量如果被推成了dim_value动态性会在中间某一层被掐断。验证的时候不能只看输入输出要打印全部中间张量的形状。1.2dim_value和dim_param是互斥的重复设置会静默失效手写代码去修形状的时候最容易踩的一个坑是这样写的import onnx model onnx.load(model.onnx) dim model.graph.input[0].type.tensor_type.shape.dim[0] dim.dim_value 1 # 先设了定值 dim.dim_param batch # 又想改成符号维 onnx.save(model, patched.onnx)这段代码不会报错onnx.checker大概率也能过取决于版本但加载进 ORT 之后你会发现动态性根本没生效或者报一个Invalid dimension之类的错。原因是 protobuf 的 oneof 语义dim_value和dim_param属于同一个 oneof 字段后写的那个会覆盖前一个但如果原来的值是默认值 0有时候序列化出来又是另一回事。稳妥的写法永远是先彻底清空再设from onnx import TensorProto def set_dim(dim, nameNone, valueNone): # 清空 oneof 里的两个字段 dim.ClearField(dim_value) dim.ClearField(dim_param) if name is not None: dim.dim_param name elif value is not None: dim.dim_value value for inp in model.graph.input: dims inp.type.tensor_type.shape.dim if len(dims) 4: set_dim(dims[0], namebatch) set_dim(dims[2], nameheight) set_dim(dims[3], namewidth)ClearField是 protobuf 的标准 API对 oneof 字段用它是唯一可靠的做法。这个细节几乎没有文档会提但它坑过的人不在少数。1.3 稀疏维和稠密维的区别别搞混ONNX 的类型系统里TypeProto有一个 oneof可以是tensor_type也可以是sparse_tensor_type、sequence_type、map_type、optional_type。你在改形状时必须确认自己改的是tensor_type否则改半天没反应。写个带校验的工具函数会省事很多def is_tensor(vi): return vi.type.HasField(tensor_type) def get_shape(vi): if not is_tensor(vi): return None tt vi.type.tensor_type if not tt.HasField(shape): return None out [] for d in tt.shape.dim: if d.HasField(dim_value): out.append(d.dim_value) elif d.HasField(dim_param): out.append(d.dim_param) else: out.append(None) # 未知维 return out这里None表示未知维注意它和符号维不是一回事。未知维在图优化阶段可能被折叠成任何值而符号维有明确的命名ORT 的 free dimension override 只能作用于符号维。这个区别在后面讲 ORT 性能调优时还会用到。2. 导出阶段就把动态轴钉死从 dummy input 到符号传播2.1torch.onnx.export里的dynamic_axes到底改了什么东西很多人对dynamic_axes有一个误解以为它会让整个模型支持动态形状。它实际做的事情只有一件修改图输入和图输出的形状声明。真正让图内部算子的形状推导变成符号式的是 PyTorch 导出器在追踪过程中做的符号传播。一个典型写法是这样的import torch model.eval() dummy torch.randn(1, 3, 640, 640) torch.onnx.export( model, dummy, det.onnx, input_names[images], output_names[preds], dynamic_axes{ images: {0: batch, 2: height, 3: width}, preds: {0: batch}, }, opset_version13, do_constant_foldingTrue, )注意dynamic_axes的键是你在input_names/output_names里起的名字不是 PyTorch 里参数的变量名。名字写错了不会报错只是那一条动态声明被静默忽略——这是最隐蔽的坑之一。写完导出脚本的第一件事应该是立刻用 ORT 把输入形状打出来看。PyTorch 2.x 之后多了一条dynamoTrue的导出路径形状控制改成了dynamic_shapes语义上用Dim对象描述表达能力更强比如可以表达高度等于宽度这种约束。如果你的环境是较新的版本两条路都值得试一下老路径对某些控制流更宽容新路径对动态 shape 的支持更彻底。具体用哪个还是看你模型里有没有 Python 风格的if/for。2.2 名字对不上会出什么乱子如果模型forward返回的是一个 tuple而output_names只给了一个名字导出器会把剩下的输出自动命名成12、13这种数字名字。后面用 C 或 Java 去取输出时按名字拿就会拿到 null。反过来如果output_names给的数量比实际输出多导出会直接失败。一个稳妥的做法是在导出后立刻做一次一致性检查import onnxruntime as ort sess ort.InferenceSession(det.onnx, providers[CPUExecutionProvider]) for i in sess.get_inputs(): print(IN , i.name, i.shape, i.type) for o in sess.get_outputs(): print(OUT, o.name, o.shape, o.type)打印出来的shape里字符串就是符号维整数就是定值维。看到[batch, 3, height, width]才说明动态轴真的生效了如果打印出来是[1, 3, 640, 640]那就是没写对名字或者被后面的优化步骤折叠回去了。2.3 实测验证同一份 session 跑两种尺寸验证动态性最直接的办法就是用同一个 session 连续跑两个不同尺寸import numpy as np for hw in [(640, 640), (1080, 1920), (384, 1280)]: x np.random.rand(1, 3, *hw).astype(np.float32) y sess.run(None, {images: x}) print(hw, -, [t.shape for t in y])这里有个实操细节值得说每次换尺寸ORT 会重新做一次内存规划和 kernel 选择。第一次跑新尺寸会比较慢后面同尺寸就快了。所以测性能的时候一定要先 warmup 几轮再计时否则你会得到动态 shape 比静态慢十倍的错误结论。2.4 已经导出成静态的了还能不能救能救但要分清改声明和改计算图两件事。onnx.load之后直接改graph.input的dim_param改的是声明而图内部的Reshape、Resize、Squeeze这些算子的 shape 输入往往是常量它们不受声明影响会继续把你锁死在原来的尺寸上。所以补救流程通常是三步先改输入输出声明再跑onnx.shape_inference.infer_shapes看符号能传播到哪一层最后用onnx-simplifier做一次常量化简和冗余清理python -m onnxsim det_static.onnx det_dyn.onnx --dynamic-input-shape --overwrite-input-shape--dynamic-input-shape会把所有输入的批次维和空间维都设成动态。但如果 sim 之后你去跑一个大尺寸输入发现还是报错那说明图内部有写死的常量在卡着只能回到导出环节重新导一遍。这种情况下改导出脚本比事后修补省事得多。3. 动态输出比动态输入麻烦得多3.1 输出为什么会变长NMS、TopK 和变长解码输入动态是你告诉模型可以多大输出动态是模型告诉你它找到了多少个。后者通常来自三类算子。第一类是非极大值抑制NMS。检测模型输出的候选框数量本身是动态的NMS 之后保留多少个完全取决于图像内容。ONNX 里对应的是NonMaxSuppression算子它的输出是[num_selected, 3]第一维天然是符号维。第二类是动态 TopK。语音识别、检索类模型里常见选出前 k 个候选k 可以动态也可以固定但排序后的索引范围是动态的。第三类是控制流。ONNX 支持If和Loop两个算子能表达数据依赖的分支和循环。带Loop的图输出形状往往要到运行时才知道。3.2 动态输出最坑的地方下游代码拿不到形状静态模型里你可以这么写out sess.run(None, feeds)[0] boxes out[0, :, :4] # 形状写死换成动态输出就不行了因为第一维是变量。必须改成从运行时结果里读outs sess.run(None, feeds) preds outs[0] print(实际输出形状:, preds.shape) # 每次都可能不同C 和 Java 侧同理。Java 里result.get(0).getValue()拿到的是嵌套数组长度必须从数组本身读不能靠预先分配的固定 buffer。如果你在 Java 里看到ArrayIndexOutOfBoundsException或者结果被截断八成就是这里写死了长度。3.3 我更推荐的做法把 NMS 挪出 ONNX 图经过这么多年折腾对检测类模型我有一个比较明确的偏好尽量让 ONNX 图的输出形状固定把 NMS 放到图外做。YOLO 系列的 anchor-free 输出就是个很好的例子模型主干输出[1, N, 85]N 只跟输入分辨率有关跟图像内容无关形状是完全可预测的。NMS 用 numpy、Java 或者 C 写一遍几十行代码性能比图内实现还好因为你可以顺便做阈值裁剪和 top-k 截断。这样做的好处是连锁的ONNX 图变简单之后量化更容易、转 TensorRT 更容易、跨平台一致性更好Java 侧也不用处理变长数组。代价是你得自己写后处理但这段代码写一次就够了。如果确实必须在图内做 NMS那就得接受输出动态同时注意NonMaxSuppression在不同 opset 版本里的输入参数有差异center_point_box这个属性的默认值在 v11 和 v13 之间是不一致的跨版本转换时经常对不上。4. 推理阶段的实战坑打包策略、内存和预分配4.1 动态 batch 不等于无脑加大 batch批处理能提升吞吐但动态 batch 有一个绕不开的问题同一个 batch 里的样本必须 padding 到相同尺寸。如果一批里有 640x640 的图和 1920x1080 的图padding 到 1920x1080 之后前者有超过 60% 的像素是无效的。算力全浪费在零上了。我一般用三种策略按场景选策略适用场景优点代价batch1 动态尺寸实时单路视频流无 padding 浪费延迟最低吞吐低GPU 利用率不高定尺分桶离线批量处理padding 浪费可控吞吐高需要先统计尺寸分布动态尺寸 batch尺寸接近的批量任务灵活排序开销大实现复杂定尺分桶是我用得最多的。做法是先跑一遍数据集统计所有图片的长宽聚成 4 到 8 个桶比如 640x640、960x544、1280x736、1920x1088同一批只放同一个桶的图。这样每个桶里可以做静态 shapeONNX Runtime 的图优化和内存复用都能吃到实测吞吐比纯动态 shape 高出不少。4.2 动态 shape 会让 ORT 的内存池失效ORT 默认用一个 arena 分配器来管理内存同一个 session 反复跑相同形状时内存是复用的几乎不产生新的分配开销。但形状一变arena 里的块就对不上需要重新申请。极端情况下比如视频流里每一帧尺寸都不同你会看到内存占用持续抖动甚至因为碎片而增长。处理方式有两个。一是限制形状的种类数——比如把输入尺寸都对齐到 32 的倍数再送入这样尺寸组合会收敛很多。二是用IOBinding自己管理输入输出内存import onnxruntime as ort import numpy as np so ort.SessionOptions() so.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL so.enable_mem_pattern False # 形状多变时关掉内存模式 sess ort.InferenceSession(det.onnx, so, providers[CPUExecutionProvider]) io sess.io_binding() x np.zeros((1, 3, 640, 640), dtypenp.float32) io.bind_cpu_input(images, x) io.bind_output(preds) sess.run_with_iobinding(io) out io.get_outputs()[0].numpy()enable_mem_pattern在形状固定时开着能省内存形状多变时反而会拖慢首次推理值得根据场景测一测。4.3 用 free dimension override 把符号维焊死如果你导出的模型是动态的但实际部署时尺寸是固定的可以在 session 层面把符号维替换成常数。这样既能复用同一份动态模型又能拿到静态图的优化效果so ort.SessionOptions() so.add_free_dimension_override_by_name(height, 640) so.add_free_dimension_override_by_name(width, 640) sess ort.InferenceSession(det.onnx, so, providers[CPUExecutionProvider])这个 API 只对dim_param生效对未知维既没有 value 也没有 param无效——这也是我在第 1.3 节强调未知维和符号维要分清的原因。注意override 之后模型就只接受这一个尺寸了喂别的尺寸会报形状不匹配。适合模型动态、部署固定这一种场景别在需要真正多尺寸的场景里用。5. 跨平台部署时的动态 shape 处理5.1 转 TensorRTprofile 的三段式必须配合业务峰值trtexec构建带动态维度的 engine 时要指定minShapes/optShapes/maxShapestrtexec --onnxyolo.onnx --saveEngineyolo_fp16.engine --fp16 \ --minShapesimages:1x3x320x320 \ --optShapesimages:1x3x640x640 \ --maxShapesimages:1x3x1280x1280三个形状的含义要搞清楚optShapes是性能拐点TensorRT 会按它来选 kernel 和调优maxShapes决定显存占用的上限因为 TensorRT 是按最大形状预分配 workspace 的。很多人只填了 max 填得很大结果显存爆了其实是可以把 max 压到业务真实峰值上线的。另外engine 一旦构建完成形状范围就固定了。如果实际输入超出了 profile 范围TensorRT 不会自动缩放会直接报错。对于分辨率变化很大的场景我的做法是构建两三个 profile 或者两个 engine按输入尺寸分流。5.2 转 RKNN动态 shape 基本要走固化这条路把 ONNX 转到 RKNN 部署在边缘设备上时动态 shape 的支持非常有限尤其是 int8 量化路径。原因是 int8 静态量化需要在构建阶段用校准数据集跑一遍统计每一层激活值的分布来确定量化参数。如果输入尺寸每次都不同激活分布没法收敛量化误差会非常不可控。所以流程上要倒过来先确定部署尺寸把 ONNX 图改成固定 shape再做量化。如果业务上确实需要多种尺寸就在同一份权重上导出多个固定尺寸的 ONNX分别转成多个 RKNN 模型运行时按需选择。听起来笨但这是目前最稳的做法。校准集也有讲究不要随便抓几张图就用最好覆盖实际业务里的各种光照、天气、角度。我一般取 200 到 500 张覆盖典型场景太多也没用反而拉长构建时间。5.3 int8 量化与动态 shape 的组合坑onnxruntime.quantization提供了两条路。quantize_dynamic只对矩阵乘法类的算子生效MatMul、Gemm、Attention 这类 Transformer 结构卷积网络基本吃不到收益。CNN 检测模型要走quantize_static需要提供校准数据读取器。这两条路和动态 shape 的关系不太一样。动态量化对形状不敏感因为它是权重的离线量化激活值在运行时动态算 scale所以动态 shape 模型也能跑。但静态量化对形状敏感校准过程中的激活统计是跟输入尺寸绑定的尺寸一变量化 scale 就失配了。实践中我的经验是如果模型必须保持动态 shape就不要做静态 int8 量化如果要做 int8就先把形状固化下来。硬要两者兼得只能自己在图里插QuantizeLinear/DequantizeLinear节点做 QDQ 格式的量化控制粒度更细但工作量翻倍。5.4 Java 端 onnxruntime 处理动态尺寸的几个要点Java 生态里做车牌识别、OCR 检测这类活儿ONNX Runtime 的 Java API 是主流选择。几个和动态 shape 直接相关的点OrtEnvironment env OrtEnvironment.getEnvironment(); OrtSession.SessionOptions opts new OrtSession.SessionOptions(); opts.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.ALL_OPT); OrtSession session env.createSession(det.onnx, opts); int h 640, w 1280; long[] shape new long[]{1, 3, h, w}; FloatBuffer buf FloatBuffer.wrap(nchwArray); OnnxTensor input OnnxTensor.createTensor(env, buf, shape); MapString, OnnxTensor feeds new HashMap(); feeds.put(images, input); try (OrtSession.Result res session.run(feeds)) { float[][][][] out (float[][][][]) res.get(0).getValue(); // 注意out[0].length 是运行时才知道的不要写死 }第一个要点是shape数组的乘积必须和FloatBuffer的长度严格相等否则会抛异常而且异常信息往往不直观。第二个要点是输出数组的维度长度要动态读res.get(0).getValue()返回的是嵌套数组第一维长度就是实际输出数量。第三个要点是OrtSession是线程安全的但OnnxTensor不是多线程推理时要每个线程各自创建输入张量。至于车牌识别里常见的 OCR 检测网络它的输入高度通常固定比如 32 或 48宽度是动态的而且宽度必须是 32 的整数倍因为检测网络里有 5 次下采样宽度不整除会导致最后一层的特征图尺寸对不上报错位置还特别深很难定位。我的做法是在送进模型之前统一做一次对齐int alignedW (int) (Math.ceil(rawW / 32.0) * 32);这个 32 是从网络的下采样倍率来的换个网络可能要改成 8 或 16看结构定。6. 一套可以直接照着走的排查顺序遇到动态 shape 相关的问题我现在的排查顺序基本固定下来了从外到内一层层剥现象大概率原因排查动作换尺寸就报形状不匹配导出时动态轴没生效打印 session 输入 shape看是不是字符串输入动态了但中间层对不上Reshape/Resize 的 shape 输入是常量跑 shape_inference检查中间张量输出形状每次都不一样导致下游崩NMS/TopK 变长输出打出运行时 shape改下游取值逻辑ORT 内存持续增长形状种类太多arena 碎片关掉 mem_pattern或做尺寸分桶转 TensorRT 后只能跑一个尺寸profile 没配或配得不对检查 min/opt/max 三段量化后精度暴跌动态 shape 下做了静态量化固化形状重新校准具体操作上我会按这个顺序走一遍。第一步用 ORT 加载模型把输入输出的名字和形状全打出来确认符号维的存在。第二步用两个差异很大的尺寸各跑一次看是否都成功顺便看输出形状的变化。第三步如果失败用onnx.shape_inference.infer_shapes生成一个带形状信息的模型然后遍历所有value_info找那些被推成常量的中间张量那个位置就是动态性的断点。第四步回到导出脚本针对性地把那个断点上游的算子换掉或者包一层。第三步是很多人会跳过的一步但它其实是最高效的定位手段。写个小脚本遍历一遍比在导出脚本里反复试参数快得多import onnx m onnx.load(det.onnx) m onnx.shape_inference.infer_shapes(m) for vi in m.graph.value_info: dims [] for d in vi.type.tensor_type.shape.dim: if d.HasField(dim_param): dims.append(d.dim_param) elif d.HasField(dim_value): dims.append(d.dim_value) else: dims.append(?) # 打印那些本该动态却变成了常量的张量 if 0 in dims or 1 in dims: print(vi.name, dims)跑完之后你会看到一批张量其中像Reshape的第二个输入、Resize的sizes输入这些如果它们是常量张量就是断点的来源。常见的修法是把尺寸计算挪到图里用Shape/Gather/Concat动态算出来而不是在导出时用 Python 的常量算好。这个改法在 torch 侧就是把x.view(x.size(0), -1)这种写法改得对 shape 算子更友好或者干脆用torch.onnx.export的dynamoTrue路径让它自己处理。最后再分享一个我踩过好几次的经验导出完的模型一定要拿两三个差异极大的尺寸各跑一遍再做交付别只用导出时那个 dummy 尺寸验证。我就遇到过一次导出脚本里动态轴配得没问题用 640 验证也过结果线上第一个 1080p 的请求就挂了——原因是模型里有一个interpolate用了固定的scale_factor图内部的尺寸在某一层被常量折叠成了导出时的值。这种问题只有用不同的尺寸实测才能暴露出来看代码是看不出来的。现在我的习惯是把尺寸实测写进 CI模型一更新就自动跑一遍多尺寸校验省得线上再踩一次。
返回列表