承接上回《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的内部节点。
import onnx
onnx_model = onnx.load("test.onnx")
graph = onnx_model.graph
node = graph.node
for i in range(len(node)):
print(node[i])
我们可以依靠上述代码输出该模型的节点个数,还有节点中的属性信息,当然也包含静态图的链路形状。
接着,依据节点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]
直接访问这个节点了。
简单来说,就像链表的插入操作一样,即是删除、新建、插入。如下列代码所示:
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模型了。
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')