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 在不序列化原型的情况下无法操作。此函数使用 ModelContainer 的 Python API。

参数：

• graph – make_graph 返回的值

• large_initializers – 字典 name: 大型张量，大型张量是任何支持 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** – 初始值的形状

返回值：

创建的张量