onnx_ir.external_data

The onnx_ir.external_data module provides utilities for handling external data in ONNX models. It enables the conversion of tensors to and from external data files, allowing for efficient storage and manipulation of large tensor data. This is particularly useful for models with large initializers that exceed memory constraints.

Functions

onnx_ir.external_data.load_to_model(model)[source]

Convert all external model initializers to memory tensors in-place.

All initializers in the main graph and subgraphs are handled.

Parameters:

model (Model) – Model to process.

Return type:

Model

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

Convert all initializers equal or above size_threshold_bytes to external tensors in-place and save data to a single data file.

It should only replace the initializers in the model with external tensors and not make any other modifications to the model.

If any existing external tensor references the provided external_data path, it will be invalidated after the external data is overwritten. To obtain a valid model, use load() to load the newly saved model, or provide a different external data path that is not currently referenced by any tensors in the model.

All initializers in the main graph and subgraphs are handled.

Parameters:

  • model (Model) – Model to process.

  • base_dir (str | PathLike) – Path the directory where the ONNX model file is.

  • relative_path (str | PathLike) – Path to which external data is to be stored, relative to the ONNX file. E.g. “model.data”

  • size_threshold_bytes (int) – Save to external data if the tensor size in bytes is larger than this threshold.

  • callback (Callable[[TensorProtocol, CallbackInfo], None] | None) – A callback function that is called for each tensor that is saved to external data for debugging or logging purposes.

Returns:

An ir.Model with all initializer data equal or above size_threshold_bytes converted to external tensors.

Return type:

Model

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

Convert a sequence of any TensorProtocol tensors to external tensors.

Existing external tensors are loaded to memory if they are referring to the same file path as the destination path.

Parameters:

  • tensors (Sequence[TensorProtocol]) – Tensors to be converted to external tensors. They can be external tensors themselves.

  • base_dir (str | PathLike) – Path of base directory.

  • relative_path (str | PathLike) – Path to which external data is to be stored, relative to the ONNX file.

  • callback (Callable[[TensorProtocol, CallbackInfo], None] | None) – A callback function that is called for each tensor that is saved to external data for debugging or logging purposes.

Returns:

A list of external tensors derived from a list of input tensors. The order should match the input tensor order.

Return type:

list[ExternalTensor]

onnx_ir.external_data.convert_tensors_from_external(tensors)[source]

Convert a sequence of external tensors to in-memory tensors.

Parameters:

tensors (Sequence[TensorProtocol]) – External tensors to be converted to in-memory tensors.

Returns:

A list of in-memory tensors derived from a list of external tensors.

Return type:

list[TensorProtocol]

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

Set the base directory for external data in a graph.

Parameters:

  • graph (Graph | GraphView) – The graph to traverse tensors on.

  • base_dir (str | PathLike) – The base directory. This is the directory where the ONNX file is.

Return type:

None