编辑ONNX模型的python代码

参考：(8条消息) 编辑ONNX模型的python代码_Alexa2077的博客-CSDN博客

编辑ONNX的python代码
一、ONNX模型的基本操作
1，加载ONNX模型
2，保存ONNX模型
3，OP节点列表
4，输入节点名称
5，输出节点名称
6，参数节点
二、ONNX模型的修改
1，修改内部的变量
2，创建tensor
3，增加OP节点
4，增加输入\输出tensor节点
5，增加参数节点
6，特殊节点-constant增加
7，读取ONNX的参数tensor格式，转换为numpy
三、例程
得到第一个Conv的节点信息
注意：
在运行程序的时候，记得安装onnx的环境，在pycharm里面切换到onnx的解释器。
关于onnx的更多信息可以查看:https://github.com/onnx
一、ONNX模型的基本操作
1，加载ONNX模型
 model = onnx.load("test.onnx") 

2，保存ONNX模型
 onnx.save(model, 'save_path.onnx') 

3，OP节点列表
 model.graph.node 是List, 可以增删改查 

4，输入节点名称
 model.graph.input 

5，输出节点名称
 model.graph.output 

6，参数节点
 model.graph.initializer 

二、ONNX模型的修改
1，修改内部的变量
可直接修改内部的变量，如改变任意OP名称或者属性。

model.graph.node[0].op_type= '自定义'
改变输入的batch=1，model.graph.input[0].type.tensor_type.shape.dim[0].dim_value=1

 

当然也可以采用下面方法创建新的节点插入

2，创建tensor
ONNX模型采用tensor与node（op）将整个网络构建起来，node中输入输出为tensor的名称，用来连接各个tensor，中间的节点tensor不用创建，网络的输入输出节点需要创建tensor。Tensor常用的2种创建方法，样例如下：

1) onnx.helper.make_tensor_value_info(name=output, elem_type=type, shape=[1,3,224])
仅根据形状创建，用于输入输出节点
2）onnx.helper.make_tensor(name='const_tensor', data_type,dims=values.shape,vals=values.flatten()))，
根据numpy数据创建

3，增加OP节点

 函数：add_node = onnx.helper.make_node(op_type='Conv', inputs=input_names, outputs=output_names,name='conv1', **attribute) 

添加OP节点到图中：

model.graph.node.append(add_node)
model.graph.node.insert(insert_pos, add_node)

其中，input_names、output_names 是字符串数组，为tensor的名称，例如,input_names=[‘name1’, ‘name2’]

attribute是字典，如 attribute[‘eps’] = 0.0001

4，增加输入\输出tensor节点
函数：

output_node = onnx.helper.make_tensor_value_info(name=output, elem_type=type, shape=[1,3,224])
model.graph.output.append(output_node)

 

5，增加参数节点

params = numpy_helper.from_array(numpy, name='p_name')
model.graph.initializer.append(params)

 

6，特殊节点-constant增加
函数：

value=onnx.helper.make_tensor(name='const_tensor', data_type=1, dims=numpy_data.shape,vals=numpy_data.flatten()))
constant_node = onnx.helper.make_node(op_type='Constant', inputs=[], outputs=[output_name],name='constant1', value = value)

 

7，读取ONNX的参数tensor格式，转换为numpy

constant节点：onnx.numpy_helper.to_array(onnx_node.attribute[0].t)
initialize的参数：onnx.numpy_helper.to_array(params)

 

特别数值，若数值为1个数，则导出的非numpy格式，也是一个数值

三、例程
本文使用打开onnx模型的软件为netron，本文使用mobilenetv2-7.onnx模型，从官网的模型库中下载的。
onnx官网：https://onnx.ai/ 进行查找和下载

使用topo_tailoring_infer_smoke_resnet18_704x384_fp16_2_Conv_161.onnx (netron.app)打开xxx.onnx文件
打开onnx模型如下图所示：

 

 

得到第一个Conv的节点信息

import onnx
model = onnx.load("mobilenetv2-7.onnx")
nodel_list = model.graph.node
input_list = model.graph.input
output_list = model.graph.output
initializer_list = model.graph.initializer
print(nodel_list[0])

 

运行之后，打印如下结果：

input: "input"
input: "475"
input: "476"
output: "474"
name: "Conv_0"
op_type: "Conv"
attribute {
name: "dilations"
ints: 1
ints: 1
type: INTS
}
attribute {
name: "group"
i: 1
type: INT
}
attribute {
name: "kernel_shape"
ints: 3
ints: 3
type: INTS
}
attribute {
name: "pads"
ints: 1
ints: 1
ints: 1
ints: 1
type: INTS
}
attribute {
name: "strides"
ints: 2
ints: 2
type: INTS
}

 

我们使用netron打开该模型的第一个Conv节点进行对照：
可以看到input,output,等等信息都可以得到，其他的修改部分按照上面给出的语法进行操作即可。
后面的操作就相当于对着List进行操作。

 

注意：
在运行程序的时候，记得安装onnx的环境，在pycharm里面切换到onnx的解释器。
关于onnx的更多信息可以查看:https://github.com/onnx