前言
自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考
本次课程我们来学习课程第三章—TensorRT 基础入门，一起来学习 onnx-graph-surgeon
课程大纲可以看下面的思维导图
0. 简述
本小节目标：学习使用 onnx-surgeon，比较与 onnx.helper 的区别，学习快速修改 onnx 以及替换算子/创建算子的技巧
这节我们学习第三章节第七小节—onnx-graph-surgeon，之前我们是利用 onnx.helper 模块来创建或者修改 onnx，在 TensorRT/tools 下面其实提供了一个更方便的 onnx 工具包 onnx_graphsurgeon
这个小节我们就来学习 onnx_graphsurgeon，并学习它与 onnx.helper 的区别，学会将复杂的 onnx 拆分成子图
1. 执行一下我们的python程序
这个小节的案例主要是 3.5-onnxsurgeon，如下所示：
我们会结合开源的 swin-tiny.onnx 来学习使用 onnx_graphsurgeon，学会修改 LayerNormalization、Multi-Head-Attention 的子图
2. onnx-graph-surgeon
我们先来看下 onnx-graph-surgeon 是什么，onnx-graph-surgeon 是一个创建和修改 onnx 的工具，它可以在 TensorRT/tools 中安装，它的主要特性：
- 更加方便地添加/修改 onnx 节点
- 更加方便地修改子图
- 更加方便地替换算子
- 底层一般是用的 onnx.helper，但是给做了一些封装
我们可以使用如下指令安装：
python3 -m pip install onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com
比如我们可以利用它将 min-max 替换成 clip 算子：
也可以用它将一系列复杂的 LayerNorm 计算替换成一个 LN 算子：
Note：LayerNorm 计算替换成 LN 是 TensorRT-8.6 以前的做法，TensorRT-8.6 之后就直接开始支持 LayerNorm 了
下面是 ChatGPT 给出的关于 onnx_graphsurgeon 模块的用途：
onnx-graphsurgeon 是 NVIDIA 的 TensorRT 工具包中一个用于操作 ONNX 图的模块。它允许用户以更高层次的抽象来修改和优化 ONNX 模型的计算图。以下是 onnx-graphsurgeon 的主要功能和用途：
- 图结构修改：可以添加、删除或替换图中的节点和边，从而对计算图进行各种修改。
- 图优化：通过合并节点、删除冗余节点等方式来优化计算图，从而提升模型的推理性能。
- 节点和张量操作：能够轻松地操作节点和张量，包括改变节点类型、修改节点属性以及调整张量形状等。
- 子图提取与合并：支持从图中提取子图或将多个子图合并，从而灵活地管理复杂的模型。
- 模型转换：可以将修改后的计算图转换回 ONNX 格式，从而便于与其他工具和框架进行集成。
- 调试与验证：提供工具来调试和验证图的结构和节点，确保修改后的图仍然有效并且能够正确运行。
onnx-graphsurgeon 适用于需要对 ONNX 模型进行深度修改和优化的场景，特别是在使用 TensorRT 进行模型加速时非常有用。
3. onnx-surgeon vs onnx.helper
我们在进入案例代码之前先看下 onnx_graphsurgeon 和 onnx.helper 有什么不同：
onnx 中的 IR 表示：
- ModelProto：描述的是整个模型的信息
- GraphProto：描述的是整个网络的信息
- NodeProto：描述的是各个计算节点，比如 conv，linear
- TensorProto：描述的是 tensor 的信息，主要包括权重
- ValueInfoProto：描述的是 input/output 信息
- GraphProto：描述的是整个网络的信息
onnx_graphsurgeon（gs）中的 IR 表示:
- Tensor：有两种类型
- Variable：主要就是那些不到推理不知道的变量
- Constant：不用推理时，而在推理前就知道的变量
- Node：跟 onnx 中的 NodeProto 差不多
- Graph：跟 onnx 中的 GraphProto 差不多
gs 帮助我们隐藏了很多信息，例如 node 的属性以前使用 AttributeProto 保存，但是 gs 中统一用 dict 来保存
onnx-graphsurgeon 和 onnx.helper 模块在功能和使用场景上有明显的区别：(from ChatGPT)
onnx-graphsurgeon：
- 高层次抽象：提供更高层次的 API，便于用户进行复杂的图结构修改和优化。
- 节点和图的高级操作：支持更灵活的节点操作、图重构和优化。
- 图优化和调试：内置了许多用于图优化和调试的功能，特别适合与 TensorRT 集成。
- 设计目标：主要设计用于深度修改和优化 ONNX 图，以提高模型的性能和可操作性。
onnx.helper：
- 基本功能：提供了创建和操作 ONNX 图基本元素（如节点、模型、图、张量）的辅助函数。
- 低层次操作：API 更基础，主要用于构建和操作基本图结构。
- 模型构建：适合用来创建和初始化 ONNX 模型，定义节点和图的基本元素。
- 设计目标：旨在简化 ONNX 模型的创建和初始化过程，帮助用户快速构建基本的 ONNX 模型。
具体区别：
抽象层次：
onnx-graphsurgeon提供更高层次的操作接口，更适合于复杂的图优化和修改。onnx.helper更侧重于基础的创建和操作，适合于构建和初始化模型。
功能侧重：
onnx-graphsurgeon侧重于图的优化和高级操作，适合需要深度修改和优化的场景。onnx.helper主要用于定义节点、创建图和模型，侧重于基础功能。
使用原生的 onnx.helper 创建 onnx 的代码如下：
使用原生的 gs 创建 onnx 的代码如下：
可以看到使用 gs 来创建 onnx 比 onnx.helper 要简单不少
gs 还可以自定义一些函数去创建 onnx 使整个 onnx 的创建更加方便，如下图所示：
这个类似于 onnx 中 symbolic 符号函数来注册算子一样，我们完全可以自己创建一些算子在这里使用
此外 gs 还可以方便我们把整个网络中的一些子图给“挖”出来，以此来分析细节，一般配合 polygraphy 使用，去寻找量化掉精度严重的子图
Note：polygraphy 是 TensorRT用来分析模型的一个非常非常重要的一个软件，会在后面的章节详细展开讲
比如现在有一个 swin-transformer.onnx 如下图所示：
我们可以使用如下代码将其中的 LayerNorm 部分给挖出来：
也可以使用如下代码将其中的 MHSA 部分给挖出来：
最后 gs 中最重要的一个特点在于我们可以使用 gs 来替换算子或者创建算子，这个会直接跟后面的 TensorRT plugin 绑定，实现算子的加速或者不兼容算子的实现
比如使用如下代码替换 LayerNormal：
修改前的 onnx 如下图所示：
修改后的 onnx 如下图所示：
4. 案例代码分析
下面我们进入代码中看下，我们先看 gs_create_conv.py 案例，代码如下：
import onnx_graphsurgeon as gs
import numpy as np
import onnx
# onnx_graph_surgeon(gs)中的IR会有以下三种结构
# Tensor
# -- 有两种类型
# -- Variable: 主要就是那些不到推理不知道的变量
# -- Constant: 不用推理时，而在推理前就知道的变量
# Node
# -- 跟onnx中的NodeProto差不多
# Graph
# -- 跟onnx中的GraphProto差不多
def main() -> None:
input = gs.Variable(
name = "input0",
dtype = np.float32,
shape = (1, 3, 224, 224))
weight = gs.Constant(
name = "conv1.weight",
values = np.random.randn(5, 3, 3, 3))
bias = gs.Constant(
name = "conv1.bias",
values = np.random.randn(5))
output = gs.Variable(
name = "output0",
dtype = np.float32,
shape = (1, 5, 224, 224))
node = gs.Node(
op = "Conv",
inputs = [input, weight, bias],
outputs = [output],
attrs = {"pads":[1, 1, 1, 1]})
graph = gs.Graph(
nodes = [node],
inputs = [input],
outputs = [output])
model = gs.export_onnx(graph)
onnx.save(model, "../models/sample-conv.onnx")
# 使用onnx.helper创建一个最基本的ConvNet
# input (ch=3, h=64, w=64)
# |
# Conv (in_ch=3, out_ch=32, kernel=3, pads=1)
# |
# output (ch=5, h=64, w=64)
if __name__ == "__main__":
main()
这段代码展示了使用 onnx_graphsurgeon 库创建了一个简单的 ONNX 模型，包含一个卷积层，可以看到相比于之前的 onnx.helper 模块创建 ONNX 要简单不少
导出的 ONNX 如下图所示：
我们再看下一个案例 gs_create_complicate_graph.py，代码如下所示：
import onnx_graphsurgeon as gs
import numpy as np
import onnx
#####################在graph注册调用的函数########################
@gs.Graph.register()
def add(self, a, b):
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)
@gs.Graph.register()
def relu(self, a):
return self.layer(op="Relu", inputs=[a], outputs=["act_out_gs"])
#####################通过注册的函数进行创建网络########################
# input (64, 64)
# |
# gemm (constant tensor A(64, 32))
# |
# add (constant tensor B(64, 32))
# |
# relu
# |
# mul (constant tensor C(64, 32))
# |
# add (constant tensor D(64, 32))
# 初始化网络的opset
graph = gs.Graph(opset=12)
# 初始化网络需要用的参数
consA = gs.Constant(name="consA", values=np.random.randn(64, 32))
consB = gs.Constant(name="consB", values=np.random.randn(64, 32))
consC = gs.Constant(name="consC", values=np.random.randn(64, 32))
consD = gs.Constant(name="consD", values=np.random.randn(64, 32))
input0 = gs.Variable(name="input0", dtype=np.float32, shape=(64, 64))
# 设计网络架构
gemm0 = graph.gemm(input0, consA, trans_b=True)
relu0 = graph.relu(*graph.add(*gemm0, consB))
mul0 = graph.mul(*relu0, consC)
output0 = graph.add(*mul0, consD)
# 设置网络的输入输出
graph.inputs = [input0]
graph.outputs = output0
for out in graph.outputs:
out.dtype = np.float32
# 保存模型
onnx.save(gs.export_onnx(graph), "../models/sample-complicated-graph.onnx")
这段代码使用 onnx_graphsurgeon 创建了一个更复杂的 ONNX 模型，展示了如何通过自定义的注册函数简化图层操作。
导出的 ONNX 如下图所示：
我们再接着看下一个案例 gs_create_subgraph.py，代码如下所示：
import onnx_graphsurgeon as gs
import numpy as np
import onnx
def load_model(model : onnx.ModelProto):
graph = gs.import_onnx(model)
print(graph.inputs)
print(graph.outputs)
def main() -> None:
model = onnx.load("../models/swin_tiny_patch4_window7_224.onnx")
graph = gs.import_onnx(model)
tensors = graph.tensors()
# LayerNorm部分
print(tensors["/backbone/patch_embed/Transpose_output_0"]) # LN的输入: 1 x 3136 x 96
print(tensors["/backbone/patch_embed/norm/Div_output_0"]) # LN的输出: 1 x 3136 x 96
graph.inputs = [
tensors["/backbone/patch_embed/Transpose_output_0"].to_variable(dtype=np.float32, shape=(1, 3136, 96))]
graph.outputs = [
tensors["/backbone/patch_embed/norm/Div_output_0"].to_variable(dtype=np.float32, shape=(1, 3136, 96))]
graph.cleanup()
onnx.save(gs.export_onnx(graph), "../models/swin-subgraph-LN.onnx")
# MHSA部分
graph = gs.import_onnx(model)
tensors = graph.tensors()
print(tensors["/backbone/layers.0/blocks.0/Reshape_3_output_0"]) # MHSA输入matmul: 64 x 49 x 96
print(tensors["onnx::MatMul_2233"]) # MHSA输入matmul的权重: 96 x 288
print(tensors["onnx::MatMul_2249"]) # MHSA输出matmul的权重: 96 x 96
print(tensors["/backbone/layers.0/blocks.0/attn/proj/MatMul_output_0"]) # MHSA输出: 64 x 49 x 96
graph.inputs = [
tensors["/backbone/layers.0/blocks.0/Reshape_3_output_0"].to_variable(dtype=np.float32, shape=(64, 49, 96))]
graph.outputs = [
tensors["/backbone/layers.0/blocks.0/attn/proj/MatMul_output_0"].to_variable(dtype=np.float32, shape=(64, 49, 96))]
graph.cleanup()
onnx.save(gs.export_onnx(graph), "../models/swin-subgraph-MSHA.onnx")
# 我们想把swin中LayerNorm中的这一部分单独拿出来
if __name__ == "__main__":
main()
这段代码使用 onnx_graphsurgeon 提取 Swin Transformer 模型中的 LayerNorm 和 MHSA 子图，具体步骤包括：加载完整模型后，通过打印和选择相关张量，设置 LayerNorm 的输入和输出，然后清理图结构并保存为 swin-subgraph-LN.onnx；同样地，为 MHSA 设置输入和输出，清理图结构后保存为 swin-subgraph-MSHA.onnx，实现了对特定子图的精准提取和导出。
swin-transformer 是一个基于 Transformer 的分类器，我们会在第六章节跟大家详细讲解，这里我们重点并不是要去理解 swin-transformer 的架构，而是要看如何把 ONNX 中的某个部分给截取出来，这个案例我们需要把 LayerNorm 和 self-attention 部分给截取出来
输出如下图所示：
导出来的 LayerNorm 子图 ONNX 如下图所示：
导出来的 MSHA 子图 ONNX 如下图所示：
我们再看后面几个案例，先看 gs_replace_op.py，代码如下所示：
import onnx_graphsurgeon as gs
import numpy as np
import onnx
import onnxruntime
import torch
#####################在graph注册调用的函数########################
@gs.Graph.register()
def min(self, *args):
return self.layer(op="Min", inputs=args, outputs=["min_output"])
@gs.Graph.register()
def max(self, *args):
return self.layer(op="Max", inputs=args, outputs=["max_output"])
@gs.Graph.register()
def identity(self, a):
return self.layer(op="Identity", inputs=[a], outputs=["identity_output"])
@gs.Graph.register()
def clip(self, inputs, outputs):
return self.layer(op="Clip", inputs=inputs, outputs=outputs)
#####################通过注册的函数进行创建网络########################
# input (5, 5)
# |
# identity
# |
# min
# |
# max
# |
# identity
# |
# output (5, 5)
def create_onnx_graph():
# 初始化网络的opset
graph = gs.Graph(opset=12)
# 初始化网络需要用的参数
min_val = np.array(0, dtype=np.float32)
max_val = np.array(1, dtype=np.float32)
input0 = gs.Variable(name="input0", dtype=np.float32, shape=(5, 5))
# 设计网络架构
identity0 = graph.identity(input0)
min0 = graph.min(*identity0, max_val)
max0 = graph.max(*min0, min_val)
output0 = graph.identity(*max0)
# 设置网络的输入输出
graph.inputs = [input0]
graph.outputs = output0
# 设置网络的输出的数据类型
for out in graph.outputs:
out.dtype = np.float32
# 保存模型
model = gs.export_onnx(graph)
model.ir_version = 9 # 设置 IR 版本为 9
onnx.save(model, "../models/sample-minmax.onnx")
#####################通过注册的clip算子替换网络节点####################
# input (5, 5)
# |
# identity
# |
# clip
# |
# identity
# |
# output (5, 5)
def change_onnx_graph():
graph = gs.import_onnx(onnx.load_model('../models/sample-minmax.onnx'))
tensors = graph.tensors()
inputs = [tensors["identity_output_0"],
tensors["onnx_graphsurgeon_constant_5"],
tensors["onnx_graphsurgeon_constant_2"]]
outputs = [tensors["max_output_6"]]
# 因为要替换子网，所以需要把子网和周围的所有节点都断开联系
for item in inputs:
# print(item.outputs)
item.outputs.clear()
for item in outputs:
# print(item.inputs)
item.inputs.clear()
# 通过注册的clip，重新把断开的联系链接起来
graph.clip(inputs, outputs)
# 删除所有额外的节点
graph.cleanup()
model = gs.export_onnx(graph)
model.ir_version = 9 # 设置 IR 版本为 9
onnx.save(model, "../models/sample-minmax-to-clip.onnx")
#####################验证模型##########################################
def validate_onnx_graph(input, path):
sess = onnxruntime.InferenceSession(path)
output = sess.run(None, {'input0': input.numpy()})
print("input is \n", input)
print("output is \n", output)
def main() -> None:
input = torch.Tensor(5, 5).uniform_(-1, 1)
# 创建一个minmax的网络
create_onnx_graph()
# 通过onnxruntime确认导出onnx是否正确生成
print("\nBefore modification:")
validate_onnx_graph(input, "../models/sample-minmax.onnx")
# 将minmax网络修改成clip网络
change_onnx_graph()
# 确认网络修改的结构是否正确
print("\nAfter modification:")
validate_onnx_graph(input, "../models/sample-minmax-to-clip.onnx")
if __name__ == "__main__":
main()
这段代码使用 onnx_graphsurgeon 库来操作和修改 ONNX 模型。它定义了一个简单的神经网络架构，其中包括一个通过最小值和最大值节点的流程，然后将这个网络转换为一个使用裁剪 (clip) 操作的网络。
值得注意的是上述代码与源码有些区别，博主进行了简单修改，将 model 的 IR 版本设置成了 9，这是因为博主在执行的过程中遇到了如下问题：
错误信息指出 “Unsupported model IR version: 10, max supported IR version: 9”，意味着模型的 IR 版本是 10，但是当前 onnxruntime 环境只支持到 IR 版本 9
博主尝试使用如下指令来更新 onnxruntime：
pip install --upgrade onnxruntime
更新后的 onnxruntime 版本是 1.17.3，执行后仍然出现如上的问题，因此博主将代码中导出的 model 的 IR 版本设置为了 9，一切输出正常
执行后输出如下：
可以看到 clip 替换 min、max 算子的前后输出基本一致，说明替换没问题
导出的 min-max.onnx 模型如下图所示：
替换后的 clip.onnx 模型如下图所示：
我们再看案例 gs_replace_LN.py，代码如下所示：
import onnx_graphsurgeon as gs
import numpy as np
import onnx
import onnxsim
import onnxruntime
import torch
import torch.nn as nn
#####################在graph注册调用的函数########################
@gs.Graph.register()
def identity(self, inputs, outputs):
return self.layer(op="Identity", inputs=inputs, outputs=outputs)
@gs.Graph.register()
def layerNorm(self, inputs, outputs, axis, epsilon):
attrs = {'axis': np.int64(axis), 'epsilon': float(epsilon)}
return self.layer(op="LayerNormalization", inputs=inputs, outputs=outputs, attrs=attrs)
@gs.Graph.register()
def layerNorm_default(self, inputs, outputs):
return self.layer(op="LayerNormalization", inputs=inputs, outputs=outputs)
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1)
self.norm = nn.LayerNorm(3)
self.act = nn.ReLU()
def forward(self, x):
_, _, H, W = x.shape
L = H * W
x = self.conv1(x)
x = x.view(x.shape[0], x.shape[1], L).permute(0, 2, 1)
x = self.norm(x)
x = self.act(x)
return x
def export_onnx_graph():
input = torch.Tensor(1, 3, 5, 5).uniform_(-1, 1)
model = Model()
model.eval()
file = "../models/sample-ln-before.onnx"
torch.onnx.export(
model = model,
args = (input,),
f = file,
input_names = ["input0"],
output_names = ["output0"],
opset_version = 12)
print("\nFinished export {}".format(file))
model_onnx = onnx.load(file)
onnx.checker.check_model(model_onnx)
print(f"Simplifying with onnx-simplifier {onnxsim.__version__}...")
model_onnx, check = onnxsim.simplify(model_onnx)
assert check, "assert check failed"
onnx.save(model_onnx, file)
#####################通过注册的LN算子替换网络节点####################
# input (5, 5)
# |
# conv
# |
# reshape
# |
# layerNorm
# |
# relu
# |
# output (5, 5)
def change_onnx_graph():
graph = gs.import_onnx(onnx.load_model('../models/sample-ln-before.onnx'))
tensors = graph.tensors()
norm_scale = gs.Constant(name="norm.weight", values=np.ones(shape=[3], dtype=np.float32))
norm_bias = gs.Constant(name="norm.bias", values=np.zeros(shape=[3], dtype=np.float32))
inputs = [tensors["/Transpose_output_0"]]
outputs = [tensors["/norm/Div_output_0"]]
# 因为要替换子网，所以需要把子网和周围的所有节点都断开联系
for item in inputs:
item.outputs.clear()
for item in outputs:
item.inputs.clear()
# 为了迎合onnx中operator中的设计，这里把scale和bias给加上
inputs = [tensors["/Transpose_output_0"],
norm_scale,
norm_bias]
# 这个onnx中的epsilon，我们给加上。当然，我们也可以选择默认的值
epsilon = [tensors["/norm/Constant_1_output_0"]]
print(type(epsilon[0].values))
# 通过注册的LayerNorm，重新把断开的联系链接起来
graph.layerNorm(inputs, outputs, axis=-1, epsilon=epsilon[0].values)
# graph.identity(inputs, outputs)
# graph.layerNorm_default(inputs, outputs)
# 删除所有额外的节点
graph.cleanup()
model = gs.export_onnx(graph)
model.ir_version = 9
onnx.save(model, "../models/sample-ln-after.onnx")
#####################验证模型##########################################
def validate_onnx_graph(input, origin_path, modified_path):
sess_origin = onnxruntime.InferenceSession(origin_path)
output_origin = sess_origin.run(None, {'input0': input.numpy()})
sess_modify = onnxruntime.InferenceSession(modified_path)
output_modify = sess_modify.run(None, {'input0': input.numpy()})
print("input is \n", input)
print("output_before is \n", output_origin)
print("output_after is \n", output_modify)
def main() -> None:
input = torch.Tensor(1, 3, 5, 5).uniform_(-1, 1)
##从pytorch导出onnx(这里为了实验，将不支持LayerNorm的opset12为例导出)
export_onnx_graph()
##手动修改LayerNorm
change_onnx_graph()
##验证修改的onnx
validate_onnx_graph(
input,
"../models/sample-ln-before.onnx",
"../models/sample-ln-after.onnx")
if __name__ == "__main__":
main()
这段代码首先定义了一个简单的 PyTorch 模型并导出为 ONNX 格式。然后，通过 onnx_graphsurgeon 替换模型中的 LayerNormalization 节点，并保存修改后的 ONNX 模型。最后，使用 onnxruntime 运行并验证修改前后的 ONNX 模型，确保两者输出一致。
在 gs_replace_LN.py 脚本中，LayerNormalization 的替换过程包括调用自定义的 LayerNormalization 函数，具体步骤如下：(from ChatGPT)
- 1. 导入模型并获取张量：使用
onnx_graphsurgeon导入已导出的 ONNX 模型，并提取模型中的所有张量。这些张量将用于重新连接新的节点。 - 2. 创建 LayerNormalization 所需的常量：创建 scale 和 bias 常量，这些常量是 LayerNormalization 操作所必需的。这些常量被定义为固定值，并命名为
norm.weight和norm.bias。 - 3. 断开旧节点的连接：找到需要替换的节点的输入和输出张量，并清除它们之间的连接。这确保了旧节点不会影响新的 LayerNormalization 操作。
- 4. 获取 epsilon 常量：提取现有网络中 LayerNormalization 操作所需的 epsilon 常量，这个值在重新连接新节点时会使用。
- 5. 调用自定义的 LayerNormalization 函数：使用先前提取的输入张量、scale 常量、bias 常量和 epsilon 常量，调用自定义的
layerNorm函数。该函数通过注册的 LayerNormalization 操作重新连接节点。调用时传入的参数包括输入张量、输出张量、axis 和 epsilon 值。 - 6. 清理图并保存模型：清理图中不再需要的节点，确保模型结构的简洁和正确。最后，将修改后的图导出并保存为新的 ONNX 模型文件。
通过这些步骤，脚本实现了使用自定义的 LayerNormalization 函数替换模型中的旧节点，成功修改了模型的结构和功能。
值得注意的是上述代码与源码有所区别，博主进行了简单修改，将 17 行的 np.float 替换成了 float，这是因为博主在执行的过程中出现了如下问题：
这个错误是因为 np.float 是一个已经被弃用的别名，使用它会导致 AttributeError。从 NumPy 1.20 开始 np.float 被弃用，并推荐使用内置的 float 或 np.float64。
此外将模型导出的 IR 版本固定在 9 版本，不然会出现之前的 “Unsupported model IR version: 10, max supported IR version: 9” 错误
执行后输出如下：
可以看到 LayerNormalization 算子替换前后输出基本一致，说明替换没问题
导出的 sample-ln-before.onnx 模型如下图所示：
导出的 sample-ln-after.onnx 模型如下图所示：
其实 LayerNormalization 算子在 opset=17 的版本已经支持，因此我们完全可以在导出 ONNX 将 opset 版本设置为 17 及以上这样导出来的就是一个完整的 LayerNormalization，这里只是为了方便演示 onnx_graphsurgeon 替换算子的功能将 opset 设置成了 12
我们来看最后一个案例 gs_remove_node.py，代码如下所示：
import onnx_graphsurgeon as gs
import numpy as np
import onnx
def load_model(model : onnx.ModelProto):
graph = gs.import_onnx(model)
print(graph.inputs)
print(graph.outputs)
def main() -> None:
model = onnx.load("../models/example_two_head.onnx")
graph = gs.import_onnx(model)
tensors = graph.tensors()
print(tensors)
print(tensors["input0"])
print(tensors["output0"])
print(tensors["output1"])
print(tensors["onnx::MatMul_8"])
graph.inputs = [
tensors["input0"].to_variable(dtype=np.float32, shape=(1, 1, 1, 4))]
graph.outputs = [
tensors["output0"].to_variable(dtype=np.float32, shape=(1, 1, 1, 3))]
# tensors["onnx::MatMul_8"].to_variable(dtype=np.float32, shape=(1, 1, 1, 3))]
graph.cleanup()
onnx.save(gs.export_onnx(graph), "../models/example_two_head_removed.onnx")
if __name__ == "__main__":
main()
该脚本文件使用 onnx_graphsurgeon 对 ONNX 模型进行修改。具体实现的功能是：导入一个名为 example_two_head.onnx 的模型，打印模型的输入、输出和特定张量的信息，修改模型的输入输出张量形状并清理图中的多余节点，最终保存修改后的模型为 example_two_head_removed.onnx。
执行后输出如下图所示：
example_two_head.onnx 如下图所示：
example_two_head_removed.onnx 如下图所示：
总结
本次课程我们主要学习了利用 onnx-graph-surgeon 来操作 ONNX，包括添加、修改节点，修改子图，替换算子，相比于 onnx.helper 更加方便，后续我们通过各种案例学习了 onnx_graphsurgeon 如何使用
OK，以上就是第 7 小节有关 onnx-graph-surgeon 的全部内容了，下节我们来学习快速分析开源代码并导出 ONNX，敬请期待😄