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, max_shard_size_bytes=None, callback=None)[source]¶

Convert all initializers equal or above size_threshold_bytes to external tensors in-place and save data to one or more data files.

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.

When max_shard_size_bytes is set, tensors are distributed across multiple shard files named like model-00001-of-00003.data. Because each ONNX tensor already carries its own location, offset, and length fields, no separate index file is required — the ONNX proto itself encodes the routing.

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”. When sharding is enabled this becomes the base name used to generate shard filenames such as “model-00001-of-00003.data”.

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

  • max_shard_size_bytes (int | None) – Maximum cumulative size in bytes for a single shard file. When None (the default) all tensors are written to a single file given by relative_path. When set, tensors are written to multiple numbered shard files. If a single tensor is larger than this value, that tensor is written in its own oversized shard.

  • 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. Under sharding the callback index reflects each shard’s size-sorted write order while remaining globally contiguous.

Returns:

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

Raises:

  • ValueError – If max_shard_size_bytes is not greater than 0.

  • FileExistsError – When max_shard_size_bytes is set and any destination shard file already exists on disk. The sharded write path never overwrites existing files (re-saving a model whose shards already exist therefore requires deleting them first or choosing a different external data path). The single-file path (max_shard_size_bytes is None) instead overwrites relative_path unconditionally.

Return type:

Model

Notes

Stale shards from a previous save (when the new layout produces fewer or differently named shard files) are the caller’s responsibility to clean up. This function will neither delete nor rename pre-existing files that are not in the new destination set.

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 (including all of its subgraphs).

Parameters:

  • graph (Graph) – 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