加载器

模块:polygraphy.backend.onnx

class GsFromOnnx(model)[source]

基类: BaseLoader

创建 ONNX ModelProto 的 ONNX-GraphSurgeon 图的函子。

从 ONNX ModelProto 创建 ONNX-GraphSurgeon 图。

参数:

model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

call_impl()[source]
返回:

ONNX 模型的 ONNX-GraphSurgeon 表示

返回类型:

onnx_graphsurgeon.Graph

__call__(*args, **kwargs)

通过将参数转发到 call_impl 来调用加载器。

注意: call_impl 不应直接调用 - 请改用此函数。

gs_from_onnx(model)

GsFromOnnx 的立即求值函数变体。

从 ONNX ModelProto 创建 ONNX-GraphSurgeon 图。

参数:

model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

返回:

ONNX 模型的 ONNX-GraphSurgeon 表示

返回类型:

onnx_graphsurgeon.Graph

class OnnxFromPath(path, external_data_dir=None, ignore_external_data=None)[source]

基类: BaseLoader

从文件加载 ONNX 模型的函子。

从文件加载 ONNX 模型。

参数:
  • path (str) – 从中加载模型的路径。

  • external_data_dir (str) – 模型外部数据存储的目录。

  • ignore_external_data (bool) – 是否忽略任何外部数据,仅加载不带任何权重的模型结构。该模型仅可用于不需要权重的目的,例如提取子图或检查模型结构。这在外部数据不可用的情况下可能很有用。默认为 False。

call_impl()[source]
返回:

ONNX 模型

返回类型:

onnx.ModelProto

__call__(*args, **kwargs)

通过将参数转发到 call_impl 来调用加载器。

注意: call_impl 不应直接调用 - 请改用此函数。

onnx_from_path(path, external_data_dir=None, ignore_external_data=None)

OnnxFromPath 的立即求值函数变体。

从文件加载 ONNX 模型。

参数:
  • path (str) – 从中加载模型的路径。

  • external_data_dir (str) – 模型外部数据存储的目录。

  • ignore_external_data (bool) – 是否忽略任何外部数据,仅加载不带任何权重的模型结构。该模型仅可用于不需要权重的目的,例如提取子图或检查模型结构。这在外部数据不可用的情况下可能很有用。默认为 False。

返回:

ONNX 模型

返回类型:

onnx.ModelProto

class OnnxFromTfGraph(graph, opset=None, optimize=None)[source]

基类: BaseLoader

加载 TensorFlow 图并使用 tf2onnx 转换器将其转换为 ONNX 的函子。

将 TensorFlow 模型转换为 ONNX。

参数:
  • graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]) – 包含 TensorFlow 图和输出名称的元组,或返回该元组的 callable。

  • opset (int) – 转换期间要使用的 ONNX opset。

  • optimize (bool) – 是否使用 tf2onnx 的图优化传递。

call_impl()[source]
返回:

ONNX 模型。

返回类型:

onnx.ModelProto

__call__(*args, **kwargs)

通过将参数转发到 call_impl 来调用加载器。

注意: call_impl 不应直接调用 - 请改用此函数。

onnx_from_tf_graph(graph, opset=None, optimize=None)

OnnxFromTfGraph 的立即求值函数变体。

将 TensorFlow 模型转换为 ONNX。

参数:
  • graph (Union[Tuple[tf.Graph, Sequence[str]], Callable() -> Tuple[tf.Graph, Sequence[str]]]) – 包含 TensorFlow 图和输出名称的元组,或返回该元组的 callable。

  • opset (int) – 转换期间要使用的 ONNX opset。

  • optimize (bool) – 是否使用 tf2onnx 的图优化传递。

返回:

ONNX 模型。

返回类型:

onnx.ModelProto

class ModifyOutputs(model, outputs=None, exclude_outputs=None, copy=None)[source]

基类: BaseLoadOnnxCopy

修改 ONNX 模型输出的函子。

修改 ONNX 模型的输出。

参数:
  • model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

  • outputs (Sequence[str]) – 要标记为输出的张量名称。如果提供,这将覆盖现有的模型输出。如果使用 constants.MARK_ALL 值而不是列表,则网络中的所有张量都将被标记。

  • exclude_outputs (Sequence[str]) – 要排除作为输出的张量名称。这可以与 outputs=constants.MARK_ALL 结合使用以省略输出。

  • copy (bool) – 是否首先创建模型副本。默认为 False。

call_impl()[source]
返回:

具有修改输出的 ONNX 模型。

返回类型:

onnx.ModelProto

__call__(*args, **kwargs)

通过将参数转发到 call_impl 来调用加载器。

注意: call_impl 不应直接调用 - 请改用此函数。

modify_outputs(model, outputs=None, exclude_outputs=None, copy=None)

ModifyOutputs 的立即求值函数变体。

修改 ONNX 模型的输出。

参数:
  • model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

  • outputs (Sequence[str]) – 要标记为输出的张量名称。如果提供,这将覆盖现有的模型输出。如果使用 constants.MARK_ALL 值而不是列表,则网络中的所有张量都将被标记。

  • exclude_outputs (Sequence[str]) – 要排除作为输出的张量名称。这可以与 outputs=constants.MARK_ALL 结合使用以省略输出。

  • copy (bool) – 是否首先创建模型副本。默认为 False。

返回:

具有修改输出的 ONNX 模型。

返回类型:

onnx.ModelProto

class ConvertToFp16(model, copy=None)[source]

基类: BaseLoadOnnxCopy

将模型中所有浮点张量转换为 16 位精度的函子。使用 TensorRT 的 fp16 精度不需要这样做,但可能对其他后端有用。

将模型中所有浮点张量转换为 16 位精度。

参数:
  • model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

  • copy (bool) – 是否首先创建模型副本。默认为 False。

call_impl()[source]
返回:

修改后的 ONNX 模型。

返回类型:

onnx.ModelProto

__call__(*args, **kwargs)

通过将参数转发到 call_impl 来调用加载器。

注意: call_impl 不应直接调用 - 请改用此函数。

convert_to_fp16(model, copy=None)

ConvertToFp16 的立即求值函数变体。

将模型中所有浮点张量转换为 16 位精度。

参数:
  • model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

  • copy (bool) – 是否首先创建模型副本。默认为 False。

返回:

修改后的 ONNX 模型。

返回类型:

onnx.ModelProto

class FoldConstants(model, num_passes=None, do_shape_inference=None, partitioning=None, fold_shapes=None, copy=None, error_ok=None, size_threshold=None, allow_onnxruntime_shape_inference=None)[source]

基类: BaseLoadOnnxCopy

折叠 ONNX 模型中常量的函子。

折叠 ONNX 模型中的常量。

参数:
  • model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

  • num_passes (int) – 要运行的常量折叠传递次数。有时,计算张量形状的子图可能无法在单次传递中折叠。默认情况下,Polygraphy 将自动确定所需的传递次数。

  • do_shape_inference (bool) – 是否在传递之间在模型中运行形状推断。这使加载器能够折叠 Shape 节点。仅当 fold_shapes 为 True 时有效。默认为 True。

  • partitioning (Union[str, None]) –

    是否/如何对图进行分区,以便折叠模型一部分时发生的错误不会影响其他部分。可用模式为

    • None:不对图进行分区。如果推理失败,则不会折叠任何常量。

    • “basic”:对图进行分区。如果在一个分区中推理失败,其他分区将保持不受影响。

    • “recursive”:递归地对图进行分区。如果在一个分区中推理失败,则将进一步对该分区进行分区。

    默认为 None。

  • fold_shapes (bool) – 是否折叠图中的 Shape 节点。这需要在图中推断形状,并且只能折叠静态形状。默认为 True。

  • copy (bool) – 是否首先创建模型副本。默认为 False。

  • error_ok (bool) – 是否禁止在常量折叠期间发生错误。如果设置为 False,则将重新引发错误。默认为 True。

  • size_threshold (int) – 折叠常量的最大大小阈值(以字节为单位)。任何大于此值的张量都不会被折叠。设置为 None 以禁用大小阈值并始终折叠常量。例如,某些模型可能会对常量应用诸如 TileExpand 之类的操作,这可能会导致非常大的张量。与其预先计算这些常量并增大模型大小,不如跳过折叠它们并允许在运行时计算它们可能更可取。默认为 None。

  • allow_onnxruntime_shape_inference (bool) – 如果可用,是否允许使用 ONNX-Runtime 的形状推断来代替 ONNX 的形状推断实用程序。前者可以提供性能或内存使用方面的优势。如果 do_shape_inference 为 False,则无效。默认为 True。

call_impl()[source]
返回:

折叠常量的新 ONNX 模型。

返回类型:

onnx.ModelProto

__call__(*args, **kwargs)

通过将参数转发到 call_impl 来调用加载器。

注意: call_impl 不应直接调用 - 请改用此函数。

fold_constants(model, num_passes=None, do_shape_inference=None, partitioning=None, fold_shapes=None, copy=None, error_ok=None, size_threshold=None, allow_onnxruntime_shape_inference=None)

FoldConstants 的立即求值函数变体。

折叠 ONNX 模型中的常量。

参数:
  • model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

  • num_passes (int) – 要运行的常量折叠传递次数。有时,计算张量形状的子图可能无法在单次传递中折叠。默认情况下,Polygraphy 将自动确定所需的传递次数。

  • do_shape_inference (bool) – 是否在传递之间在模型中运行形状推断。这使加载器能够折叠 Shape 节点。仅当 fold_shapes 为 True 时有效。默认为 True。

  • partitioning (Union[str, None]) –

    是否/如何对图进行分区,以便折叠模型一部分时发生的错误不会影响其他部分。可用模式为

    • None:不对图进行分区。如果推理失败,则不会折叠任何常量。

    • “basic”:对图进行分区。如果在一个分区中推理失败,其他分区将保持不受影响。

    • “recursive”:递归地对图进行分区。如果在一个分区中推理失败,则将进一步对该分区进行分区。

    默认为 None。

  • fold_shapes (bool) – 是否折叠图中的 Shape 节点。这需要在图中推断形状,并且只能折叠静态形状。默认为 True。

  • copy (bool) – 是否首先创建模型副本。默认为 False。

  • error_ok (bool) – 是否禁止在常量折叠期间发生错误。如果设置为 False,则将重新引发错误。默认为 True。

  • size_threshold (int) – 折叠常量的最大大小阈值(以字节为单位)。任何大于此值的张量都不会被折叠。设置为 None 以禁用大小阈值并始终折叠常量。例如,某些模型可能会对常量应用诸如 TileExpand 之类的操作,这可能会导致非常大的张量。与其预先计算这些常量并增大模型大小,不如跳过折叠它们并允许在运行时计算它们可能更可取。默认为 None。

  • allow_onnxruntime_shape_inference (bool) – 如果可用,是否允许使用 ONNX-Runtime 的形状推断来代替 ONNX 的形状推断实用程序。前者可以提供性能或内存使用方面的优势。如果 do_shape_inference 为 False,则无效。默认为 True。

返回:

折叠常量的新 ONNX 模型。

返回类型:

onnx.ModelProto

class SetUpperBound(model, upper_bounds, copy=None)[source]

基类: BaseLoadOnnxCopy

为 ONNX 模型中具有无界 DDS 的张量设置上限的函子。

要求模型已常量折叠并已推断形状。

为 ONNX 模型中具有无界 DDS 的张量设置上限。

参数:
  • model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

  • upper_bounds (Union[int, Dict[str, int]]) – 具有无界 DDS 的张量的上限。如果提供单个整数,它将用作所有具有无界 DDS 的张量的默认上限。也可以使用字典按张量提供。在这种情况下,使用空字符串(“”)作为键来指定未显式列出的张量的默认上限。

  • copy (bool) – 是否首先创建模型副本。默认为 False。

call_impl()[source]
返回:

新的 ONNX 模型。

返回类型:

onnx.ModelProto

__call__(*args, **kwargs)

通过将参数转发到 call_impl 来调用加载器。

注意: call_impl 不应直接调用 - 请改用此函数。

set_upper_bound(model, upper_bounds, copy=None)

SetUpperBound 的立即求值函数变体。

为 ONNX 模型中具有无界 DDS 的张量设置上限。

参数:
  • model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

  • upper_bounds (Union[int, Dict[str, int]]) – 具有无界 DDS 的张量的上限。如果提供单个整数,它将用作所有具有无界 DDS 的张量的默认上限。也可以使用字典按张量提供。在这种情况下,使用空字符串(“”)作为键来指定未显式列出的张量的默认上限。

  • copy (bool) – 是否首先创建模型副本。默认为 False。

返回:

新的 ONNX 模型。

返回类型:

onnx.ModelProto

class InferShapes(model, error_ok=None, external_data_dir=None, save_to_disk_threshold_bytes=None, allow_onnxruntime=None)[source]

基类: BaseLoader

对 ONNX 模型运行形状推断的函子。

对 ONNX 模型运行形状推断。

参数:
  • model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto, str, Callable() -> str]) – ONNX 模型或返回 ONNX 模型的 callable,或模型路径。支持大于 2 GB protobuf 限制的模型。

  • error_ok (bool) – 是否应禁止形状推断期间的错误。默认为 True。

  • external_data_dir (str) – 模型外部数据存储的目录。仅当通过路径而不是加载器提供模型时使用。

  • save_to_disk_threshold_bytes (int) – ModelProto 将被序列化到磁盘然后运行形状推断的大小(以字节为单位)阈值。这可以用于解决 2 GB protobuf 限制。默认为 2 GB。

  • allow_onnxruntime (bool) – 如果可用,是否允许使用 ONNX-Runtime 的形状推断来代替 ONNX 的形状推断实用程序。前者可以提供性能或内存使用方面的优势。默认为 True。

call_impl()[source]
返回:

推断形状的新 ONNX 模型。

返回类型:

onnx.ModelProto

__call__(*args, **kwargs)

通过将参数转发到 call_impl 来调用加载器。

注意: call_impl 不应直接调用 - 请改用此函数。

infer_shapes(model, error_ok=None, external_data_dir=None, save_to_disk_threshold_bytes=None, allow_onnxruntime=None)

InferShapes 的立即求值函数变体。

对 ONNX 模型运行形状推断。

参数:
  • model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto, str, Callable() -> str]) – ONNX 模型或返回 ONNX 模型的 callable,或模型路径。支持大于 2 GB protobuf 限制的模型。

  • error_ok (bool) – 是否应禁止形状推断期间的错误。默认为 True。

  • external_data_dir (str) – 模型外部数据存储的目录。仅当通过路径而不是加载器提供模型时使用。

  • save_to_disk_threshold_bytes (int) – ModelProto 将被序列化到磁盘然后运行形状推断的大小(以字节为单位)阈值。这可以用于解决 2 GB protobuf 限制。默认为 2 GB。

  • allow_onnxruntime (bool) – 如果可用,是否允许使用 ONNX-Runtime 的形状推断来代替 ONNX 的形状推断实用程序。前者可以提供性能或内存使用方面的优势。默认为 True。

返回:

推断形状的新 ONNX 模型。

返回类型:

onnx.ModelProto

class ExtractSubgraph(model, input_metadata=None, output_metadata=None, check_meta=None)[source]

基类: BaseLoader

从 ONNX 模型提取子图的函子。

从 ONNX 模型提取子图。

参数:
  • model (Union[Union[onnx.ModelProto, onnx_graphsurgeon.Graph], Callable() -> Union[onnx.ModelProto, onnx_graphsurgeon.Graph]]) – ONNX 模型或 ONNX-GraphSurgeon Graph 或返回 ONNX 模型或 ONNX-GraphSurgeon Graph 的 callable。

  • input_metadata (TensorMetadata) – 子图输入的元数据。需要名称、形状和数据类型。如果未提供,则不会修改图输出。

  • output_metadata (TensorMetadata) – 子图输出的元数据。需要名称和数据类型。如果未提供,则不会修改图输出。

  • check_meta (bool) – 是否检查提供的输入和输出元数据是否包含所有预期的字段。默认为 True。

call_impl()[source]
返回:

新的 ONNX 模型或 ONNX-GraphSurgeon Graph。

返回类型:

Union[onnx.ModelProto, onnx_graphsurgeon.Graph]

__call__(*args, **kwargs)

通过将参数转发到 call_impl 来调用加载器。

注意: call_impl 不应直接调用 - 请改用此函数。

extract_subgraph(model, input_metadata=None, output_metadata=None, check_meta=None)

ExtractSubgraph 的立即求值函数变体。

从 ONNX 模型提取子图。

参数:
  • model (Union[Union[onnx.ModelProto, onnx_graphsurgeon.Graph], Callable() -> Union[onnx.ModelProto, onnx_graphsurgeon.Graph]]) – ONNX 模型或 ONNX-GraphSurgeon Graph 或返回 ONNX 模型或 ONNX-GraphSurgeon Graph 的 callable。

  • input_metadata (TensorMetadata) – 子图输入的元数据。需要名称、形状和数据类型。如果未提供,则不会修改图输出。

  • output_metadata (TensorMetadata) – 子图输出的元数据。需要名称和数据类型。如果未提供,则不会修改图输出。

  • check_meta (bool) – 是否检查提供的输入和输出元数据是否包含所有预期的字段。默认为 True。

返回:

新的 ONNX 模型或 ONNX-GraphSurgeon Graph。

返回类型:

Union[onnx.ModelProto, onnx_graphsurgeon.Graph]

class SaveOnnx(model, path, external_data_path=None, size_threshold=None, all_tensors_to_one_file=None)[source]

基类: BaseLoader

将 ONNX 模型保存到指定路径的函子。

将 ONNX 模型保存到指定路径。

参数:
  • model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

  • path (str) – 写入 ONNX 模型的路径。

  • external_data_path (str) – 保存外部数据的路径。这始终是相对路径;外部数据始终写入与模型相同的目录。设置为空字符串以使用默认路径。设置为 None 以禁用。如果模型在 protobuf 大小阈值内,则默认为 None,否则默认为空字符串。

  • size_threshold (int) – 张量大小阈值(以字节为单位),超过该阈值,张量数据将存储在外部文件中。小于此阈值的张量将保留在 ONNX 文件中。如果未设置 external_data_path,则无效。默认为 1024。

  • all_tensors_to_one_file (bool) – 保存外部数据时,是否将所有张量写入一个文件。如果未设置 external_data_path,则无效。默认为 True。

call_impl()[source]
返回:

模型,在保存之后。

返回类型:

onnx.ModelProto

__call__(*args, **kwargs)

通过将参数转发到 call_impl 来调用加载器。

注意: call_impl 不应直接调用 - 请改用此函数。

save_onnx(model, path, external_data_path=None, size_threshold=None, all_tensors_to_one_file=None)

SaveOnnx 的立即求值函数式变体。

将 ONNX 模型保存到指定路径。

参数:
  • model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

  • path (str) – 写入 ONNX 模型的路径。

  • external_data_path (str) – 保存外部数据的路径。这始终是相对路径;外部数据始终写入与模型相同的目录。设置为空字符串以使用默认路径。设置为 None 以禁用。如果模型在 protobuf 大小阈值内,则默认为 None,否则默认为空字符串。

  • size_threshold (int) – 张量大小阈值(以字节为单位),超过该阈值,张量数据将存储在外部文件中。小于此阈值的张量将保留在 ONNX 文件中。如果未设置 external_data_path,则无效。默认为 1024。

  • all_tensors_to_one_file (bool) – 保存外部数据时,是否将所有张量写入一个文件。如果未设置 external_data_path,则无效。默认为 True。

返回:

模型,在保存之后。

返回类型:

onnx.ModelProto

class BytesFromOnnx(model)[source]

基类: BaseLoader

序列化 ONNX 模型的函子。

序列化 ONNX 模型。

参数:

model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

call_impl()[source]
返回:

序列化后的模型。

返回类型:

bytes

__call__(*args, **kwargs)

通过将参数转发到 call_impl 来调用加载器。

注意: call_impl 不应直接调用 - 请改用此函数。

bytes_from_onnx(model)

BytesFromOnnx 的立即求值函数式变体。

序列化 ONNX 模型。

参数:

model (Union[onnx.ModelProto, Callable() -> onnx.ModelProto]) – ONNX 模型或返回 ONNX 模型的 callable。

返回:

序列化后的模型。

返回类型:

bytes

class OnnxFromBytes(serialized_onnx)[source]

基类: BaseLoader

反序列化 ONNX 模型的函子。

反序列化 ONNX 模型。

参数:

serialized_onnx (Union[bytes, Callable() -> bytes]) – 序列化的 ONNX 模型或返回模型的 callable 对象。

call_impl()[source]
返回:

ONNX 模型。

返回类型:

onnx.ModelProto

__call__(*args, **kwargs)

通过将参数转发到 call_impl 来调用加载器。

注意: call_impl 不应直接调用 - 请改用此函数。

onnx_from_bytes(serialized_onnx)

OnnxFromBytes 的立即求值函数式变体。

反序列化 ONNX 模型。

参数:

serialized_onnx (Union[bytes, Callable() -> bytes]) – 序列化的 ONNX 模型或返回模型的 callable 对象。

返回:

ONNX 模型。

返回类型:

onnx.ModelProto