onnx.model_container

ModelContainer

class onnx.model_container.ModelContainer[源代码]

实现一个 API,用于将大型张量存储在主 ModelProto 之外,它避免了在定义模型时复制大型初始化器,并且这些初始化器永远不会通过 protobuf 进行序列化。在用户明确保存模型之前,不会有任何张量存储在磁盘上。

enumerate_graph_protos() Iterable[onnx.GraphProto][源代码]

枚举模型中的所有 GraphProto。

is_in_memory_external_initializer(name: str) bool[源代码]

判断一个初始化器名称是否为存储在内存中的外部初始化器。在这种情况下,该名称必须以“#”开头。

load(file_path: str, load_large_initializers: bool = True)[源代码]

加载大型模型。

参数:
  • file_path – 模型文件

  • load_large_initializers – 加载大型初始化器,如果未完成,模型将不完整,但可用于查看模型而不执行它,并且可以使用 _load_large_initializers() 方法稍后加载它们

save(file_path: str, all_tensors_to_one_file: bool = False) ModelProto[源代码]

保存大型模型。该函数返回一个 ModelProto,如果模型不需要任何修改,则返回当前模型;如果需要更改(例如为每个外部张量指定文件名),则返回修改后的副本。

参数:
  • file_path – 模型文件

  • all_tensors_to_one_file – 将所有大型张量保存在一个文件中,或者每个大型张量一个文件

返回:

保存的 ModelProto

set_large_initializers(large_initializers: dict[str, ndarray])[源代码]

添加所有大型张量(未存储在模型中)。

make_large_model

onnx.model_container.make_large_model(graph: GraphProto, large_initializers: dict[str, ndarray] | None = None, **kwargs: Any) ModelContainer[源代码]

构建一个 ModelContainer

protobuf 的 C API 和 Python API 必须通过序列化 protos 来操作。此函数使用 ModelContainer 的 Python API。

参数:
  • graphmake_graph 返回值

  • large_initializers – 字典 name: large tensor,大型张量是任何支持 DLPack 协议的 python 对象,张量的所有权被转移到 ModelContainer,该张量必须定义类似 numpy 张量的 tobytes 方法

  • **kwargs – 要添加到返回实例的任何属性

返回:

ModelContainer

make_large_tensor_proto

onnx.model_container.make_large_tensor_proto(location: str, tensor_name: str, tensor_type: int, shape: tuple[int, ...]) TensorProto[源代码]

创建一个外部张量。

参数:
  • location – 唯一标识符(不一定是路径)

  • tensor_name – 图中的张量名称

  • tensor_type – onnx 类型

  • shape – 初始化器的形状

返回:

创建的张量