ONNX系列: ONNX模型修改

ONNX 模型修改

        当我们熟悉了ONNX模型各个层级的结构后，我们便可以针对各个结构来对模型进行修改，从而使其更好的适配后端运行时或者特定硬件平台的编译器。对模型的修改通常可以概括为"增删改查"的操作。"增"是增加相应结构，"删"是删除相应结构，"改"是修改相应结构，"查"是获取到指定的模型结构。修改ONNX模型通常有两种思路，一是使用ONNX官方提供的Python API；二是使用第三方ONNX模型修改工具，例如onnx-graphsurgeon工具。本文将聚焦第一种方案，介绍如何使用ONNX官方API来对ONNX模型进行"增删改查"修改。完整的ONNX官方API文档可以参考：https://onnx.ai/onnx/index.html。

1. ONNX 模型的"查"

我们想要修改ONNX模型，首先需要知道如何定位到自己感兴趣的位置，比如如何找到具体某个节点、某个 initializer,、计算图的input/output, 某个节点的 input/output以及某个 value_info。参考下面的代码，我们可以发现定位某个元素的基本思路就是遍历该元素的列表，然后根据该元素在计算图中独有的属性名称来实现定位。下面的代码实现了定位下图模型的各个元素。

# 根据算子的名字来找到目标节点

for item in model.graph.node:

    if item.name == 'Conv_1':

        print(item)

# 有的onnx模型中算子没有name属性，可以根据算子类型和输出的名字来组合找到目标节点

for item in model.graph.node:

    if item.op_type == 'Conv':

        if '1338' in item.output:

            print(item)

# 找到目标 intializer

for i in model.graph.initializer:

    if i.name == '1339':

        print(i.dims)

        print(i.dims)

        print(i.data_type)

        # 二进制形式打印，可能比较长

        print(i.raw_data)

# 找到 graph 的input和output

for i in model.graph.input:

    if i.name == 'input':

        print(i.name)

        print(i.type)

# 找到 graph 的valueinfo

for i in model.graph.value_info:

    if i.name == '9':

        print(i.name)

        print(i.type)

2. ONNX 模型的"删"

在了解了如何定位到需要修改的部分后，我们就可以对ONNX模型进行魔改了。我们首先了解如何删除ONNX模型中的指定节点或元素。下面的代码实现了删除图中标注的节点。

import onnx

# 加载模型

model = onnx.load('./super-resolution-10.onnx')

# 根据输入获取指定节点

def get_node_with_input(model, input_name):

    res = []

    for i in model.graph.node:

        if input_name in i.input:

            res.append(i)

    return res

# 根据输出获取指定节点

def get_node_with_output(model, output_name):

    res = []

    for i in model.graph.node:

        if output_name in i.output:

            res.append(i)

    return res

# 删除指定节点并将前后节点连接起来

remove_nodes = []

p = None

n = None

for i in model.graph.node:

    if '10' in i.input:

        # p = find_node_with_output(i.input[0])

        p = get_node_with_output(model, i.input[0])[0]

        remove_nodes.append(i)

    if '11' in i.input:

        # n = find_node_with_input(i.output[0])

        n = get_node_with_input(model, i.output[0])[0]

        remove_nodes.append(i)

n.input[0] = p.output[0]

for i in remove_nodes:

    model.graph.node.remove(i)

onnx.checker.check_model(model)

onnx.save(model, 'super-resolution-10-delete.onnx')

3. ONNX 模型的"增"

"增"是指在ONNX模型指定位置添加节点。在了解添加节点之前，我们首先需要了解如何创建 ONNX 节点。下面以创建一个2D卷积算子和一个ReLu算子为例，并尝试将上一步骤中删除的这两个节点重新添加回模型当中（注意我们权重没有与原模型保持一致）。

node1 = onnx.helper.make_node(

        name="Conv_0",   # 节点名字，不要和op_type搞混了

        op_type="Conv",  # 节点的算子类型, 比如'Conv'、'Relu'、'Add'这类，详细可以参考onnx给出的算子列表

        inputs=["image", "conv.weight", "conv.bias"],  # 各个输入的名字，结点的输入包含：输入和算子的权重。必有输入X和权重W，偏置B可以作为可选。

        outputs=["11"], 

        pads=[1, 1, 1, 1], # 其他字符串为节点的属性，attributes在官网被明确的给出了，标注了default的属性具备默认值。

        group=1,

        dilations=[1, 1],

        kernel_shape=[3, 3],

        strides=[1, 1]

    )

initializer_w = onnx.helper.make_tensor(

        name="conv.weight",

        data_type=onnx.helper.TensorProto.DataType.FLOAT,

        dims=[64, 64, 3, 3],

        vals=np.ones([64,64,3,3], dtype=np.float32).tobytes(),

        raw=True

    )

initializer_b = onnx.helper.make_tensor(

        name="conv.bias",

        data_type=onnx.helper.TensorProto.DataType.FLOAT,

        dims=[64],

        vals=np.ones([64], dtype=np.float32).tobytes(),

        raw=True

    )

node2 = onnx.helper.make_node(

        name="ReLU_1",

        op_type="Relu",

        inputs=["11"],

        outputs=["12"]

    )

下面代码将上述创建的两个节点插入到模型指定位置。

for i in range(len(model.graph.node)):

    if '10' in model.graph.node[i].output:

        model.graph.node[i].output[0] = 'pre_output'

        model.graph.node[i+1].input[0] = 'relu_output'

        model.graph.node.insert(i+1, node1)

        model.graph.node.insert(i+2, node2)

model.graph.initializer.append(initializer_w)

model.graph.initializer.append(initializer_b)

input = model.graph.input[0]

new_input = onnx.helper.make_tensor_value_info(input.name, onnx.TensorProto.FLOAT, [1,1,224,224])

model.graph.input[0].CopyFrom(new_input)

onnx.checker.check_model(model)

model = onnx.shape_inference.infer_shapes(model)

onnx.save(model, 'super-resolution-10-insert.onnx')

4. ONNX 模型的"改"

通常来说修改 ONNX 模型可以概括为一下两种：

• 修改模型节点
• 修改权重（initializer）

修改模型的节点可以通过上述的删除 + 添加节点组合操作来实现，这里不再赘述。下面将介绍如何修改节点权重。节点权重通常保存在initializer中，下面代码尝试将Conv算子中的bias缩小10倍。

import onnx

model = onnx.load("./super-resolution-10.onnx")

# 得到所有 initializer

all_initializer = model.graph.initializer

# 定位到目标 initializer

target_initializer = 'conv1.bias'

idx = ''

scale_factor = 10

for i, j in enumerate(all_initializer):

    if j.name == target_initializer:

        idx = i

        break

# 将 conv1 算子的 bias 缩小10倍

model.graph.initializer[idx].raw_data = (onnx.numpy_helper.to_array(all_initializer[idx]) / scale_factor).tobytes()

onnx.save(model,'super-resolution-10-scale.onnx')

总结

当我们在实际部署模型时，会根据具体硬件特性来在 ONNX 模型层面做相应的优化修改，使其能在特定的硬件平台上获得更好的推理性能。本文简单介绍了如何调用 ONNX 官方API来对 ONNX 模型进行增删改查，更加复杂的模型修改操作通常是上述四种操作的各种组合。

使用ONNX 官方API需要我们对 ONNX 模型的定义和Proto结构足够熟悉，并且通过本文中的示例代码可以看到，繁多复杂的API在使用过程中也不是很方便。在实际工作中，我们一般使用NV提供的onnx-graphsurgeon工具来快速对ONNX模型进行修改验证。这个工具在官方ONNX API的基础上提供了更为友好的高级API封装，大大提升了我们修改ONNX模型的效率，在之后的文章中我们将进一步详细介绍这个工具的使用。

作者：高通工程师，阮慧源（Huiyuan Ruan）