onnx_ir.external_data

onnx_ir.external_data 模块提供了用于处理 ONNX 模型中外部数据的实用工具。它实现了张量与外部数据文件之间的转换,从而可以高效地存储和操作大型张量数据。这对于具有超出内存限制的大型初始值设定项的模型特别有用。

函数

onnx_ir.external_data.load_to_model(model)[source]

将所有外部模型初始值设定项就地转换为内存张量。

处理主图和子图中的所有初始值设定项。

参数:

model (Model) – 要处理的模型。

返回类型:

Model

onnx_ir.external_data.unload_from_model(model, base_dir, relative_path, *, size_threshold_bytes=0, callback=None)[source]

将所有等于或大于 size_threshold_bytes 的初始值设定项就地转换为外部张量,并将数据保存到单个数据文件。

它应该只用外部张量替换模型中的初始值设定项,而不对模型进行任何其他修改。

如果任何现有外部张量引用了提供的 external_data 路径,则在外部数据被覆盖后,它将失效。要获取有效的模型,请使用 load() 加载新保存的模型,或者提供一个模型中任何张量当前未引用的不同外部数据路径。

处理主图和子图中的所有初始值设定项。

参数:
  • model (Model) – 要处理的模型。

  • base_dir (str | PathLike) – ONNX 模型文件所在的目录路径。

  • relative_path (str | PathLike) – 外部数据存储的路径,相对于 ONNX 文件。例如“model.data”

  • size_threshold_bytes (int) – 如果张量大小(以字节为单位)大于此阈值,则保存到外部数据。

  • callback (Callable[[TensorProtocol, CallbackInfo], None] | None) – 一个回调函数,为每个保存到外部数据以进行调试或日志记录的张量调用。

返回:

一个 ir.Model,其中所有等于或大于 size_threshold_bytes 的初始值设定项数据都转换为外部张量。

返回类型:

Model

onnx_ir.external_data.convert_tensors_to_external(tensors, base_dir, relative_path, callback=None)[source]

将 TensorProtocol 张量序列转换为外部张量。

如果现有外部张量引用与目标路径相同的文件路径,则将其加载到内存中。

参数:
  • tensors (Sequence[TensorProtocol]) – 要转换为外部张量的张量。它们本身可以是外部张量。

  • base_dir (str | PathLike) – 基目录的路径。

  • relative_path (str | PathLike) – 外部数据存储的路径,相对于 ONNX 文件。

  • callback (Callable[[TensorProtocol, CallbackInfo], None] | None) – 一个回调函数,为每个保存到外部数据以进行调试或日志记录的张量调用。

返回:

从输入张量列表派生出的外部张量列表。顺序应与输入张量顺序匹配。

返回类型:

list[ExternalTensor]

onnx_ir.external_data.convert_tensors_from_external(tensors)[source]

将外部张量序列转换为内存张量。

参数:

tensors (Sequence[TensorProtocol]) – 要转换为内存张量的外部张量。

返回:

从外部张量列表派生出的内存张量列表。

返回类型:

list[TensorProtocol]

onnx_ir.external_data.set_base_dir(graph, base_dir)[source]

设置图中外部数据的基目录。

参数:
  • graph (Graph | GraphView) – 要遍历张量的图。

  • base_dir (str | PathLike) – 基目录。这是 ONNX 文件所在的目录。

返回类型: