ONNX内部节点修改方法

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

我们可以依靠上述代码输出该模型的节点个数，还有节点中的属性信息，当然也包含静态图的链路形状。

2. 搜索目标节点

接着，依据节点ID找到我们需要修改的Resize节点，这里需要注意的是，Netron可视化出来的id需要经过转换才可以得到ONNX的实际ID，就像相对路径之于绝对路径一样，为了方便，这里就推荐直接将打印出来的节点信息拷贝出来，进行关键字查找。

比如这里，我的Resize输出id为450，那么就用450作为关键字进行搜索，得到下图结果。

看起来，这个449就是对应的scales的节点onnx.Constant，所以按照下列代码，将这个节点的真实ID搜索出来，得到的结果是i=157

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

我们就可以直接使用node[157]直接访问这个节点了。

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模型了。

4. 检查图与保存

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

可以看到，已成功修改。

完整代码

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

 

 

参考链接：https://blog.csdn.net/github_28260175/article/details/105736654