onnx合二为一

双输入双输出单输入单输出

# -*- coding:utf-8 -*-
import onnx
import copy
from onnx import TensorProto
yolox_model = 'yolox.onnx'
eff_model = 'effl1.onnx'
list_modelname = [yolox_model, eff_model]
# 模型中节点和输入输出名字存在相同的，合并后要加个前缀
list_prefix = ['m1_', 'm2_']

model_num = len(list_modelname)
input_name = 'img'
output_name = 'out'

list_model = []
list_inputs = []
list_outputs = []
for i in range(model_num):
    model = onnx.load(list_modelname[i])
    list_model.append(model)

input_1 = list_model[0].graph.input
output_1 = list_model[0].graph.output
input_2 = list_model[1].graph.input
output_2 = list_model[1].graph.output

list_inputs.append(input_1)
list_inputs.append(input_2)
list_outputs.append(output_1)
list_outputs.append(output_2)

# 0、构造新模型，以list_model[1]结构特点为新模型结构特征
modelX = onnx.ModelProto(ir_version=list_model[1].ir_version,
                         producer_name=list_model[1].producer_name,
                         producer_version=list_model[1].producer_version,
                         opset_import=list_model[1].opset_import)
# 1、添加 input，可以选择list_model[0]或者list_model[1]的输入格式
model1_input_tensor_type = list_model[0].graph.input[0].type.tensor_type
input_elem_type = model1_input_tensor_type.elem_type
# 固定输入的shape
input_shape = [1,3,540,960]
# 按照相应模型的shape为基础
# input_shape= []
# for s in model1_input_tensor_type.shape.dim:
#     if (s.dim_value > 0):
#         input_shape.append(s.dim_value)
#     else:
#         input_shape.append(s.dim_param)
#2 添加输入节点
modelX_input = onnx.helper.make_tensor_value_info(
    input_name,
    input_elem_type,
    input_shape,
)
modelX.graph.input.append(modelX_input)

# 3、添加训练好的模型中的节点和权重
for idx in range(model_num):
    model = list_model[idx]
    for node in model.graph.node:
        for i in range(len(node.input)):
            if (node.input[i] != input_name):
                node.input[i] = list_prefix[idx] + node.input[i]

        for i in range(len(node.output)):
            node.output[i] = list_prefix[idx] + node.output[i]

        node.name = list_prefix[idx] + node.name
        modelX.graph.node.append(node)

    for weight in model.graph.initializer:
        weight.name = list_prefix[idx] + weight.name
        modelX.graph.initializer.append(weight)
    # 这里打开即为双输入双出，注释掉即为单输入双输出
    # for inputss in model.graph.input:
    #     inputss.name = list_prefix[idx] +inputss.name
    #     modelX.graph.input.append(inputss)

    for outputss in model.graph.output:
        outputss.name =list_prefix[idx]+ outputss.name
        modelX.graph.output.append(outputss)
#
# # 5、单输出时，可以将该节点添加 ，即模型的output
# model1_output_tensor_type = list_model[0].graph.output[0].type.tensor_type
# output_elem_type = model1_output_tensor_type.elem_type
# output_shape = []
# for s in model1_output_tensor_type.shape.dim:
#     if (s.dim_value > 0):
#         output_shape.append(s.dim_value)
#     else:
#         output_shape.append(s.dim_param)
#
# modelX_output = onnx.helper.make_tensor_value_info(
#     output_name,
#     output_elem_type,
#     output_shape
# )
# modelX.graph.output.append(modelX_output)

# 6、保存新模型
onnx.checker.check_model(modelX)
onnx.save(modelX, output_path)
model = onnx.load(output_path)
node = model.graph.node
# 将新增节点和之前的节点链接，按照节点位置和名称进行修改
for i in range(len(node)):
    if node[i].op_type == "Shape" and node[i].name == "m2_Shape_544":
        print(i)
        print(node[i].input[0])
        node[i].input[0] = "m2_pre_Mul"
    if node[i].op_type == "Resize" and node[i].name == "m1_":
        print(i)
        print(node[i].input[0])
        node[i].input[0] = "img"
        node[i].name = "m1_Resize"

onnx.checker.check_model(model)
onnx.save(model, './effl1_yolox_s2.onnx')