如何拆分ONNX格式DNN模型以实现边云协同分段推理

第一步：理清节点结构与拆分点 #

不需要依赖第三方工具做节点梳理，直接用ONNX原生API遍历所有节点即可，你可以打印每个节点的类型、名称、输入输出张量名，快速定位你要拆分的位置：

import onnx

# 加载你导出的模型
onnx_model = onnx.load("mobilenet_v2.onnx")
graph = onnx_model.graph

# 遍历所有节点，打印关键信息
for idx, node in enumerate(graph.node):
    print(f"节点序号: {idx}, OP类型: {node.op_type}, 节点名称: {node.name}")
    print(f"输入张量: {node.input}")
    print(f"输出张量: {node.output}")
    print("-"*50)

运行之后记录你要拆分位置的输出张量名，比如你要把前N层放边缘端，就记录第N层节点的output张量名，假设这里记录为splitter_tensor。

第二步：用原生API拆分ONNX模型 #

ONNX自带onnx.utils.extract_model接口，直接支持按输入输出张量拆分模型，不需要手动修改计算图结构，自动处理依赖的常量、初始化参数。

拆分边缘端运行的前半段（从原输入到拆分点）的代码：

# 原模型输入张量名，你可以通过print(graph.input)查看，MobileNetV2一般是["input_1:0"]
original_input = [i.name for i in graph.input]
# 拆分点的输出张量名，就是你上一步记录的splitter_tensor
edge_output = ["splitter_tensor"]

# 导出边缘端模型
onnx.utils.extract_model(
    "mobilenet_v2.onnx",
    "edge_part.onnx",
    input_names=original_input,
    output_names=edge_output
)

拆分云端运行的后半段（从拆分点到原输出）的代码：

# 后半段输入就是拆分点的张量名
cloud_input = ["splitter_tensor"]
# 原模型输出张量名，print(graph.output)可查看
original_output = [o.name for o in graph.output]

# 导出云端模型
onnx.utils.extract_model(
    "mobilenet_v2.onnx",
    "cloud_part.onnx",
    input_names=cloud_input,
    output_names=original_output
)

第三步：分段运行验证 #

用ONNX Runtime分别加载两段模型，验证结果和原模型一致即可：

import onnxruntime as ort
import numpy as np

# 生成测试输入，和原模型输入shape一致，MobileNetV2是(1,224,224,3)
test_input = np.random.randn(1,224,224,3).astype(np.float32)

# 运行原模型对比
ort_original = ort.InferenceSession("mobilenet_v2.onnx")
original_out = ort_original.run(None, {original_input[0]: test_input})[0]

# 运行边缘端前半段
ort_edge = ort.InferenceSession("edge_part.onnx")
edge_out = ort_edge.run(None, {original_input[0]: test_input})[0]

# 运行云端后半段
ort_cloud = ort.InferenceSession("cloud_part.onnx")
cloud_out = ort_cloud.run(None, {cloud_input[0]: edge_out})[0]

# 验证结果误差，数值接近说明拆分正确
print(np.mean(np.abs(cloud_out - original_out)))

注意事项：

• 拆分时保持原模型的opset版本一致，不要手动修改算子版本
• 如果拆分后模型运行报错，可先调用onnx.checker.check_model()检查拆分后的模型合法性
• 动态维度的模型拆分后如果运行出错，可以在导出时手动指定输入输出的shape固定维度

内容的提问来源于stack exchange，提问作者Starivore