ONNX内部节点修改方法

承接上回《PyTorch转ONNX之F.interpolate》，因为op10的计算输出大小问题，导致我上采样的结果的大小出现小数，由预期输出结果output_size=[1., 3., 9., 9.]变成了output_size=[1., 3., 8.999, 8.999]，经过后续强制转换操作抹平成为了output_size=[1, 3, 8, 8]，这就很气了。

如下图所示，输入大小为input_size=[1, 3, 5, 5]，scales为[1, 1, 1.799, 1.799]，根据input_size x scales = output_size，输出大小应为output_size=[1., 3., 8.996, 8.996]，按之前的描述，后续操作会向下取整得到实际输出大小output_size=[1, 3, 8, 8]。那么，如果我将scales人工修改为[1, 1, 1.801, 1.801]不就可以避开这个问题了吗。因此，接下来的问题就是，如何修改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

 
 
  
  
 
 
  
  
   
   1
   
   2
   
   3
   
   4
   
   5

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

扫描二维码关注公众号，回复： 15245382 查看本文章

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)  # 插入新节点

 
 
  
  
 
 
  
  
   
   1
   
   2
   
   3
   
   4
   
   5
   
   6
   
   7
   
   8
   
   9

具体onnx.helper.make_node的使用方法，可以去github上查找doc，然后就可以愉快地随意修改ONNX模型了。

4. 检查图与保存

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

 
 
  
  
 
 
  
  
   
   1
   
   2

可以看到，已成功修改。
在这里插入图片描述

完整代码

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

ONNX内部节点修改方法

在这里插入图片描述

1. 载入ONNX文件

import onnx

ONNX模型解析修改结点

ONNX内部节点修改方法

1. 载入ONNX文件

2. 搜索目标节点

3. 修改目标节点

4. 检查图与保存

完整代码

ONNX内部节点修改方法

1. 载入ONNX文件

猜你喜欢