双输入双输出单输入单输出
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)
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)
model1_input_tensor_type = list_model[0].graph.input[0].type.tensor_type
input_elem_type = model1_input_tensor_type.elem_type
input_shape = [1,3,540,960]
modelX_input = onnx.helper.make_tensor_value_info(
input_name,
input_elem_type,
input_shape,
)
modelX.graph.input.append(modelX_input)
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 outputss in model.graph.output:
outputss.name =list_prefix[idx]+ outputss.name
modelX.graph.output.append(outputss)
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')