解剖刀Onnx-GraphSurgeon：对onnx模型的末端进行增，删，改操作（三）

自动驾驶实战(AIFighting)

于 2024-07-20 22:48:47 发布

阅读量363

点赞数 3

文章标签：自动驾驶

本文链接：https://blog.csdn.net/laukal/article/details/140578796

版权

一、引言

ONNX-GraphSurgeon 是一个专为 ONNX 模型设计的 Python 工具包，它的主要功能是提供对 ONNX 计算图的高级编辑能力。这使得开发者能够根据具体的应用场景，对模型进行定制化修改与性能优化。

特点概述：

高度可塑性 - 开发者可以利用 ONNX-GraphSurgeon 轻松调整计算图的结构，包括但不限于节点和边的增加、移除或替换，从而实现模型的个性化定制。
卓越的效率 - 该工具包内置了多种优化机制，比如层融合和模型剪枝，能够显著提升模型运行时的效率和资源利用率。
用户友好 - 提供了一套直观且强大的 API 接口，使得开发者即使没有深入的底层知识也能迅速掌握并使用。

代码地址：

https://github.com/NVIDIA/TensorRT/tree/release/10.1/tools/onnx-graphsurgeon

二、Onnx-graphsurgeon 安装

安装命令如下：

pip install onnx-graphsurgeon

三、对onnx输入端进行处理

1、onnx为啥需要剪切呢？

你以为的模型导出的onnx,

实际导出的onnx.

使用ONNX-GraphSurgeon 剪切后的onnx.

2、生成模型

建立一个模型。

import onnx_graphsurgeon as gs
import numpy as np
import onnx

# Register functions to make graph generation easier
@gs.Graph.register()
def min(self, *args):
    return self.layer(op="Min", inputs=args, outputs=["min_out"])[0]


@gs.Graph.register()
def max(self, *args):
    return self.layer(op="Max", inputs=args, outputs=["max_out"])[0]


@gs.Graph.register()
def identity(self, inp):
    return self.layer(op="Identity", inputs=[inp], outputs=["identity_out"])[0]


# Generate the graph
graph = gs.Graph()

graph.inputs = [gs.Variable("input", shape=(4, 4), dtype=np.float32)]

# Clip values to [0, 6]
MIN_VAL = np.array(0, np.float32)
MAX_VAL = np.array(6, np.float32)

# Add identity nodes to make the graph structure a bit more interesting
inp = graph.identity(graph.inputs[0])
max_out = graph.max(graph.min(inp, MAX_VAL), MIN_VAL)
graph.outputs = [
    graph.identity(max_out),
]

# Graph outputs must include dtype information
graph.outputs[0].to_variable(dtype=np.float32, shape=(4, 4))

onnx.save(gs.export_onnx(graph), "model.onnx")

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

在Identity结点后面添加Mul结点

import onnx_graphsurgeon as gs
import onnx
import numpy as np

# 加载现有的ONNX模型
graph = gs.import_onnx(onnx.load("model.onnx"))
tmap = graph.tensors()

# 找到第二个 Identity 节点
identity_nodes = [node for node in graph.nodes if node.op == "Identity"]
second_identity_node = identity_nodes[1]

# 创建新的 Mul 节点
shape = (4, 4)
a = gs.Constant("a", values=np.ones(shape=shape, dtype=np.float32))
mul_out = gs.Variable(name="mul_out", dtype=np.float32, shape=shape)
mul_node = gs.Node(op="Mul", inputs=[second_identity_node.outputs[0], a], outputs=[mul_out])

# 更新图中引用第二个 Identity 节点输出的节点
for output in graph.outputs:
    if output == second_identity_node.outputs[0]:
        graph.outputs.remove(output)
        graph.outputs.append(mul_out)
        break

# 将新的 Mul 节点添加到图中
graph.nodes.append(mul_node)

# 清理和顶排序
graph.cleanup().toposort()

# 保存修改后的ONNX模型
onnx.save(gs.export_onnx(graph), "model_add.onnx")

4、修改结点的输入

将identity结点改成Mul结点

import onnx_graphsurgeon as gs
import onnx
import numpy as np

# 加载现有的ONNX模型
graph = gs.import_onnx(onnx.load("model.onnx"))
tmap = graph.tensors()

# 找到第二个 Identity 节点
identity_nodes = [node for node in graph.nodes if node.op == "Identity"][1]

identity_nodes.op= "Mul"

graph.cleanup()

# 保存修改后的ONNX模型
onnx.save(gs.export_onnx(graph), "model_modify.onnx")

5、删除结点

将前面添加的Mul结点删除

import onnx_graphsurgeon as gs
import onnx
import numpy as np

# 加载现有的ONNX模型
graph = gs.import_onnx(onnx.load("model_add.onnx"))
tmap = graph.tensors()

identity_nodes = [node for node in graph.nodes if node.op == "Identity"][1]

graph.outputs = [identity_nodes.outputs[0]]

# 清理和顶排序
graph.cleanup()

# 保存修改后的ONNX模型
onnx.save(gs.export_onnx(graph), "model_delete.onnx")