ONNX GraphSurgeon

目录

export_onnx

import_onnx

onnx_graphsurgeon.Graph

register(opsets=None)

node_ids()

copy（）

toposort（）

tensors()

layer（）

Node

Variable

to_constant(values: numpy.ndarray)

Constant

---

在进行TensorRT推理模型的时候，有些网络TensorRT是没有进行实现的，那么就需要进行TensorRT插件的定义，在使用onnx模型进行导入的时候需要针对模型进行修改，故本文针对onnx修改的工具进行学习。

官网包括ONNX GraphSurgeon的Python API文档。ONNX GraphSurgeon提供了一种方便的方法来创建和修改ONNX模型。

export_onnx

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

将onnx-graphsurgeon Graph导出到ONNX模型。

参数
        graph（图形）-要导出的onnx modle
        do_type_check（bool）-是否检查输入和输出张量是否定义了数据类型，如果没有则失败。
        kwargs -onnx.helper.make_model的附加参数
返回
        相应的ONNX model。
返回类型
        onnx.ModelProto

import_onnx

onnx_graphsurgeon.import_onnx(onnx_model: onnx.onnx_ml_pb2.ModelProto)→ onnx_graphsurgeon.ir.graph.Graph

从提供的ONNX模型导入onnx-graphsurgeon Graph。
参数
        onnx_model（onnx.ModelProto）-ONNX模型。
返回
        对应的onnx-graphsurgeon Graph。
返回类型
        Graph

onnx_graphsurgeon.Graph

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)

表示包含节点和张量的图。

参数
        nodes（Sequence[Node]）-此图形中的节点列表。
        inputs（Sequence[Tensor]）-图形输入张量的列表。
        outputs（Sequence[Tensor]）-图形输出张量的列表。
        name（str）-图形的名称。默认为“onnx_graphsurgeon_graph”。
        doc_string（str）-图形的doc_string。默认为“"。
        opset（int）-导出此图形时使用的ONNX操作集。
        producer_name（str）-用于生成模型的工具的名称。默认为“"。
        producer_version（str）-生成工具的版本。默认为“"。

register(opsets=None)

Registers a function with the Graph class for the specified group of opsets. After registering the function, it can be accessed like a normal member function.

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

参数
        opsets（Sequence[int]）-要为其注册函数的一组opset。具有相同名称的多个函数可以被同时注册，如果它们被注册用于不同的opset。为相同的opset注册同名函数将覆盖先前为这些opset注册的任何函数。默认情况下，为所有opset注册该函数。

node_ids()

Returns a context manager that supplies unique integer IDs for Nodes in the Graph.

with graph.node_ids():
    assert graph.nodes[0].id != graph.nodes[1].id

返回
        为节点提供唯一整数ID的上下文管理器。
返回类型
        NodeIDAdder

cleanup( )

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

从图中删除未使用的节点和张量。如果节点或张量对任何图形输出都没有贡献，则将其视为未使用。

另外，从图中移除不在图中的图输入张量的任何生产者节点以及图输出张量的用户者节点。
注意：此函数永远不会修改图形输出张量。

参数
        remove_unused_node_outputs（bool）-是否删除节点的未使用输出张量。这将永远不会删除空张量（即可选，但省略）输出。默认为False。
        recurse_subgraphs（bool）-是否递归地清除子图。默认为True。
        remove_unused_graph_inputs（bool）-是否删除未使用的图形输入。默认为False。
返回
        self

copy（）

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

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

参数
tensor_map（OrderedDict[str，Tensor]）-将张量名称映射到外部图中的张量。如果这是最外面的图形，则应为“无”。
返回
        A copy of the graph.
返回类型
       graph

toposort（）

toposort(recurse_subgraphs=True)

在适当的位置对图进行拓扑排序。
参数
        recurse_subgraphs（bool）-是否对子图进行递归拓扑排序。默认为True。
返回
        self

tensors()

tensors(check_duplicates=False)

通过遍历所有节点，创建此图使用的所有张量的张量映射。空张量从该映射中省略。
张量保证按照图中节点的顺序。因此，如果图是拓扑排序的，张量映射也会是。

参数
        check_duplicates（bool）-如果遇到多个同名的张量，是否失败。
Raises
        OnnxGraphSurgeonException -如果check_duplicates为True并且图中的多个不同张量共享相同的名称。
返回
        张量名称到张量的映射。
返回类型
        OrderedDict[str，Tensor]

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)

在图形中就地折叠常数。在调用此函数之前，必须对图进行拓扑排序（参见toposort（））。
此函数在折叠常量后不会删除常量。为了摆脱这些挂起的节点，可以运行cleanup（）函数。
注：由于此功能的实现方式，图形必须可导出到ONNX，并可在ONNX-Runtime中评估。此外，还必须安装ONNX-Runtime。

参数
fold_shapes（bool）-是否折叠图形中的Shape节点。这需要在图形中推断形状，并且只能折叠静态形状。默认为True。
recurse_subgraphs（bool）-是否递归地折叠子图中的常数。默认为True。
partitioning（Union[str，None]）-
是否/如何划分图形，以便折叠模型的一部分时的错误不会影响其他部分。可用模式包括：
无：不对图形进行分区。如果推理失败，则不折叠任何常数。
“basic”：对图进行分区。如果在一个分区中推理失败，其他分区也会失败
不受影响。
“recursive”：递归地划分图。如果推理在分区中失败，则该分区
将被进一步瓜分。
默认为无。
error_ok（bool）-是否应该抑制推理错误。如果此值为False，则将重新引发推断期间遇到的任何错误。默认为True。
flatten_subgraphs（bool）-是否在可能的情况下展平子图。例如，如果具有恒定条件的节点可以被展平到父图中。
size_threshold（int）-折叠常量的最大大小阈值，以字节为单位。任何大于此值的张量都不会折叠。设置为“无”可禁用大小阈值，并始终禁用折叠常量。例如，一些模型可能会将Tile或Expand等操作应用于常量，这可能会导致非常大的张量。不是预先计算那些常数并膨胀模型大小，而是可能期望跳过折叠它们并允许它们在运行时被计算。默认为无。
should_exclude_node（Callable[[gs.Node]，bool]）-接受图中的onnx-graphsurgeon节点并报告是否应从折叠中排除该节点的可调用项。这仅对否则可折叠的节点调用。注意，防止节点被折叠也防止其消费者被折叠。默认为始终返回False的可调用项。

layer（）

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

创建一个节点，将其添加到此图形，并可以选择创建其输入和输出张量。
输入和输出列表可以包括各种不同的类型：

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

Node

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, 
                domain: Optional[str] = None)

节点表示图中的操作，并且消耗零个或多个张量，并且产生零个或多个张量。

参数
op（str）-此节点执行的操作。onnx/Operators.md at main · onnx/onnx · GitHub，支持的算子
name（str）-此节点的名称。
attrs（Dict[str，object]）-将属性名称映射到其值的字典。
inputs（List[Tensor]）-零个或多个输入张量的列表。
outputs（List[Tensor]）-零个或多个输出张量的列表。
domain（str）-此节点的域，

i(tensor_idx=0, producer_idx=0)

便利函数，用于获取该节点的输入张量之一的生产者节点。请注意，与o（）函数相比，参数是交换的;这是因为张量可能只有一个生产者
例子

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

参数
        tensor_idx（int）-该节点的输入张量的索引。默认为0。
        producer_idx（int）-如果张量有多个生产者，则输入张量的生产者的索引。默认为0
返回
        指定的生产者（输入）节点。
返回类型
        node

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（int）-输入张量的消费者的索引。默认为0。
        tensor_idx（int）-如果节点有多个输出，则该节点的输出张量的索引。默认为0。
返回
        指定的消费者（输出）节点
返回类型
        NODE

Variable

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

表示一个张量，其值直到推理时才知道。

参数
        name（str）-张量的名称。
        dtype（numpy.dtype）-张量的数据类型。
        shape（Sequence[Union[int，str]]）-张量的形状。如果模型使用标注参数，则该参数可能包含字符串。

to_constant(values: numpy.ndarray)

就地修改此张量以将其转换为常数。这意味着张量的所有消费者/生产者都将看到更新。
参数
        values（np.ndarray）-此张量中的值
        data_location（int）-指示张量数据存储位置的枚举值。通常，这将来自onnx. TensorProto. DataLocation。
返回
        self

Constant

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

Represents a Tensor whose value is known.

参数
name（str）-张量的名称。
values（numpy.ndarray）-此张量中的值，以NumPy数组的形式。
data_location（int）-指示张量数据存储位置的枚举值。通常，这将来自onnx. TensorProto. DataLocation。