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 overwrites destination data file.

  • Sharded mode is stricter and can raise FileExistsError for collisions.

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,
)

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.