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)

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

All initializers and tensor attributes in the main graph and subgraphs are handled.

Changed in version 1.2.0: Also materializes external TENSOR and TENSORS node attributes in the main graph and subgraphs. Tensor docstrings and metadata properties are preserved during conversion.

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, convert_attributes=False, max_shard_size_bytes=None, callback=None, max_workers=None, max_in_flight_bytes=1073741824, alignment=None, align_threshold=1048576)

Convert model tensors equal or above size_threshold_bytes to external tensors in-place and save their data.

Added in version 1.1.0: Added max_shard_size_bytes, max_workers, max_in_flight_bytes, alignment, and align_threshold.

Changed in version 1.1.0: Initializers are packed densely in declaration order by default. Single-file writes now replace the destination atomically. Added concurrent writing, sharding, and optional alignment controls.

Added in version 1.2.0: Added convert_attributes to optionally store tensors in node attributes as external data.

Changed in version 1.2.0: Tensors exactly at size_threshold_bytes are now externalized; string tensors remain embedded. Tensor docstrings and metadata properties are preserved during conversion. With convert_attributes=False, existing external attribute tensors stay external and are streamed when relocation or an overwrite requires it. External source and output paths are validated before writing. Preserved attribute references must also pass the containment policy in the destination directory.

It should only replace initializers and requested tensor attributes in the model with external tensors and not make any other modifications to the model. When attribute conversion is disabled, existing external attributes stay external. Attributes needing relocation or backed by an overwritten destination are streamed into the new external data file regardless of the size threshold.

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. Tensor attributes are also handled when convert_attributes is true.

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 equal to or larger than this threshold.

• convert_attributes (bool) – Whether to store tensors in node attributes as external data according to size_threshold_bytes, in addition to initializers. When true, non-string attribute tensors at or above the threshold are stored as external tensors; smaller tensors are embedded in the model. When false, existing external attributes remain external and are streamed to the destination if relocation or an overwrite requires it. In-memory attributes remain embedded.

• 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 write order while remaining globally contiguous. When max_workers enables concurrency the callback is serialized with a lock but is no longer invoked in index order.

• max_workers (int | None) – Number of threads used to materialize and write tensors. None (the default) or 1 writes serially, preserving the previous behavior. Values above 1 overlap tensor materialization (dtype conversion, lazy tensor evaluation) with disk writes and parallelize both. Shards are also written concurrently.

• max_in_flight_bytes (int) – Upper bound on materialized tensor bytes held in memory when writing concurrently. Peak memory is bounded by this value plus the size of the largest single tensor. Shared across shards so that peak memory does not grow with the shard count.

• alignment (int | None) – Alignment to apply to the offsets of large tensors, in bytes. None (the default) packs tensors densely with no padding. When set, offsets are aligned to max(4096, alignment); 65536 matches the Windows allocation granularity used for memory mapping.

• align_threshold (int) – Only tensors strictly larger than this many bytes are aligned. Ignored when alignment is None.

Returns:

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

Raises:

• ValueError – If an external data path violates the containment policy.

• ValueError – If a numeric write option is outside its supported range.

• 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) atomically replaces relative_path only after the new file is complete.

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, max_workers=None, max_in_flight_bytes=1073741824, alignment=None, align_threshold=1048576, _budget=None, _tensor_write_locks=None)

Convert a sequence of any TensorProtocol tensors to external tensors.

Added in version 1.1.0: Added max_workers, max_in_flight_bytes, alignment, and align_threshold.

Changed in version 1.1.0: Tensors are packed densely in input order by default, and writes now atomically replace the destination after succeeding. Added concurrent writing and optional alignment controls.

Changed in version 1.2.0: Tensor docstrings and metadata properties are preserved during conversion.

Data is written to a temporary file in the destination directory and atomically replaces the destination only after the write succeeds. Existing external tensors backed by the destination stream directly into the temporary file and are invalidated after replacement.

Tensors are written in the order they are given, packed densely unless alignment is set. Preserving the input order matters: initializers are declared in topological order by every mainstream exporter, so the weights of one layer stay adjacent on disk. Runtimes memory-map each tensor separately and do not prefetch, so page faults are served in execution order and adjacency turns that into a sequential read.

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. When max_workers enables concurrency the callback is serialized with a lock but is no longer invoked in index order.

• max_workers (int | None) – Number of threads used to materialize and write tensors. None or 1 writes serially.

• max_in_flight_bytes (int) – Upper bound on materialized tensor bytes held in memory when writing concurrently.

• alignment (int | None) – Alignment to apply to the offsets of large tensors, in bytes. None (the default) packs tensors densely with no padding. When set, offsets are aligned to max(4096, alignment); 65536 matches the Windows allocation granularity used for memory mapping.

• align_threshold (int) – Only tensors strictly larger than this many bytes are aligned. Ignored when alignment is None.

• _budget (_ByteBudget | None)

• _tensor_write_locks (dict[int, allocate_lock] | None)

Returns:

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

Raises:

ValueError – If a numeric write option is outside its supported range.

Return type:

list[ExternalTensor]

onnx_ir.external_data.convert_tensors_from_external(tensors)

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

Changed in version 1.2.0: Tensor docstrings and metadata properties are preserved during conversion.

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)

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

See Model I/O and external data workflows for saving, sharding, callbacks, and safetensors, and Performance and large-model workflows for memory guidance.