onnx 创建、修改、增加、修改、删除节点

目录

onnx创建示例

onnx推理示例

报错At least one output should be requested：

1. 修改目标节点

1.1 载入ONNX文件

3. Op裁剪

4. PRelu参数修改

---

onnx创建示例

import onnx
from onnx import helper
from onnx import TensorProto
import numpy as np

weight = np.random.randn(36)
X = helper.make_tensor_value_info('input', TensorProto.FLOAT, [1, 2, 4, 4])
W = helper.make_tensor('W', TensorProto.FLOAT, [2, 2, 3, 3], weight)
B = helper.make_tensor('B', TensorProto.FLOAT, [2], [1.0, 2.0])
output = helper.make_tensor_value_info('output', TensorProto.FLOAT, [1, 2, 2, 2])
node_def = helper.make_node('Conv',  # node name
    ['input', 'W', 'B'], ['output'],  # outputs
    # attributes
    strides=[2, 2], )
graph_def = helper.make_graph([node_def], 'test_conv_mode', [X],  # graph inputs
    [output],  # graph outputs
    initializer=[W, B], )
mode_def = helper.make_model(graph_def, producer_name='onnx-example')
onnx.checker.check_model(mode_def)
onnx.save(mode_def, "./Conv.onnx")

有output:

删除output：

onnx推理示例

import os

import cv2

import onnxruntime
import time
os.environ['TF_CPP_MIN_LOG_LEVEL']='3'
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

import numpy as np

if __name__ == '__main__':

    dummy_input = np.ones([1, 4, 4, 2], dtype=np.float32)
    # img_o = cv2.resize(img, (192, 192))
    # img = img_o.astype(np.float32)
    # img/=255.0


    model_path = r'Conv2.onnx'
    # model_path = r'C:\Users\Administrator\Pictures\mm\resources.tar\resources\face_landmark_with_attention_192x192/model_float32.onnx'
    model_file_name = model_path.split(".")[0]
    session_option = onnxruntime.SessionOptions()
    session_option.optimized_model_filepath = f"{model_file_name}_cudaopt.onnx"
    session_option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
    session = onnxruntime.InferenceSession(
        model_path,
        session_option,
        providers=['CUDAExecutionProvider']
    )

    input_name = session.get_inputs()[0].name
    output_names = [o.name for o in session.get_outputs()]
    input_shape = session.get_inputs()[0].shape
    # Warmup
    output = session.run(
        output_names,
        # {input_name: dummy_input}
        {input_name: dummy_input.transpose((0,3,1,2))}
    )
    # Inference

    for i in range(1):
        start = time.time()
        output = session.run(output_names,{input_name: dummy_input.transpose((0,3,1,2))})
        print(output)



        # cv2.imshow("result",img_o)
        # cv2.waitKey(0)
        # print(f'elapsed_time : {(time.time()-start)*1000} ms')

报错At least one output should be requested：

onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : At least one output should be requested.

原因：make_graph时没有指定output：

graph_def = helper.make_graph([node_def], 'test_conv_mode', [X],  # graph inputs
    [],  # graph outputs
    # [output],  # graph outputs
    initializer=[W, B], )
mode_def = helper.make_model(graph_def, producer_name='onnx-example')

解决方法：

加上output即可：

graph_def = helper.make_graph([node_def], 'test_conv_mode', [X],  # graph inputs
    [output],  # graph outputs
    # [output],  # graph outputs
    initializer=[W, B], )
mode_def = helper.make_model(graph_def, producer_name='onnx-example')

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.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)

原文链接：https://blog.csdn.net/weixin_41521681/article/details/112724867