save_safetensors

onnx_ir.save_safetensors(model, path, /, *, format=None, size_threshold_bytes=256, max_shard_size_bytes=None, callback=None, convert_attributes=False)

Save an ONNX model to a file with external data in a safetensors file.

The model object is unmodified after this operation, except that existing external tensors backed by overwritten files are invalidated.

When sharding is enabled, multiple safetensors files will be created with names like “model-00001-of-00003.safetensors”, and an index file “model.safetensors.index.json” will be created to map tensors to their respective shard files. The shards will be created only if the total size of tensors exceeds the specified max_shard_size_bytes.

Note

Because the safetensors data format uses key-value mapping to store tensors, all initializer names in the model (across subgraphs) must be unique. Tensor attributes use generated unique safetensors keys, independent of their ONNX tensor names. Initializer names are reserved before generating attribute keys, and original tensor names are preserved in the ONNX file.

Example:

import onnx_ir as ir

model = ir.load("model.onnx")

# Save model with tensors larger than 100 bytes to safetensors external data,
# sharding files larger than 5GB.
ir.save_safetensors(
    model,
    "model.onnx",
    size_threshold_bytes=100,
    max_shard_size_bytes=int(5 * 1000**3),  # Shard safetensors files larger than 5GB
)

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 {metadata.filename}: {tensor.name} ({tensor.dtype}, {tensor.shape})")

    ir.save_safetensors(
        ...,
        callback=callback,
    )

Added in version 0.1.15.

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

Changed in version 1.2.0: Existing external attributes are preserved or streamed when relocation or an overwrite requires it. Tensor names, docstrings, and metadata properties are preserved. External initializers below the size threshold are embedded in the model. Single-file writes replace destinations atomically; sharded writes refuse to overwrite existing files. Callback offsets identify the tensor’s physical byte offset within its shard, including the header. External source and output paths are validated before writing. Preserved attribute references must also pass the containment policy in the destination directory.

Parameters:

• model (Model) – ONNX model to save.

• path (str | PathLike) – Path to the ONNX model file. 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.

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

• max_shard_size_bytes (int | None) – Maximum size in bytes (as int) a safetensors file before being sharded. If None, no sharding is performed.

• callback (Callable[[TensorProtocol, CallbackInfo], None] | None) – A callback function invoked before each tensor is written. The callback must have signature Callable[[ir.TensorProtocol, ir.external_data.CallbackInfo], None], where the first argument is the tensor being saved and the second contains metadata such as filename, progress, and the byte offset within the shard.

• 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.

Raises:

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

• ValueError – If initializer names are missing, duplicated, or reserved by safetensors.

• ValueError – If max_shard_size_bytes is not positive.

• ValueError – If the ONNX model and safetensors data paths are the same.

• FileExistsError – If a destination shard or index file already exists when max_shard_size_bytes is set.

Return type:

None