onnx文件新增 sub/mul前处理节点

发布时间 2023-09-25 15:25:33作者: WenJXUST
#更改onnx里的节点参数内容
import onnx

model = onnx.load("best20230908sim.onnx")
graph = model.graph
initializer = graph.initializer

node = graph.node

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

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

#插入sub
# sub_const_node = onnx.helper.make_tensor(name='const_sub',
#                       data_type=onnx.TensorProto.FLOAT,
#                       dims=[1],
#                       vals=[-0])
# graph.initializer.append(sub_const_node)
# sub_node = onnx.helper.make_node(
#                 'Add',
#                 name='pre_sub',
#                 inputs=['data', 'const_sub'],
#                 outputs=['pre_sub'])
# graph.node.insert(0, sub_node)

# 插入mul
mul_const_node = onnx.helper.make_tensor(name='const_mul',
                                         data_type=onnx.TensorProto.FLOAT,
                                         dims=[1,3,1,1],
                                         vals=[1.0 / 255.0,1.0/255.0,1.0/255.0])

graph.initializer.insert(0,mul_const_node)

mul_node = onnx.helper.make_node(
    'Mul',
    name='pre_mul',
    inputs=['data', 'const_mul'],
    outputs=['pre_mul'])
graph.node.insert(0, mul_node)

# 第一层卷积的输入修改
for id, dnode in enumerate(graph.node):
    for i, input_node in enumerate(dnode.input):
        if 'data' == input_node and dnode.name != 'pre_mul':
            dnode.input[i] = 'pre_mul'



onnx.checker.check_model(model)
onnx.save(model, 'best20230908sim_out.onnx')

参考:onnx 创建、修改、增加、修改、删除节点_修改 onnx 节点_AI视觉网奇的博客-CSDN博客