ONNX GraphSurgeon
项目描述
ONNX GraphSurgeon
目录
简介
ONNX GraphSurgeon是一个允许您轻松生成新的ONNX图或修改现有图的工具。
安装
使用预构建的wheel
python3 -m pip install onnx_graphsurgeon --extra-index-url https://pypi.ngc.nvidia.com
从源码构建
使用make目标
make install
手动构建
- 构建wheel
make build
- 从存储库外手动安装wheel
python3 -m pip install onnx_graphsurgeon/dist/onnx_graphsurgeon-*-py2.py3-none-any.whl
示例
示例目录包含ONNX GraphSurgeon的几个常用案例示例。
提供的可视化是由Netron生成的。
理解基础知识
ONNX GraphSurgeon由三个主要组件组成:导入器、IR和导出器。
导入器
导入器用于将图导入到ONNX GraphSurgeon IR中。导入器接口定义在base_importer.py中。
ONNX GraphSurgeon还提供了高级导入器API,以便于使用。
graph = gs.import_onnx(onnx.load("model.onnx"))
IR
中间表示(IR)是所有图修改的地方。它也可以用于从头开始创建新图。IR涉及三个组件:张量、节点和图。
每个组件的几乎所有成员变量都可以自由修改。有关这些类的各种属性的详细信息,您可以在交互式shell中使用help(<class_or_instance>)
查看帮助输出,或在脚本中使用print(help(<class_or_instance>))
,其中<class_or_instance>
是ONNX GraphSurgeon类型,或该类型的实例。
张量
张量分为两个子类:Variable
和Constant
。
- 常量是一个值预先已知的张量,可以像NumPy数组一样检索和修改。注意:常量的
values
属性是按需加载的。如果未访问该属性,则不会将其作为NumPy数组加载。 - 变量是一个张量,其值在推理时未知,但可能包含有关数据类型和大小的信息。
张量的输入和输出始终是节点。
ResNet50的一个示例常量张量
>>> print(tensor)
Constant (gpu_0/res_conv1_bn_s_0)
[0.85369843 1.1515082 0.9152944 0.9577646 1.0663182 0.55629414
1.2009839 1.1912311 2.2619808 0.62263143 1.1149117 1.4921428
0.89566356 1.0358194 1.431092 1.5360111 1.25086 0.8706703
1.2564877 0.8524589 0.9436758 0.7507614 0.8945271 0.93587166
1.8422242 3.0609846 1.3124607 1.2158023 1.3937513 0.7857263
0.8928106 1.3042281 1.0153942 0.89356416 1.0052011 1.2964457
1.1117343 1.0669073 0.91343874 0.92906713 1.0465593 1.1261675
1.4551278 1.8252873 1.9678202 1.1031747 2.3236883 0.8831993
1.1133649 1.1654979 1.2705412 2.5578163 0.9504889 1.0441847
1.0620039 0.92997414 1.2119316 1.3101407 0.7091761 0.99814713
1.3404484 0.96389204 1.3435135 0.9236031 ]
ResNet50的一个示例变量张量
>>> print(tensor)
Variable (gpu_0/data_0): (shape=[1, 3, 224, 224], dtype=float32)
节点
Node
定义图中的操作。节点可以指定属性;属性值可以是任何Python原始类型,也可以是ONNX GraphSurgeon的Graph
或Tensor
。
节点的输入和输出始终是张量。
ResNet50的一个示例ReLU节点
>>> print(node)
(Relu)
Inputs: [Tensor (gpu_0/res_conv1_bn_1)]
Outputs: [Tensor (gpu_0/res_conv1_bn_2)]
在这种情况下,节点没有属性。否则,属性将显示为OrderedDict
。
关于修改输入和输出的说明
节点和张量的inputs
/outputs
成员具有特殊逻辑,当您更改时,将更新所有受影响的节点和张量的输入/输出。这意味着,例如,当您更改输入张量的输出时,无需更新节点的inputs
。
考虑以下节点
>>> print(node)
(Relu).
Inputs: [Tensor (gpu_0/res_conv1_bn_1)]
Outputs: [Tensor (gpu_0/res_conv1_bn_2)]
可以像这样访问输入张量
>>> tensor = node.inputs[0]
>>> print(tensor)
Tensor (gpu_0/res_conv1_bn_1)
>>> print(tensor.outputs)
[ (Relu).
Inputs: [Tensor (gpu_0/res_conv1_bn_1)]
Outputs: [Tensor (gpu_0/res_conv1_bn_2)]
如果我们从张量的输出中删除节点,这也会反映在节点输入中
>>> del tensor.outputs[0]
>>> print(tensor.outputs)
[]
>>> print(node)
(Relu).
Inputs: []
Outputs: [Tensor (gpu_0/res_conv1_bn_2)]
图
图包含零个或多个Node
和输入/输出Tensor
。
中间张量不会被显式跟踪,而是从图中包含的节点中检索。
图类公开了几个函数。这里列出了其中的一小部分
cleanup()
:删除图中的未使用节点和张量toposort()
:对图进行拓扑排序。tensors()
:返回一个映射张量名称到张量的Dict[str, Tensor]
,通过遍历图中的所有张量。这是一个O(N)
操作,因此对于大型图可能会很慢。
要查看完整的图API,可以在交互式Python外壳中查看help(onnx_graphsurgeon.Graph)
。
导出器
导出器用于将ONNX GraphSurgeon IR导出到ONNX或其他类型的图。导出器接口定义在base_exporter.py中。
ONNX GraphSurgeon还提供了高级导出器API,以简化使用。
onnx.save(gs.export_onnx(graph), "model.onnx")
高级
处理具有外部数据的模型
使用ONNX-GraphSurgeon与外部存储的数据模型几乎与使用不带外部数据的ONNX模型相同。有关如何加载此类模型的详细信息,请参阅官方ONNX文档。要将模型导入ONNX-GraphSurgeon,可以使用正常的import_onnx
函数。
在导出时,只需执行一个额外的步骤
-
像平常一样从ONNX-GraphSurgeon导出模型
model = gs.export_onnx(graph)
-
更新模型,使其将数据写入外部位置。如果未指定位置,则默认为与ONNX模型相同的目录
from onnx.external_data_helper import convert_model_to_external_data convert_model_to_external_data(model, location="model.data")
-
然后您可以像平常一样保存模型
onnx.save(model, "model.onnx")