TensorRT – onnx_graphsurgeon工具库的安装与API简介

1 onnx_grapgsurgeon

1.1 onnx_grapgsurgeon简介

onnx_grapgsurgeon是NVIDIA提供的创建和修改onnx模型的便捷易用工具库。

Github：https://github.com/NVIDIA/TensorRT/tree/master/tools/onnx-graphsurgeon

官方文档：https://docs.nvidia.com/deeplearning/tensorrt/onnx-graphsurgeon/docs/index.html

1.2 onnx_grapgsurgeon安装

参考官方Github：https://github.com/NVIDIA/TensorRT/tree/master/tools/onnx-graphsurgeon

使用以下命令安装

pip install onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com

如遇到以下提示：

ERROR: Could not find a version that satisfies the requirement protobuf<=3.20.1,>=3.12.2 (from onnx) (from versions: none)
ERROR: No matching distribution found for protobuf<=3.20.1,>=3.12.2

则使用以下命令安装指定版本的protobuf包

pip install protobuf==3.20.1

1.3 onnx_grapgsurgeon的API

1.3.1 导出模型

函数原型

onnx_graphsurgeon.export_onnx(graph: onnx_graphsurgeon.ir.graph.Graph, do_type_check=True, **kwargs)

函数功能

将onnx-graphsurgeon导出为onnx模型

函数参数

• graph：需要导出的图
• do_type_check：是否检查输入和输出张量定义了数据类型，如果没有则导出失败
• kwargs：onnx.helper.make_mode的其他参数

函数返回值

相关的onnx模型

1.3.2 导入模型

函数原型

onnx_graphsurgeon.import_onnx(onnx_model: onnx.onnx_ml_pb2.ModelProto)

函数功能

从onnx模型中导入onnx-graphsurgeon所需要的graph

函数参数

• onnx_model：需导入的onnx模型

函数返回值

与导入的onnx模型相关的onnx_graphsurgeon图

1.3.3 Graph类

onnx_graphsurgeon.Graph类主要用来表示onnx模型，包含了其中的节点和tensor。

classonnx_graphsurgeon.Graph(nodes: Optional[Sequence[onnx_graphsurgeon.ir.node.Node]] = None, inputs: Optional[Sequence[onnx_graphsurgeon.ir.tensor.Tensor]] = None, outputs: Optional[Sequence[onnx_graphsurgeon.ir.tensor.Tensor]] = None, name=None, doc_string=None, opset=None, import_domains=None, producer_name: Optional[str] = None, producer_version: Optional[str] = None)

参数

• node：graph中的node列表
• inputs：graph的输入tenosr列表
• outputs：graph的输出tensor列表
• name：graph的name，默认为onnx_graphsurgeon_graph
• doc_string：graph的doc_string，默认为""
• opset：导出graph时使用的onnx opset
• producer_name：生成模型所使用的工具名称，默认为""
• producer_version：生成工具的版本，默认为""

1.3.3.1 register

函数原型

static register(opsets=None)

函数功能

在Graph类中为指定的操作集组注册函数，注册函数之后可以像正常的成员函数一样访问。

例如

@Graph.register()
def add(self, a, b):
    return self.layer(op="Add", inputs=[a, b], outputs=["add_out_gs"])

graph.add(a, b)

函数参数

• opset：为其注册函数的一组操作集。

1.3.3.2 node_ids()

返回为Graph中的节点，提供唯一整数ID的上下文管理器。

1.3.3.3 cleanup

函数原型

cleanup(remove_unused_node_outputs=False, recurse_subgraphs=True, remove_unused_graph_inputs=False)

函数功能

从graph中删除未使用的node和tensor。如果node和tensor对graph中任何输出都没有贡献，则该node或者tensor被认为是未使用的。

函数参数

• remove_unused_node_outputs：是否删除节点的未使用输出张量，默认为False
• recurse_subgraphs：是否递归清理子图，默认为False
• remove_unused_graph_inputs：是否删除未使用的图输入，默认为False

1.3.3.4 toposort

函数原型

toposort(recurse_subgraphs=True)

函数功能

对Graph进行拓扑排序

函数参数

• recurse_subgraphs：是否对子图进行递归拓扑排序，默认为False

1.3.3.5 tensors

函数原型

tensors(check_duplicates=False)

函数功能

通过遍历所有节点，创建图使用的所有张量的张量映射。张量保证与图中节点的顺序一致。因此，如果图形是拓扑排序的，张量映射也会是。

函数参数

• check_duplicates：是否检查重复的同名tensor，默认为False

函数返回值

返回张量名称到张量的映射

1.3.3.6 fold_constants

函数原型

fold_constants(fold_shapes=True, recurse_subgraphs=True, partitioning=None, error_ok=True, flatten_subgraphs=True, size_threshold=None, should_exclude_node=None)

函数功能

将常量折叠到Graph中的适当位置，调用此函数之前必须对Graph进行拓扑排序。折叠常量之后不会删除多余常量，可以使用cleanup消除这些多余的常量。

函数参数

• fold_shapes：是否折叠图形中的形状节点。这需要在图形中推断形状，并且只能折叠静态形状。默认为True。
• recurse_subgraphs：是否递归折叠子图中的常量。默认为True。
• partitioning：是否/如何划分图形，以便折叠模型的一部分时的不会影响其他部分，可选模式有：1）None：不要对图形进行分区。如果推断失败，则不折叠常数；2）：basic：对图形进行分区。如果一个分区中的推理失败，则其他分区将不受影响；3）recursive：递归地对图形进行奇偶校验。如果分区中的推理失败，则将进一步分区。
• error_ok：是否应抑制推断错误。当该值为False时，将重新引发推理过程中遇到的任何错误。默认为True。
• flatten_subgraphs：是否在可能的情况下展平子图。例如，具有恒定条件的If节点可以展平到父图形中。默认为True。
• size_threshold：要折叠常数的最大大小阈值（以字节为单位）。任何大于该值的张量都不会被折叠。设置为“无”可禁用大小阈值并始终折叠常量。
• should_exclude_node：一个可调用的函数，它从图形中接受Graph中的节点，并报告是否应将其排除在折叠之外。

1.3.3.7 layer

函数原型

layer(inputs=[], outputs=[], *args, **kwargs)

函数功能

创建一个节点，将其添加到此图中，并可选地创建其输入和输出张量。

输入和输出列表可以包括各种不同的类型：

• Tensor：所提供的任何张量将按原样用于所创建节点的输入/输出。
• str：如果提供了字符串，则此函数将使用该字符串生成一个新的张量以生成名称。它将在提供的字符串末尾附加一个索引，以避免重复的张量名称，但由于这不能保证名称是唯一的，因此您应该尽量确保提供的字符串尽可能唯一。为了避免重复名称的问题，您可以自己生成名称并提供Tensor。
• numpy.ndarray：如果提供了NumPy数组，则此函数将使用名称前缀生成常量张量：onnx_graphsurgeon_constant
• Union[List[Number],Tuple[Number]]：如果提供了数字列表或元组（int或float），则此函数将使用名称前缀“onnx_graphsurgeon_lst_Constant”生成常量张量。张量的值将是包含指定值的1D数组。数据类型将是np.float32或np.int64。

函数参数

• inputs：输入列表
• outputs：输出列表
• args/kwargs：直接传递给Node类的构造函数的参数

函数返回值

node的输出tensor

1.3.3.8 copy

函数原型

copy(tensor_map: Optional[collections.OrderedDict[str, onnx_graphsurgeon.ir.tensor.Tensor]] = None)

函数功能

复制graph。

这将复制图形中的所有节点和张量，但不会对权重或属性进行深度复制（graph属性除外，将使用其复制方法进行复制）。

函数参数

• tensor_map：张量名称到外部图中张量的映射。如果这是最外层的图，则应为“None”。

函数返回值

graph的复制图

1.3.4 Node类

Node表示Graph中的一个操作，并消耗零个或多个张量，并产生零个或更多个张量。

classonnx_graphsurgeon.Node(op: str, name: Optional[str] = None, attrs: Optional[Dict[str, object]] = None, inputs: Optional[List[onnx_graphsurgeon.ir.tensor.Tensor]] = None, outputs: Optional[List[onnx_graphsurgeon.ir.tensor.Tensor]] = None)

参数：

• op：此node执行的操作
• name：此node的名字
• attrs：将属性名称映射到其值的字典
• inputs：零或者多个输入tensor的列表
• outputs：零或者多个输出tensor的列表

1.3.4.1 i

函数原型

i(tensor_idx=0, producer_idx=0)

函数功能

用于获取该节点的一个输入张量的生产者节点。

例如

assert node.i() == node.inputs[0].inputs[0]
assert node.i(1, 2) == node.inputs[1].inputs[2]

函数参数

• tensor_idx：此节点的输入张量的索引。默认值为0。
• producer_idx：如果张量有多个生产者，则输入张量生产者的索引。默认值为0。

函数返回值

指定的生产者（输入）节点。

1.3.4.2 o

函数原型

o(consumer_idx=0, tensor_idx=0)

函数功能

用于获取该节点的一个输出张量的消费者节点。

例如

assert node.o() == node.outputs[0].outputs[0]
assert node.o(2, 1) == node.outputs[1].outputs[2]

函数参数

• consumer_idx：输入张量的使用者的索引。默认值为0。
• tensor_idx：如果节点有多个输出，则为此节点的输出张量的索引。默认值为0。

函数返回值

指定的消费者（输出）节点。

1.3.4.3 copy

函数原型

copy(inputs: Optional[List[onnx_graphsurgeon.ir.tensor.Tensor]] = None, outputs: Optional[List[onnx_graphsurgeon.ir.tensor.Tensor]] = None, tensor_map=None)

函数功能

生成此节点的浅拷贝副本，覆盖输入和输出信息。

1.3.5 Tensor类

Tensor为Graph中张量的抽象类。

1.3.5.1 is_empty

函数原型

is_empty()

函数功能

返回此张量在Graph中是否为空。

1.3.5.2 to_constant

函数原型

to_constant(values: numpy.ndarray, data_location: Optional[int] = None)

函数功能

修改此张量以将其转换为常量。

函数参数

• values：tensor中的值
• data_location：一个枚举值，指示存储张量数据的位置。通常为onnx.TensorProto.DataLocation。

1.3.5.3 to_variable

函数原型

to_variable(dtype: Optional[numpy.dtype] = None, shape: Sequence[Union[int, str]] = [])

函数功能

修改此张量以将其转换为变量。

函数参数

• dtype：tensor的dtype
• shape：tensor的shape

1.3.5.4 i

函数原型

i(tensor_idx=0, producer_idx=0)

函数功能

用于获取该张量的输入节点之一的输入张量。

示例

assert tensor.i() == tensor.inputs[0].inputs[0]
assert tensor.i(1, 2) == tensor.inputs[2].inputs[1]

函数参数

• tensor_idx：输入节点的输入张量的索引。默认值为0。
• producer_idx：如果张量具有多个生产者，则输入张量的生产者节点的索引。默认值为0。

1.3.5.5 o

函数原型

o(consumer_idx=0, tensor_idx=0)

函数功能

用于获取该张量输出节点之一的输出张量。

示例

assert tensor.o() == tensor.outputs[0].outputs[0]
assert tensor.o(2, 1) == tensor.outputs[2].outputs[1]

函数参数

• consumer_idx：输入张量的使用者的索引。默认值为0。
• tensor_idx：如果节点有多个输出，则节点输出张量的索引。默认值为0。

1.3.6 Variable类

表示值未知的tensor。

classonnx_graphsurgeon.Variable(name: str, dtype: Optional[numpy.dtype] = None, shape: Optional[Sequence[Union[int, str]]] = None)

参数

• name：tensor的名字
• dtype：tensor的dtype
• shape：tensor的shape

1.3.7 Constant类

表示值已知的tensor。

classonnx_graphsurgeon.Constant(name: str, values: Union[numpy.ndarray, onnx_graphsurgeon.ir.tensor.LazyValues], data_location: Optional[int] = None)

参数

• name：tensor的名字
• values：tensor中的值，使用NumPy数组的形式表示
• data_location：一个枚举值，指示存储张量数据的位置。通常为onnx.TensorProto.DataLocation。