加载器
模块: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__(*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__(*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__(*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__(*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__(*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
以禁用大小阈值并始终折叠常量。例如,某些模型可能会对常量应用诸如 Tile 或 Expand 之类的操作,这可能会导致非常大的张量。与其预先计算这些常量并增大模型大小,不如跳过折叠它们并允许在运行时计算它们可能更可取。默认为 None。allow_onnxruntime_shape_inference (bool) – 如果可用,是否允许使用 ONNX-Runtime 的形状推断来代替 ONNX 的形状推断实用程序。前者可以提供性能或内存使用方面的优势。如果
do_shape_inference
为 False,则无效。默认为 True。
- __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
以禁用大小阈值并始终折叠常量。例如,某些模型可能会对常量应用诸如 Tile 或 Expand 之类的操作,这可能会导致非常大的张量。与其预先计算这些常量并增大模型大小,不如跳过折叠它们并允许在运行时计算它们可能更可取。默认为 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__(*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__(*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__(*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__(*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__(*args, **kwargs)
通过将参数转发到
call_impl
来调用加载器。注意:
call_impl
不应直接调用 - 请改用此函数。
- onnx_from_bytes(serialized_onnx)
OnnxFromBytes
的立即求值函数式变体。反序列化 ONNX 模型。
- 参数:
serialized_onnx (Union[bytes, Callable() -> bytes]) – 序列化的 ONNX 模型或返回模型的 callable 对象。
- 返回:
ONNX 模型。
- 返回类型:
onnx.ModelProto