如何修改已有的ONNX模型

简单来说，我们只需要学习一下把大象如何放进冰箱的就行了：

1、把冰箱门打开

使用onnx的原生接口：

onnx_model = onnx.load(onnx_path)
graph = onnx_model.graph

这样我们就可以将模型load出来，并且到到graph信息。

2、把大象放进去

这一步相对来说选择就比较多了，比如你可以选择删除一些节点，修改一下节点，增加一些节点。

删除：这个是最容易的，直接一句话 graph.node.remove(xxx_node)

修改：举个例修改一下input的名称

for input_node in onnx_model.graph.input:
    if 'input_xxx' == input_node.name:
        print("change input data name")
        input_node.name = 'data'

就是拿到某个属性或者信息，然后直接修改就行了。

增加：举个例增加一组图像预处理操作(减均值，除方差)

这一步稍微复杂一点，我们首先要生成一个node或者tensor，然后将这个node或者tensor加入graph中，然后选择性的增加一个node来操作刚刚加入graph的node或者tensor。

首先我们生成一个tensor，就是需要减去的均值

sub_const_node = onnx.helper.make_tensor(name='const_sub',
                      data_type=onnx.TensorProto.FLOAT,
                      dims=[1],
                      vals=[-127.5])

然后我们将刚刚生成的tensor插入graph中

graph.initializer.append(sub_const_node)

然后我们再增加一个减均值的node

sub_node = onnx.helper.make_node(
                'Add',
                name='pre_sub',
                inputs=['data', 'const_sub'],
                outputs=['pre_sub'])

然后将node加入graph中

graph.node.insert(0, sub_node)

仿造这样的流程我们继续加入除以方差的操作

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

sub_node = onnx.helper.make_node(
               'Mul',
               name='pre_mul',
               inputs=['pre_sub', 'const_mul'],
               outputs=['pre_mul'])
graph.node.insert(1, sub_node)

这样操作之后，我们还需要一步，就是将第一个卷积层的输入改动一下：

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

这样能我们加入node或者tensor的过程基本就结束了

3、把冰箱门关上

这一步我们就可以简单的重组一下graph，然后save模型就行了

graph = onnx.helper.make_graph(graph.node, graph.name, graph.input, graph.output, graph.initializer)
info_model = onnx.helper.make_model(graph)
onnx_model = onnx.shape_inference.infer_shapes(info_model)
 
onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, onnx_path.replace('nopre', 'fix'))

完事我们就拿到一个修改过的模型了，找一个人脸的模型来示例一下：

红圈就是增加的节点，当然你也可以增加其它节点。注意按照onnx的op要求实现就行了，给一个参考的路径：

出处：如何修改已有的ONNX模型 - 知乎 (zhihu.com)