ONNX 模型拆分完整技巧工具、场景、坑点、示例代码拆分 ONNX ≠ 单纯提速拆分是手段提速只是其中一种场景很多时候拆分反而会变慢。一、需要拆分 ONNX 模型的核心场景拆分模型不是单纯为了提速是解决各类工程痛点的手段分为性能优化类、工程兼容类、业务调试类三大场景一拆分后可提升推理速度的场景异构硬件流水线并行设备同时具备 GPU/CPU/NPU/DLA/DSP 多种算力单元将预处理、主干网络、NMS 后处理分配给不同硬件异步并行执行硬件资源不闲置整体推理耗时下降。规避算子不兼容减少 CPU 降级推理TensorRT、RKNN、Tengine 等端侧加速引擎对 Loop、If、复杂索引、动态 NMS 等算子支持差拆分后把不兼容算子单独切出交由 CPU 运行主干网络完整走硬件加速避免整张模型全部软推理。多任务多头按需推理减少无效计算一体化多任务模型检测 分割 分类每次推理会执行全部分支拆分后可按需只运行需要的子模型跳过无关分支降低计算量。超大模型分片解决显存 / 内存溢出BEV、大分割模型一次性加载权重会出现 OOM拆分后分段加载、推理后释放权重保证模型正常运行间接避免卡顿、崩溃带来的效率损失。二拆分不提速仅解决工程 / 调试问题的场景精度问题定位调试单独提取中间特征子图对比框架导出与 ONNX 中间输出快速定位误差层。业务模块化复用主干网络通用检测头、分割头可单独替换、迭代无需重新导出完整模型。分层交付与权限隔离模型主干、后处理分开交付区分不同模块使用权限。自定义算子改造单独切出目标算子所在子图单独替换为 CUDA / 硬件自定义算子。多工具链适配主干用 TensorRT 编译加速后处理用 ONNXRuntime 运行两类工具无法同时处理一张完整模型。三不建议拆分的场景拆分反而减速仅单 GPU 硬件、拆分后串行运行多个子模型会产生多重损耗多次创建推理会话、GPU/CPU 数据来回拷贝、算子融合优化碎片化、权重重复加载推理耗时上升 20%~50%。二、4 种主流拆分工具对比工具优势适用场景局限onnx.utils.extract_modelonnx 原生零额外依赖、极简 API简单直线图、无分支 / 残差残差跳跃连接、循环子图易报错onnx-graphsurgeonTensorRT 配套图编辑能力最强可改节点、重连线、处理残差复杂网络、残差、多分支、跨硬件拆分需额外安装仅 Pythononnxslim轻量化、支持 CLI 命令行、一键提取子图快速裁剪、删除冗余输出复杂图修改能力弱于 graphsurgeon框架侧预拆分PyTorch/Paddle 导出时分段无 ONNX 图修复、张量名天然对齐模型还未导出 onnx 阶段拆分已有的 onnx 文件无法使用三、基础拆分方法 1原生 onnx extract_model最简单原理指定子图输入张量名、子图输出张量名自动提取中间子图自动携带所需权重 initializerONNX。示例代码import onnx # 原始模型、输出子模型 src_onnx full_model.onnx sub_onnx backbone_sub.onnx # 1. 用Netron打开onnx找到分割边界张量名 # 例如原图输入是[images], 分割点张量为feat_out sub_inputs [images] sub_outputs [feat_out] # 提取子模型 onnx.utils.extract_model( src_onnx, sub_onnx, input_namessub_inputs, output_namessub_outputs ) # 校验子图合法性 model onnx.load(sub_onnx) onnx.checker.check_model(model) print(子模型拆分完成)适用限制不能切割If/Loop等带子图的控制流算子残差跳跃连接跨分割边界会丢失张量直接报错仅适合单向无分支的简单网络分类主干。四、基础拆分方法 2OnnxSlim命令行 Python快速裁剪CLI 一行拆分推荐快速调试# 只保留从images输入到feat_out输出的子图 onnxslim full.onnx sub.onnx --inputs images --outputs feat_out # 只删除多余输出保留原图输入 onnxslim full.onnx head.onnx --outputs det_out0,det_out1Python APIimport onnxslim onnxslim.slim(full.onnx, sub.onnx, inputs[images], outputs[feat_out])五、高级拆分方法 3onnx-graphsurgeon复杂网络首选处理残差 / 多分支核心优势手动遍历张量、重定向输入输出、清理孤立节点、完美兼容残差、跳跃连接、多分支网络解决 extract_model 残差报错问题。实战代码切分 Backbone 与检测头import onnx import onnx_graphsurgeon as gs # 1. 加载图 model onnx.load(yolov8_full.onnx) graph gs.import_onnx(model) tensors graph.tensors() # 2. 定义分割边界张量Netron查看 split_tensor tensors[backbone_feat] # 主干输出头分支输入 origin_input tensors[images] # 拆分1主干子图 # 新图输入原图输入新图输出分割张量 graph_backbone graph.copy() graph_backbone.inputs [origin_input] graph_backbone.outputs [split_tensor] # 清理无用节点、权重 graph_backbone.cleanup().toposort() onnx.save(gs.export_onnx(graph_backbone), backbone.onnx) # 拆分2检测头子图 graph_head graph.copy() # 头模型输入改为分割张量输出保留原图所有输出 graph_head.inputs [split_tensor] graph_head.cleanup().toposort() onnx.save(gs.export_onnx(graph_head), det_head.onnx)关键技巧处理跨分支残差拆分前先graph.cleanup()移除冗余 Identity所有跳跃连接张量不能跨分割边界分割线必须放在残差 Add 之后多输出场景将所有分支末端张量统一设为子图 output。六、框架侧预拆分导出前拆分最优方案在 PyTorch 导出 ONNX 前直接拆分子网络张量名天然对齐无图修复问题推荐新项目使用。import torch from torchvision import models model models.resnet50(pretrainedTrue).eval() dummy torch.randn(1,3,640,640) # 拆分1主干 backbone torch.nn.Sequential(*list(model.children())[:-2]) torch.onnx.export( backbone, dummy, backbone.onnx, input_names[img], output_names[feat], opset_version17 ) # 拆分2分类头 feat_dummy torch.randn(1,2048,20,20) head torch.nn.Sequential(model.avgpool, model.fc) torch.onnx.export( head, feat_dummy, cls_head.onnx, input_names[feat], output_names[pred], opset_version17 )七、工程通用拆分技巧1. 分割线选择黄金规则禁止切残差 / 跳跃连接中间分割点放在 Add/Concat 之后不要切 Shortcut 支路避开控制流算子内部If、Loop、Scan 的子图不能被分割线截断边界张量尽量选 Identity 输出Netron 中插入 Identity 节点作为分割标记方便定位多头模型统一在分支起点分割多检测头、多分割头从主干输出处一刀切。2. 拆分前预处理大幅降低报错简化模型onnxsim full.onnx sim.onnx消除冗余 Reshape、Identity推理 shape 推导onnx.shape_inference.infer_shapes(model)子图 shape 校验外部权重分离超大模型拆分前导出外部数据避免 onnx 文件过大from onnx.external_data_helper import convert_model_to_external_data model onnx.load(big.onnx) convert_model_to_external_data(model, locationweights.bin) onnx.save(model, big_split.onnx)3. 拆分后校验三板斧# 1. 格式合法性校验 onnx.checker.check_model(sub_model) # 2. 推理输出对齐原始图vs子图拼接输出 import onnxruntime as ort # 原图推理 ori_sess ort.InferenceSession(full.onnx) ori_out ori_sess.run(None, {images: rand_input}) # 子图串联推理 b_sess ort.InferenceSession(backbone.onnx) feat b_sess.run(None, {images: rand_input})[0] h_sess ort.InferenceSession(head.onnx) sub_out h_sess.run(None, {backbone_feat: feat}) # 输出误差校验 import numpy as np print(np.allclose(ori_out[0], sub_out[0], atol1e-5))4. 特殊场景拆分方案场景 A前后处理拆分CPU 预处理GPU 推理预处理Resize/Normalize/Transpose拆为独立子模型CPU 执行推理主干 GPU 执行后处理 NMS/ArgMax 单独拆分DLA/CPU 执行场景 B多输出多头拆分# extract_model同时提取多个输出拆出单头 onnx.utils.extract_model( full.onnx, seg_head.onnx, input_names[feat], output_names[seg_out] # 只保留分割输出丢弃检测输出 )场景 C算子不兼容拆分如 TensorRT 不支持循环Netron 定位不支持算子的输入输出张量分割线放在该算子前后拆出不兼容子图子图用 ONNXRuntime/CPU 自定义算子推理其余用 TensorRT。八、高频报错与解决方案extract_model 报错张量不存在原因分割线跨残差支路shortcut 张量未包含在子图 解决改用 onnx-graphsurgeon 复制完整图后裁剪或移动分割点到 Add 后。子图推理 shape 不匹配原因未做 shape 推理、动态维度冲突 解决拆分前执行infer_shapes固定输入 shape 或统一 dynamic 维度。拆分后权重丢失原因initializer 仅被分支使用裁剪时被清理 解决使用graph.copy()完整复制原图再裁剪不要直接修改原图。Loop/If 子图拆分失败官方限制分割线不能切割控制流子图需将整个 Loop 作为一个完整子图拆分。九、工具选型速记简单直线网络、快速调试 →onnx.utils.extract_model/onnxslimResNet、YOLO、U-Net 带残差 / 多分支 →onnx-graphsurgeon模型还未导出、原生 PyTorch/Paddle → 框架内分段导出最优命令行批量拆分脚本 → onnxslim CLI以上内容要是对您有用请给个赞和关注感谢您的支持。