使用 Polygraphysurgeon sanitize修改 ONNX 模型输入形状为部署打造动态 Batch 输入的实战指南【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT本文以 NVIDIA TensorRT 开源仓库中 Polygraphy 工具集的 CLI 示例为核心讲解如何在不改动模型中间层的前提下用一条polygraphy surgeon sanitize命令把 ONNX 模型的输入形状改写为带动态维度的形状如动态 Batch。读完本文你将掌握--override-input-shapes参数的完整语法、inspect model的验证方法以及该功能在 Polygraphy 源码中的实现原理与适用边界可直接应用于模型部署前的输入形状适配。为什么要修改 ONNX 模型的输入形状在深度学习模型部署中一个常见需求是模型在训练/导出时输入是固定形状例如[1, 3, 224, 224]但在线服务往往需要动态 Batch、可变分辨率等灵活性。理想的方案是在导出 ONNX 时就使用期望的形状——正如关联文档 03_modifying_input_shapes/README.md 强调的强烈建议以期望的形状重新导出 ONNX 模型。本文展示的方法只应在无法重新导出时使用。但在不少现实场景中模型由第三方提供、导出工具链不可用、只拿到成品 ONNX 文件重新导出并不可行。此时 Polygraphy 的surgeon sanitize子工具提供了一条轻量路径只改写模型输入的 shape 信息而不触碰任何中间计算层。这种改写的原理决定了它的边界它只是告诉模型新的输入形状并不会真正重写算子逻辑。因此如果模型内部对输入形状做了硬编码假设例如一个Reshape节点的 new shape 是写死的具体数值改写输入形状后模型可能无法正确工作。这一点将在后文注意事项与风险边界中展开。认识surgeon sanitize一把针对 ONNX 图的手术刀Polygraphy 的surgeon工具基于 ONNX-GraphSurgeon 对 ONNX 模型进行图级修改其总览见 surgeon/README.md。它提供多个子工具extract提取子图也可通过重新指定输入/输出来改变形状与数据类型sanitize清理、优化 ONNX 模型包括移除无用节点、常量折叠以及本文的主角——修改输入形状insert实验性向模型中插入节点prune实验性生成 2:4 结构化稀疏模型weight-strip/weight-reconstruct实验性权重剥离与重建。其中sanitize子工具在源码 sanitize.py 中的类注释将其定位为Clean up, optimize, and/or change input shapes in an ONNX model.即清理、优化和/或修改 ONNX 模型的输入形状。从Sanitize类的get_subscriptions_impl()可以看到它组合的能力模块模型参数支持输入形状覆盖、数据加载器、ONNX 形状推断、ONNX 加载/保存以及常量折叠参数组。这意味着sanitize在修改输入形状的同时还能顺带完成常量折叠--fold-constants、死节点清理--cleanup、拓扑排序--toposort等操作。核心参数--override-input-shapes语法详解修改输入形状对应的命令行参数在sanitize中名为--override-input-shapes其短别名是--override-inputs。该命名规则来自源码 model.py 中ModelArgs的input_shapes_opt_name机制传入override-inputs后工具会自动去掉末尾的s并追加-shapes生成第二个选项名因此两个写法等价。参数格式为--override-input-shapes name:shape [name:shape ...]nameONNX 模型中输入张量的名称shape目标形状支持固定数值、动态维度占位符与符号维度三种写法。形状中的维度写法规则parse_meta解析逻辑见 util.py 的parse_meta实现写法含义示例固定数值该维度保持固定大小[1,2,2]-1动态维度运行时可变化[-1,3,224,224]name符号维度用带引号的标识符命名动态维度[batch,1,2,2]多个输入可以一次指定多个name:shape对中间以空格分隔未列出的输入保持不变这一点有源码保证见下文源码原理部分。实战将固定形状模型改写为动态 Batch示例使用的模型是identity.onnx位于 03_modifying_input_shapes/identity.onnx。它本质上是一个恒等Identity网络输出等于输入不改变张量内容非常适合用来观察输入形状改写这一纯元数据操作。步骤一执行形状改写目标是把输入x的 batch 维度变为动态同时保持其余维度不变polygraphy surgeon sanitize identity.onnx \ --override-input-shapes x:[batch,1,2,2] \ -o dynamic_identity.onnx命令逐段解读identity.onnx输入模型路径sanitize的必需位置参数--override-input-shapes x:[batch,1,2,2]将输入x的形状改写为[batch, 1, 2, 2]其中batch是一个符号维度名称代表该维度为动态-o dynamic_identity.onnx保存改写后的模型到新文件不会原地覆盖原模型。执行后得到dynamic_identity.onnx其输入x带有符号动态维度batch中间层与输出层不做任何改动。文档明确指出输出形状可以被推断出来因此既不需要也不应该去修改它们——输出形状由 ONNX 形状推断自动重新推导。步骤二验证改写结果可选但推荐改写完成后用inspect model查看模型结构确认输入形状是否符合预期polygraphy inspect model dynamic_identity.onnx --show layers--show layers会展开显示模型中的每一层及其输入输出张量信息可以直观地看到x的新形状以及各中间/输出张量经过形状推断后的结果。若发现形状不合理可以调整--override-input-shapes的参数后重新执行。源码原理sanitize内部是如何改写输入形状的理解实现细节有助于判断什么时候能用、什么时候不能用。相关代码位于 sanitize.py 的run_impl_surgeon()其输入形状处理流程为解析用户输入元数据从ModelArgs读取--override-input-shapes解析出的user_input_metadata图处理若用户提供了形状覆盖调用override_input_shapes(graph, user_input_metadata)完成改写并标记rerun_shape_inference True重新形状推断改写后对模型重新执行 ONNX 形状推断使所有中间张量与输出张量获得正确的新 shape后续处理视参数执行常量折叠ConstFoldArgs.fold最后保存模型。关键的override_input_shapes函数实现在 tools/util.pydef override_input_shapes(graph, user_input_metadata): # 复用 extract_subgraph 逻辑前提是保留当前所有图输入 input_metadata onnx_util.meta_from_gs_tensors(graph.inputs) input_metadata.update(user_input_metadata) graph onnx_backend.extract_subgraph(graph, input_metadata) # 清空所有中间张量的形状避免旧的静态形状引发问题 for tensor in tensors.values(): if tensor not in graph.inputs and isinstance(tensor, gs.Variable): tensor.shape None return graph其中蕴含三个重要实现事实未指定的输入保持不变input_metadata.update(user_input_metadata)只覆盖用户显式给出的输入其他输入沿用原始形状中间层形状被清空所有非输入的中间张量shape被置为None这正是不改中间层的具体体现——旧形状信息被清除交由后续形状推断重新生成从而避免残留的静态形状与新的动态输入冲突重跑形状推断改写后的图由OnnxInferShapesArgs.infer_shapes重新推断sanitize默认开启形状推断可用--no-shape-inference关闭详见 loader.py。这也是输出形状可以推断、无需修改的由来。注意事项与风险边界使用此方法前必须明确其适用前提与限制模型内部不得对输入形状做硬编码假设。最典型的是带固定 new shape 的Reshape节点、依赖输入尺寸的ConstantOfShape、写死的Tile次数等。这类节点在输入形状改变后会产生错误结果甚至直接失效文档对此给出了明确警告。改写的只是输入张量的 shape 元数据不会也没有能力去修正这些依赖输入形状的内部算子。优先重新导出模型。文档明确强烈建议以期望形状重新导出 ONNX本方法只是无法重新导出时的兜底手段。常量折叠与形状改写可组合使用。sanitize同时支持--fold-constants常量折叠、--fold-size-threshold限制参与折叠的常量大小避免Tile/Expand等算子撑爆模型体积见 sanitize.py 的ConstFoldArgs、--cleanup死节点清理与--toposort拓扑排序在实际生产管线中可一并使用。延伸阅读从修改输入形状到完整形状工程surgeon系列示例围绕形状这一主题构成了完整的工作流值得一并参考01_isolating_subgraphs/README.mdsurgeon extract提取子图支持用auto让 Polygraphy 自动推断形状/数据类型也可用--inputs x1:[-1,-1]:auto的方式将固定形状改回动态02_folding_constants/README.mdsurgeon sanitize --fold-constants常量折叠可减少运行时计算量04_setting_upper_bounds/README.mdsurgeon sanitize --set-unbounded-dds-upper-bound为无界的数据相关形状DDS如Range、NonZero算子的输出设置上界避免 TensorRT 构建期优化困难甚至构建失败。在 TensorRT 部署管线中典型的组合拳是先用surgeon sanitize --override-input-shapes将输入改为动态形状如动态 Batch再用--fold-constants折叠常量必要时用--set-unbounded-dds-upper-bound收敛无界动态形状最后交给 TensorRT 构建引擎。这种先改形状、再优化、后构建的流程能让 ONNX 模型以最小的改动成本适配多样的在线推理负载。【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考