saveΒΆ
- onnx_ir.save(model, path, format=None, external_data=None, size_threshold_bytes=256, max_shard_size_bytes=None, callback=None, max_workers=None, max_in_flight_bytes=1073741824, alignment=None, align_threshold=1048576)ΒΆ
Save an ONNX model to a file.
The model remains unchanged after the call. If any existing external tensor references the provided
external_datapath, it will be invalidated after the external data is overwritten. To obtain a valid model, useload()to load the newly saved model, or provide a different external data path that is not currently referenced by any tensors in the model.Added in version 1.1.0: Added the
max_workers,max_in_flight_bytes,alignment, andalign_thresholdparameters.Changed in version 1.1.0: External tensors are packed densely by default instead of aligning tensors larger than 1 MiB to 64 KiB offsets. Pass
alignment=65536to retain the previous layout.Single-file external data writes now use a same-filesystem temporary file and atomically replace the destination after a successful write. Failures leave an existing destination unchanged. Replacement preserves symlinks and file permissions, but creates a new inode, so other hardlinks keep the old file contents.
Tip
A simple progress bar can be implemented by passing a callback function as the following:
import onnx_ir as ir import tqdm with tqdm.tqdm() as pbar: total_set = False def callback(tensor: ir.TensorProtocol, metadata: ir.external_data.CallbackInfo) -> None: nonlocal total_set if not total_set: pbar.total = metadata.total total_set = True pbar.update() pbar.set_description(f"Saving {tensor.name} ({tensor.dtype}, {tensor.shape}) at offset {metadata.offset}") ir.save( ..., callback=callback, )
- Parameters:
model (Model) β The model to save.
path (str | PathLike) β The path to save the model to. E.g. βmodel.onnxβ.
format (str | None) β The format of the file (e.g.
protobuf,textproto,json, etc.). If None, the format is inferred from the file extension.external_data (str | PathLike | None) β The relative path to save external data to. When specified, all initializers in the model will be converted to external data and saved to the specified directory. If None, all tensors will be saved unmodified. That is, if a tensor in the model is already external, it will be saved with the same external information; if the tensor is not external, it will be serialized in the ONNX Proto message.
size_threshold_bytes (int) β Save to external data if the tensor size in bytes is larger than this threshold. Effective only when
external_datais set.max_shard_size_bytes (int | None) β Maximum cumulative size in bytes for a single external data shard file. When
None(the default) all external tensors are written to the single file given byexternal_data. When set, tensors are distributed across numbered shard files (e.g.model-00001-of-00003.data). Because the ONNX format storeslocation,offset, andlengthper tensor, no separate index file is created β the saved ONNX proto itself encodes which shard each tensor lives in. If a single tensor is larger than this value, it is written in its own oversized shard file. Effective only whenexternal_datais set.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_workersenables 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 write external data.
None(the default) or1writes serially. Values above 1 overlap tensor materialization (lazy tensor evaluation, dtype conversion) with disk writes and parallelize both, which is significantly faster for large models. Peak memory stays bounded regardless of the worker count. Effective only whenexternal_datais set.max_in_flight_bytes (int) β Upper bound, in bytes, on the total size of tensors held in memory at once while writing. This caps peak memory use. A tensor larger than this budget is still admitted on its own, so the effective peak is roughly this value plus the size of the largest tensor. Effective only when
external_dataandmax_workersare set.alignment (int | None) β Alignment to apply to the offsets of large tensors, in bytes.
None(the default) packs tensors densely with no padding, producing smaller files. When set, offsets are aligned tomax(4096, alignment); 65536 matches the Windows allocation granularity used for memory mapping. Effective only whenexternal_datais set.align_threshold (int) β Only tensors strictly larger than this many bytes are aligned. Ignored when
alignmentisNone.
- Raises:
ValueError β If the external data path is absolute or a numeric write option is outside its supported range.
ValueError β If
max_shard_size_bytesis set withoutexternal_data.FileExistsError β When
max_shard_size_bytesis set and any destination shard file already exists on disk. The sharded write path never overwrites existing files; delete the conflicting files or choose a different external data path to re-save. The single-file path (max_shard_size_bytes is None) atomically replacesexternal_dataonly after the new file is complete.
- Return type:
None