ONNX 模型的修改与调试_onnx.modelproto

主要问题：

ONNX 模型在底层是用什么格式存储的？
如何不依赖深度学习框架，只用 ONNX 的 API 来构造一个 ONNX 模型？
如果没有源代码，只有一个 ONNX 模型，该如何对这个模型进行调试？

ONNX 的底层实现

存储格式

ONNX的底层是用 Protobuf 定义的。是 Google 提出的一套表示和序列化数据的机制。
对于 ONNX ，它的 Protobuf 数据定义文件在其开源库中，这些文件定义了神经网络中模型、节点、张量的数据类型规范。
ONNX 提供了很多实用 API，我们可以在完全不了解 Protobuf 的前提下，构造和读取 ONNX 模型。

结构定义

图一般是用一个节点集和一个边集表示的。
神经网络本质上是一个计算图。计算图的节点是算子，边是参与运算的张量。
可视化 ONNX 模型， 能看到所有算子节点的属性信息，参与运算的张量信息存储在算子节点的输入输出信息中。

一个 ONNX 模型可以用 ModelProto 类表示。
ModelProto 包含了版本、创建者等日志信息，还包含了存储计算图结构的 graph。
GraphProto 类则由输入张量信息、输出张量信息、节点信息组成。
张量信息 ValueInfoProto 类包括张量名、基本数据类型、形状。
节点信息 NodeProto 类包含了算子名、算子输入张量名、算子输出张量名。
ModelProto模型信息{GraphProto图信息{NodeProto节点信息, ValueInfoProto节点信息}}
下面演示实现 output=a*x+b 功能的ONNX模型

ir_version:8, 
graph{
    node {input: "a",input: "x",output: "c",op_type: "Mul"}, 
node {input: "c",input: "b",output: "output",op_type: "Add"},
name: "linear_func",
input {name: "a",
       type {tensor_type {
           elem_type: 1, 
           shape {dim {dim_value: 10}, dim {dim_value: 10}}
       }}
      },
input {name: "x",
       type {tensor_type {
           elem_type: 1, 
           shape {dim {dim_value: 10}, dim {dim_value: 10}}
       }}
      },
input {name: "b",
       type {tensor_type {
           elem_type: 1, 
           shape {dim {dim_value: 10}, dim {dim_value: 10}}
       }}
      }, 
output {name: "output",
        type {tensor_type {
            elem_type: 1, 
            shape {dim { dim_value: 10}, dim { dim_value: 10}}
        }}
             }
       },
opset_import {version: 15}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

读写 ONNX 模型

构造 ONNX 模型

完全用 ONNX 的 Python API 构造一个描述线性函数 output=a*x+b 的 ONNX 模型。(自底向上实现)

import onnx
from onnx import helper
from onnx import TensorProto

# 用 helper.make_tensor_value_info 构造描述张量信息的 ValueInfoProto 对象
# 形参列表： 张量名、张量的基本数据类型、张量形状
# 用此的方式为三个输入 a, x, b 和一个输出 output 构造 ValueInfoProto 对象。
a = helper.make_tensor_value_info('a', TensorProto.FLOAT, [10, 10])
x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [10, 10])
b = helper.make_tensor_value_info('b', TensorProto.FLOAT, [10, 10])
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [10, 10])

# 可以通过 helper.make_node 构造算子节点信息 NodeProto
# 形参列表：算子类型、输入张量名、输出张量名
# 先构造了描述 c=a*x 的乘法节点，再构造了 output=c+b 的加法节点
mul = helper.make_node('Mul', ['a', 'x'], ['c'])
add = helper.make_node('Add', ['c', 'b'], ['output'])

# 用 helper.make_graph 来构造计算图 GraphProto
# 形参列表：节点、图名称、输入张量信息、输出张量信息
# make_graph 的节点参数要求：计算图的节点必须以拓扑序给出
# 把之前构造出来的 NodeProto 对象和 ValueInfoProto 对象按照顺序传入
graph = helper.make_graph([mul, add], 'linear_func', [a, x, b], [output])

# 用 helper.make_model 把计算图 GraphProto 封装进模型 ModelProto 里，ONNX 模型构造完成
# 可选参数：模型制作者、版本等信息
model = helper.make_model(graph)

onnx.checker.check_model(model) # 检查模型正确性
print(model) # 把模型以文本形式输出
onnx.save(model, 'linear_func.onnx') # 导出 ONNX 模型
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

注意事项：

• ONNX 对节点的输入有一定的要求：一个节点的输入，要么是整个模型的输入，要么是之前某个节点的输出
• 不满足标准的 ONNX 模型可能无法被推理引擎正确识别。使用 onnx.checker.check_model 来判断一个 ONNX 模型是否满足标准

用 ONNX Runtime 运行模型，测试模型正确性：

import onnxruntime
import numpy as np

sess = onnxruntime.InferenceSession('linear_func.onnx')
a = np.random.rand(10, 10).astype(np.float32)
b = np.random.rand(10, 10).astype(np.float32)
x = np.random.rand(10, 10).astype(np.float32)

output = sess.run(['output'], {'a': a, 'b': b, 'x': x})[0]

assert np.allclose(output, a * x + b)
1
2
3
4
5
6
7
8
9
10
11

读写并修改 ONNX 模型

import onnx
# 读取一个 ONNX 模型
model = onnx.load('linear_func.onnx')
# print(model)

graph = model.graph
node = graph.node
input_ = graph.input
output_ = graph.output
print(node)
print(input_)
print(output_)

# node, input_, output_都是列表
node_0 = node[0]
node_0_inputs = node_0.input
node_0_outputs = node_0.output
input_0 = node_0_inputs[0]
input_1 = node_0_inputs[1]
output = node_0_outputs[0]
op_type = node_0.op_type

print(input_0)
print(input_1)
print(output)
print(op_type)

# 修改模型属性
node = model.graph.node
node[1].op_type = 'Sub'
onnx.checker.check_model(model)
onnx.save(model, 'linear_func_2.onnx')

# 添加节点时, 先创建节点, 然后添加到图中。它的输入是上个节点的输出, 输出是下个节点的输入
node = model.graph.node
mul = helper.make_node('Mul', ['c', 'b'], ['d'])

node[1].input[0] = mul.output[0]  # 输出是下个节点的输入
model.graph.node.append(mul)  # 新节点加入model的node列表中
node[0].output[0] = mul.input[0]  # 输入是上个节点的输出
# 修改后模型的功能: (a * x) * b + b
onnx.save(model, 'linear_func_3.onnx')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

调试 ONNX 模型

学习如何巧妙利用 ONNX 提供的子模型提取功能，对 ONNX 模型进行调试。

子模型提取

就是从一个给定的 ONNX 模型中，拿出一个子模型。这个子模型的节点集、边集都是原模型中对应集合的子集。

import torch

class Model(torch.nn.Module):

    def __init__(self):
        super().__init__()
        self.convs1 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3),
                                          torch.nn.Conv2d(3, 3, 3),
                                          torch.nn.Conv2d(3, 3, 3))
        self.convs2 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3),
                                          torch.nn.Conv2d(3, 3, 3))
        self.convs3 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3),
                                          torch.nn.Conv2d(3, 3, 3))
        self.convs4 = torch.nn.Sequential(torch.nn.Conv2d(3, 3, 3),
                                          torch.nn.Conv2d(3, 3, 3),
                                          torch.nn.Conv2d(3, 3, 3))
    def forward(self, x):
        x = self.convs1(x)
        x1 = self.convs2(x)
        x2 = self.convs3(x)
        x = x1 + x2
        x = self.convs4(x)
        return x

model = Model()
input = torch.randn(1, 3, 20, 20)

torch.onnx.export(model, input, 'whole_model.onnx')

# 编号22的边对应/convs1/convs1.1/Conv_output_0，编号28的边对应Add_output_0
# 参数分别是原模型路径、输出模型路径、子模型的输入边（输入张量）、子模型的输出边（输出张量）
onnx.utils.extract_model('whole_model.onnx', 'partial_model.onnx', ['/convs1/convs1.1/Conv_output_0'], ['/Add_output_0'])

# 先使用netron打开onnx模型, 然后查看对应边的name, 将编号修改为name才能运行
# 下面的代码将会有两个输出
onnx.utils.extract_model('whole_model.onnx', 'submodel_1.onnx', ['22'], ['27', '31'])
# 下面代码将不会达到理想中的功能
onnx.utils.extract_model('whole_model.onnx', 'submodel_2.onnx', ['22', 'input.1'], ['28'])
# 下面代码报错
# onnx.utils.extract_model('whole_model.onnx', 'submodel_3.onnx', ['24'], ['28'])

# 将原模型的多个中间输出作为最终输出
onnx.utils.extract_model('whole_model.onnx', 'more_output_model.onnx', ['input.1'], ['31', '23', '25', '27'])
# 继续导出更小的模型子图
onnx.utils.extract_model('more_output_model.onnx', 'debug_model_1.onnx', ['input.1'], ['23'])
onnx.utils.extract_model('more_output_model.onnx', 'debug_model_2.onnx', ['23'], ['25'])
onnx.utils.extract_model('more_output_model.onnx', 'debug_model_3.onnx', ['23'], ['27'])
onnx.utils.extract_model('more_output_model.onnx', 'debug_model_4.onnx', ['25', '27'], ['31'])
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

总结：

• ONNX 使用 Protobuf 定义规范和序列化模型。
• 一个 ONNX 模型主要由 ModelProto,GraphProto,NodeProto,ValueInfoProto 这几个数据类的对象组成。
• 使用 onnx.helper.make_xxx，我们可以构造 ONNX 模型的数据对象。
• onnx.save() 可以保存模型，onnx.load() 可以读取模型，onnx.checker.check_model() 可以检查模型是否符合规范。
• onnx.utils.extract_model() 可以从原模型中取出部分节点，和新定义的输入、输出边构成一个新的子模型。
• 利用子模型提取功能，我们可以输出原 ONNX 模型的中间结果，实现对 ONNX 模型的调试。