onnx修改教程

1. 修改目标节点

1.1 载入ONNX文件

输出该模型的节点个数，还有节点中的属性信息，当然也包含静态图的链路形状。

import onnx

onnx_model = onnx.load("test.onnx")
graph = onnx_model.graph
node  = graph.node

for i in range(len(node)):
	print(node[i])

1.2 搜索目标节点

for i in range(len(node)):
    if node[i].op_type == 'Constant':
        node_rise = node[i]
        if node_rise.output[0] == '449':
            print(i)  # 157

1.3 修改目标节点

就像链表的插入操作一样，即是删除、新建、插入。如下列代码所示：

old_scale_node = node[157]
new_scale_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=['449'],
    value=onnx.helper.make_tensor('value', onnx.TensorProto.FLOAT, [4], [1, 1, 1.81, 1.81])
)  # 新建新节点
graph.node.remove(old_scale_node)  # 删除旧节点
graph.node.insert(157, new_scale_node)  # 插入新节点

具体onnx.helper.make_node的使用方法，可以去github上查找doc，然后就可以愉快地随意修改ONNX模型了。

1.4 检查图与保存

onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, 'out.onnx')

1.5 完整代码

import onnx

onnx_model = onnx.load("test.onnx")
graph = onnx_model.graph
node  = graph.node

old_scale_node = node[157]
new_scale_node = onnx.helper.make_node(
    "Constant",
    inputs=[],
    outputs=['449'],
    value=onnx.helper.make_tensor('value', onnx.TensorProto.FLOAT, [4], [1, 1, 1.81, 1.81])
)
graph.node.remove(old_scale_node)  
graph.node.insert(157, new_scale_node) 

onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, 'out.onnx')

2. 更改网络输入输出

import onnx
import math

input_size =(1080,1920)
model = onnx.load_model("centerface.onnx")
d = model.graph.input[0].type.tensor_type.shape.dim
print(d)
rate = (int(math.ceil(input_size[0]/d[2].dim_value)),int(math.ceil(input_size[1]/d[3].dim_value)))
print("rare",rate)
d[0].dim_value = 1
d[2].dim_value *= rate[0]
d[3].dim_value *= rate[1]
for output in model.graph.output:
    d = output.type.tensor_type.shape.dim
    print(d)
    d[0].dim_value = 1
    d[2].dim_value  *= rate[0]
    d[3].dim_value  *= rate[1]

onnx.save_model(model,"centerface_1088_1920.onnx" )

3. Op裁剪

思路很简单，我们录入的模型是一个DAG，裁剪后也要保持网络为一个DAG。如原始的预处理的数据流为data—>Identity—>Sub—>Mul—>Conv0，我们在进行Op裁剪以后，数据流则变成：data—>Conv0。核心代码也非常简单，我们只取从Conv0往后的nodes，然后再修改Conv0的input为data即可。

# remove preprocessing node
new_nodes = old_nodes[3:]
del model.graph.node[:]
model.graph.node.extend(new_nodes)

conv0_node = model.graph.node[0]
conv0_node.input[0] = "data"
graph = model.graph

4. PRelu参数修改

PRelu的参数为一维的，我们在onnxruntime中进行inference的时候可能无法正常进行broadcast（graph optimization阶段也无法进行），所以我们的思路很直接：直接修改slope参数的shape信息，若原来是1x64的向量，则将其shape信息改成(64, 1, 1)。修改后可以正常进行broadcast。核心代码如下，我们需要构建input maps（节点信息）和initializer maps（权值信息），然后遍历所有node去找出prelu的节点进行修改。

input_maps = {
    
    }
init_maps = {
    
    }
keys = []

for inp in model.graph.input:
	input_maps[inp.name] = inp
	keys.append(inp.name)

for init in model.graph.initializer:
	init_maps[init.name] = init

for key in keys:
	if "relu" in key:
		## revise input
		inp = input_maps[key]
		dim_value = inp.type.tensor_type.shape.dim[0].dim_value
		new_shape = [dim_value, 1, 1]
		graph.input.remove(inp)
		
		new_inp = onnx.helper.make_tensor_value_info(inp.name, onnx.TensorProto.FLOAT, new_shape)
		graph.input.extend([new_inp])
		
		## revise init
		init = init_maps[key]
		new_init = onnx.helper.make_tensor(init.name, onnx.TensorProto.FLOAT, new_shape, init.float_data)
		graph.initializer.remove(init)
		graph.initializer.extend([new_init])

5. BatchNormalization参数修改

迭代遍历所有节点找到BN层，然后对其spatial参数进行修改就好。下面为核心代码部分：

## revise batch norm op type
for node in model.graph.node:
	if (node.op_type == "BatchNormalization"):
		for attr in node.attribute:
			if (attr.name == "spatial"):
				attr.i = 1
			

6. 模型导出验证

最后在进行模型导出的时候，可以借助onnx的工具进行graph的有效性验证，调用如下的api即可：onnx.checker.check_model。此外进行裁剪和修改的模型，需要与原有模型进行数值对比，我采用的方案为计算原始graph和修改后graph的cos similarity，一般如果流程没有出错的话，cos similarity会非常接近1.0。如下为核心代码片：

import cv2
import numpy as np
from numpy import linalg as LA

orig_session = onnxruntime.InferenceSession("./r18_orig.onnx", None)
revs_session = onnxruntime.InferenceSession("./r18_revised.onnx", None)

image = cv2.imread("./Tom_Hanks_574745.png")
image = cv2.resize(image, (112, 112))
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

orig_data = np.array(image).transpose(2, 0, 1)
orig_data = orig_data.reshape(1, 1, 112, 112).astype("float32")
input_name = orig_session.get_inputs()[0].name
orig_result = orig_session.run([], {
    
    input_name: orig_data})

revs_data = (iamge - 127.5) / 128.0
revs_data = np.array(revs_data).transpose(2, 0, 1)
revs_data = revs_data.reshape(1, 3, 112, 112).astype("float32")
revs_result = revs_session.run([], {
    
    input_name: revs_data})

# calculate cos similarity
result_dot = np.dot(orig_result[0], revs_result[0])
result_norm = LA.norm(orig_result[0], 2) * LA.norm(revs_result[0], 2)
cos_similarity = result_dot / result_norm
print(cos_similarity)

ref:

https://blog.csdn.net/github_28260175/article/details/105736654
https://zhuanlan.zhihu.com/p/212893519