onnx手术刀（ONNX-GraphSurgeon）：对模型的输入端进行增，删，改操作（一）

一、引言

ONNX（Open Neural Network Exchange）是一种开放格式，用于表示深度学习模型。它旨在促进不同框架之间的模型互操作性。然而，在实际应用中，我们可能需要对模型进行定制和优化，以满足特定场景的需求。ONNX-GraphSurgeon正是为此而生，它允许开发者轻松地修改和优化ONNX模型。

二、ONNX-GraphSurgeon简介

ONNX-GraphSurgeon是一个Python库，用于操作ONNX计算图。它提供了丰富的API，支持对计算图进行增删改查等操作。以下是ONNX-GraphSurgeon的主要特点：

灵活性：可以轻松地修改计算图结构，如添加、删除、替换节点和边。
高效性：支持在计算图中进行层融合、模型剪枝等优化操作。
易用性：提供了简洁的API，便于开发者快速上手。

官方代码地址：
https://github.com/NVIDIA/TensorRT/tree/release/10.1/tools/onnx-graphsurgeon

三、安装ONNX-GraphSurgeon

在开始使用ONNX-GraphSurgeon之前，需要先安装以下依赖：

Python 3.6及以上版本

ONNX 1.6.0及以上版本

numpy

安装命令如下：

pip install onnx-graphsurgeon

四、对onnx输入端进行处理

1、onnx为啥需要剪切呢？

你以为的模型导出的onnx,

实际导出的onnx.

使用ONNX-GraphSurgeon 剪切后的onnx.

2、生成模型

建立一个模型：output = ReLU((A * X^T) + B) (.) C + D

#!/usr/bin/env python3
 
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
 
#
 
# Licensed under the Apache License, Version 2.0 (the "License");
 
# you may not use this file except in compliance with the License.
 
# You may obtain a copy of the License at
 
#
 
#     http://www.apache.org/licenses/LICENSE-2.0
 
#
 
# Unless required by applicable law or agreed to in writing, software
 
# distributed under the License is distributed on an "AS IS" BASIS,
 
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 
# See the License for the specific language governing permissions and
 
# limitations under the License.
 
#
 
 
import onnx_graphsurgeon as gs
 
import numpy as np
 
import onnx
 
 
print("Graph.layer Help:\n{}".format(gs.Graph.layer.__doc__))
 
 
# We can use `Graph.register()` to add a function to the Graph class. Later, we can invoke the function
 
# directly on instances of the graph, e.g., `graph.add(...)`
 
@gs.Graph.register()
 
def add(self, a, b):
 
    # The Graph.layer function creates a node, adds inputs and outputs to it, and finally adds it to the graph.
 
    # It returns the output tensors of the node to make it easy to chain.
 
    # The function will append an index to any strings provided for inputs/outputs prior
 
    # to using them to construct tensors. This will ensure that multiple calls to the layer() function
 
    # will generate distinct tensors. However, this does NOT guarantee that there will be no overlap with
 
    # other tensors in the graph. Hence, you should choose the prefixes to minimize the possibility of
 
    # collisions.
 
    return self.layer(op="Add", inputs=[a, b], outputs=["add_out_gs"])
 
 
 
@gs.Graph.register()
 
def mul(self, a, b):
 
    return self.layer(op="Mul", inputs=[a, b], outputs=["mul_out_gs"])
 
 
 
@gs.Graph.register()
 
def gemm(self, a, b, trans_a=False, trans_b=False):
 
    attrs = {"transA": int(trans_a), "transB": int(trans_b)}
 
    return self.layer(op="Gemm", inputs=[a, b], outputs=["gemm_out_gs"], attrs=attrs)
 
 
 
# You can also specify a set of opsets when regsitering a function.
 
# By default, the function is registered for all opsets lower than Graph.DEFAULT_OPSET
 
@gs.Graph.register(opsets=[11])
 
def relu(self, a):
 
    return self.layer(op="Relu", inputs=[a], outputs=["act_out_gs"])
 
 
 
# Note that the same function can be defined in different ways for different opsets.
 
# It will only be called if the Graph's opset matches one of the opsets for which the function is registered.
 
# Hence, for the opset 11 graph used in this example, the following function will never be used.
 
@gs.Graph.register(opsets=[1])
 
def relu(self, a):
 
    raise NotImplementedError("This function has not been implemented!")
 
 
 
##########################################################################################################
 
# The functions registered above greatly simplify the process of building the graph itself.
 
 
graph = gs.Graph(opset=11)
 
 
# Generates a graph which computes:
 
# output = ReLU((A * X^T) + B) (.) C + D
 
X = gs.Variable(name="X", shape=(64, 64), dtype=np.float32)
 
graph.inputs = [X]
 
 
# axt = (A * X^T)
 
# Note that we can use NumPy arrays directly (e.g. Tensor A),
 
# instead of Constants. These will automatically be converted to Constants.
 
A = np.ones(shape=(64, 64), dtype=np.float32)
 
axt = graph.gemm(A, X, trans_b=True)
 
 
# dense = ReLU(axt + B)
 
B = np.ones((64, 64), dtype=np.float32) * 0.5
 
dense = graph.relu(*graph.add(*axt, B))
 
 
# output = dense (.) C + D
 
# If a Tensor instance is provided (e.g. Tensor C), it will not be modified at all.
 
# If you prefer to set the exact names of tensors in the graph, you should
 
# construct tensors manually instead of passing strings or NumPy arrays.
 
C = gs.Constant(name="C", values=np.ones(shape=(64, 64), dtype=np.float32))
 
D = np.ones(shape=(64, 64), dtype=np.float32)
 
graph.outputs = graph.add(*graph.mul(*dense, C), D)
 
 
# Finally, we need to set the output datatype to make this a valid ONNX model.
 
#   In our case, all the data types are float32.
 
for out in graph.outputs:
 
    out.dtype = np.float32
 
 
onnx.save(gs.export_onnx(graph), "model.onnx")

3、在初始结点处增加操作

在Gemm操作前添加减均值，除均差的操作。

import onnx_graphsurgeon as gs
 
import numpy as np
 
import onnx
 
 
#增加减均值，除方差的操作 X^T) + B) (.) C + D
 
 
graph = gs.import_onnx(onnx.load("model.onnx"))
 
 
tamps = graph.tensors()
 
 
X = gs.Variable(name="X", shape=(64, 64), dtype=np.float32)
 
 
# 定义均值和方差
 
mean_value = np.array([0.5], dtype=np.float32)  # 替换 YOUR_MEAN_VALUE
 
std_value = np.array([0.2], dtype=np.float32)   # 替换 YOUR_STD_VALUE
 
 
mean = gs.Constant(name="mean", values=mean_value)
 
std = gs.Constant(name="std", values=std_value)
 
 
# 创建减均值和除方差的节点
 
sub_output = gs.Variable(name="X_minus_mean", shape=(64, 64), dtype=np.float32)
 
div_output = gs.Variable(name="X_normalized", shape=(64, 64), dtype=np.float32)
 
 
sub_node = gs.Node(op="Sub", inputs=[X, mean], outputs=[sub_output])
 
div_node = gs.Node(op="Div", inputs=[sub_output, std], outputs=[div_output])
 
 
# 将新创建的节点添加到图中
 
graph.nodes.extend([sub_node, div_node])
 
 
first_node = [node for node in graph.nodes if node.op == "Gemm"][0]
 
first_node.inputs[1] = div_node.outputs[0]
 
 
# 清理和顶排序
 
graph.cleanup().toposort()
 
 
onnx.save(gs.export_onnx(graph), "model_add.onnx")

4、修改结点的输入

将输入X改成Y。

import onnx_graphsurgeon as gs
 
import numpy as np
 
import onnx
 
 
# output = ReLU((A * X^T) + B) (.) C + D
 
 
graph = gs.import_onnx(onnx.load("model.onnx"))
 
 
tamps = graph.tensors()
 
 
# modify the input from X to Y
 
Y = gs.Variable(name="Y", shape=(64, 64), dtype=np.float32)
 
 
graph.inputs = [Y]
 
 
first_node = [node for node in graph.nodes if node.op == "Gemm"][0]
 
first_node.inputs[1] = Y
 
 
# 清理和顶排序
 
graph.cleanup().toposort()
 
 
onnx.save(gs.export_onnx(graph), "model_modify.onnx")

​

5、删除结点

删除Gemm操作

import onnx_graphsurgeon as gs
 
import numpy as np
 
import onnx
 
​
 
​
 
# output = ReLU((A * X^T) + B) (.) C + D
 
​
 
graph = gs.import_onnx(onnx.load("model.onnx"))
 
​
 
tamps = graph.tensors()
 
​
 
#delete Gemm
 
first_node = [node for node in graph.nodes if node.op == "Add"][0]
 
first_node.inputs[0] = tamps["X"]
 
​
 
# 清理和顶排序
 
graph.cleanup().toposort()
 
​
 
onnx.save(gs.export_onnx(graph), "model_delete.onnx")

总结：

ONNX GraphSurgeon 是一个强大的深度学习模型优化工具，它可以帮助我们提高模型的推理速度和资源利用率。通过合理地使用 ONNX GraphSurgeon，我们可以使深度学习模型在各种硬件平台上发挥出更好的性能。

欢迎关注我的公众号auto_driver_ai(Ai fighting), 第一时间获取更新内容。