Model I/O and external data workflows¶

This page covers user-facing I/O features implemented in onnx_ir._io, onnx_ir.external_data, and onnx_ir._safetensors.

Load and save a model¶

import onnx_ir as ir

model = ir.load("model.onnx")
# ... mutate model ...
ir.save(model, "updated.onnx")

Use normalization and validation when needed¶

The IR permits temporarily incomplete or invalid states so multi-step transformations can be expressed naturally. Prefer transformations that preserve names, topological order, and type/shape information when possible. Normalization and inference passes traverse the model, so do not add them to every save path by default.

See Preserve invariants and use targeted repair for the authoritative decision checklist.

Apply only the operations required by the preceding transformations:

import onnx_ir.passes.common as common_passes

# Only if a rewrite may have disturbed node order.
model = common_passes.TopologicalSortPass()(model).model

# Only if stale shapes/types are needed by a later stage.
model = common_passes.ShapeInferencePass()(model).model

# At an explicit validation boundary, if desired.
common_passes.CheckerPass(full_check=True)(model)

ir.save(model, "updated.onnx")

These stages serve different purposes:

  • TopologicalSortPass stably sorts the main graph, nested subgraphs, and model functions. A cycle raises ValueError.

  • ShapeInferencePass asks ONNX shape inference to update value types and shapes. If inference fails, it logs a warning and leaves the model unchanged.

  • CheckerPass serializes through the ONNX boundary and runs onnx.checker.check_model; it does not modify the model.

  • ir.save performs final serialization and external-data handling.

Use symbolic shape inference¶

For richer symbolic inference, use the optional onnx-shape-inference package. It operates directly on ONNX IR, represents dimension arithmetic with SymPy expressions, propagates shape-tensor data through patterns such as Shape -> Slice -> Concat -> Reshape, and supports custom operator inference functions.

pip install onnx-shape-inference

Use it instead of the built-in ShapeInferencePass when a later stage needs richer symbolic relationships:

from onnx_shape_inference import infer_symbolic_shapes

model = infer_symbolic_shapes(model)

The default refine merge policy preserves compatible existing information while adding inferred details. Use policy="strict" to report conflicts between declared and inferred shapes. Other policies and extension APIs are documented in the onnx-shape-inference project.

Shape and type information is not automatically recomputed after every graph edit. Run inference when downstream transformations or consumers depend on updated information, and use pass preconditions or postconditions for requirements that are specific to your pipeline.

Save large initializers as ONNX external data¶

ir.save(
    model,
    "model.onnx",
    external_data="model.data",       # relative to model path
    size_threshold_bytes=1024,        # externalize tensors >= 1KB
    max_shard_size_bytes=256 * 1024,  # optional sharding
)

Notes:

  • external_data must be a relative path.

  • max_shard_size_bytes requires external_data.

  • Single-file mode writes in the destination directory and atomically replaces the destination only after the new data file is complete. A failed write leaves the previous file unchanged.

  • Sharded mode is stricter and can raise FileExistsError for collisions. This is a preflight check; another process creating a destination shard concurrently can still race with the final replacement.

By default, tensors are stored densely in initializer declaration order, without padding between them. Set alignment=65536 to retain the previous layout, which aligns tensors larger than 1 MiB to Windows’ 64 KiB allocation granularity:

ir.save(
    model,
    "model.onnx",
    external_data="model.data",
    alignment=65536,
)

Save external data with progress callback¶

def callback(tensor, info):
    print(f"[{info.index + 1}/{info.total}] {info.filename} :: {tensor.name}")

ir.save(
    model,
    "model.onnx",
    external_data="model.data",
    size_threshold_bytes=0,
    callback=callback,
)

Speed up saving with multiple threads¶

Pass max_workers to overlap tensor materialization (lazy tensor evaluation, dtype conversion) with disk writes and to write in parallel. This is most effective when initializers need work before they can be written, such as a ir.LazyTensor that casts from bfloat16:

ir.save(
    model,
    "model.onnx",
    external_data="model.data",
    max_workers=8,
)

Weights that live on an accelerator benefit the most. The device-to-host copy runs inside the write with the GIL released, so it overlaps with other threads’ writes instead of serializing behind them.

Peak memory stays bounded regardless of the worker count: at most max_in_flight_bytes (1GB by default) plus the size of the largest single tensor. Lower it when memory is tight:

ir.save(
    model,
    "model.onnx",
    external_data="model.data",
    max_workers=8,
    max_in_flight_bytes=64 * 1024**2,
)

Note

When max_workers is greater than 1, the callback is called from worker threads. Calls are serialized with a lock, so the callback does not need to be thread-safe itself, but it is no longer invoked in info.index order. Write progress callbacks as counters rather than assuming index increases:

import threading

lock = threading.Lock()
done = 0

def callback(tensor, info):
    global done
    with lock:
        done += 1
        print(f"[{done}/{info.total}] {tensor.name}")

Save with safetensors backend¶

ir.save_safetensors(
    model,
    "model.onnx",
    size_threshold_bytes=0,
    max_shard_size_bytes=5 * 1000**3,
)

save_safetensors writes model weights to .safetensors side files while keeping the ONNX graph in model.onnx.

Important safetensors constraints¶

  • All initializer names across graphs/subgraphs must be unique.

  • Tensor attributes in nodes are not externalized to safetensors.

  • For large constant nodes you want externalized, consider lifting constants to initializers before saving.