Getting started with ONNX IR 🌱

The ONNX IR is available as onnx_ir. Use ir.load("model.onnx") to load a model from disk. This notebook starts from ONNX text so it is self-contained, then demonstrates inspection, transformation, validation, and saving.

# Define an example model for this example
MODEL_TEXT = r"""
<
   ir_version: 8,
   opset_import: ["" : 18],
   producer_name: "pytorch",
   producer_version: "2.0.0"
>
torch_jit (float[5,5,5] input_0) => (float[5,5] val_19, float[5,5] val_6) {
   [node_1] val_1 = Constant <value_int: int = 1> ()
   [node_2] val_2 = Shape <start: int = 0> (val_1)
   [node_3] val_3 = Size (val_2)
   [node_4] val_4 = Constant <value: tensor = int64 {0}> ()
   [node_5] val_5 = Equal (val_3, val_4)
   [node_6] val_6 = ReduceMean <keepdims: int = 0, noop_with_empty_axes: int = 0> (input_0, val_1)
   [node_7] val_7 = ReduceMean <keepdims: int = 1, noop_with_empty_axes: int = 0> (input_0, val_1)
   [node_8] val_8 = Shape <start: int = 0> (input_0)
   [node_9] val_9 = Gather <axis: int = 0> (val_8, val_1)
   [node_10] val_10 = ReduceProd <keepdims: int = 0, noop_with_empty_axes: int = 0> (val_9)
   [node_11] val_11 = Sub (input_0, val_7)
   [node_12] val_12 = Mul (val_11, val_11)
   [node_13] val_13 = ReduceMean <keepdims: int = 0, noop_with_empty_axes: int = 0> (val_12, val_1)
   [node_14] val_14 = Cast <to: int = 1> (val_10)
   [node_15] val_15 = Mul (val_13, val_14)
   [node_16] val_16 = Constant <value: tensor = int64 {1}> ()
   [node_17] val_17 = Sub (val_10, val_16)
   [node_18] val_18 = Cast <to: int = 1> (val_17)
   [node_19] val_19 = Div (val_15, val_18)
}
"""

import onnx

import onnx_ir as ir

# Load the model as onnx.ModelProto
# You can also load the model from a file using onnx.load("model.onnx")
model_proto = onnx.parser.parse_model(MODEL_TEXT)

# Create an IR object from the model
model = ir.from_proto(model_proto)

Now we can explore the IR object

print(f"The main graph has {len(model.graph)} nodes.")

The main graph has 19 nodes.

All inputs

print(model.graph.inputs)

[Value(name='input_0', type=Tensor(FLOAT), shape=Shape([5, 5, 5]))]

All outputs

print(model.graph.outputs)

[Value(name='val_19', type=Tensor(FLOAT), shape=Shape([5, 5]), producer='node_19', index=0), Value(name='val_6', type=Tensor(FLOAT), shape=Shape([5, 5]), producer='node_6', index=0)]

Nodes that uses the first input

print(list(model.graph.inputs[0].uses()))

[Usage(node=Node(name='node_6', domain='', op_type='ReduceMean', inputs=(Value(name='input_0', type=Tensor(FLOAT), shape=Shape([5, 5, 5])), Value(name='val_1', producer='node_1', index=0)), attributes={'keepdims': Attr('keepdims', INT, 0), 'noop_with_empty_axes': Attr('noop_with_empty_axes', INT, 0)}, overload='', outputs=(Value(name='val_6', type=Tensor(FLOAT), shape=Shape([5, 5]), producer='node_6', index=0),), version=None, doc_string=None), idx=0), Usage(node=Node(name='node_7', domain='', op_type='ReduceMean', inputs=(Value(name='input_0', type=Tensor(FLOAT), shape=Shape([5, 5, 5])), Value(name='val_1', producer='node_1', index=0)), attributes={'keepdims': Attr('keepdims', INT, 1), 'noop_with_empty_axes': Attr('noop_with_empty_axes', INT, 0)}, overload='', outputs=(Value(name='val_7', producer='node_7', index=0),), version=None, doc_string=None), idx=0), Usage(node=Node(name='node_8', domain='', op_type='Shape', inputs=(Value(name='input_0', type=Tensor(FLOAT), shape=Shape([5, 5, 5])),), attributes={'start': Attr('start', INT, 0)}, overload='', outputs=(Value(name='val_8', producer='node_8', index=0),), version=None, doc_string=None), idx=0), Usage(node=Node(name='node_11', domain='', op_type='Sub', inputs=(Value(name='input_0', type=Tensor(FLOAT), shape=Shape([5, 5, 5])), Value(name='val_7', producer='node_7', index=0)), attributes={}, overload='', outputs=(Value(name='val_11', producer='node_11', index=0),), version=None, doc_string=None), idx=0)]

The node that produces the last output (as the i-th output)

print(model.graph.outputs[-1].producer())
print(model.graph.outputs[-1].index())

%"val_6"<FLOAT,[5,5]> ⬅️ ::ReduceMean(%"input_0", %"val_1") {keepdims=0, noop_with_empty_axes=0}
0

Print the graph

print(model.graph)

graph(
    name=torch_jit,
    inputs=(
        %"input_0"<FLOAT,[5,5,5]>
    ),
    outputs=(
        %"val_19"<FLOAT,[5,5]>,
        %"val_6"<FLOAT,[5,5]>
    ),
) {
     0 |  # node_1
          %"val_1"<?,?> ⬅️ ::Constant() {value_int=1}
     1 |  # node_2
          %"val_2"<?,?> ⬅️ ::Shape(%"val_1") {start=0}
     2 |  # node_3
          %"val_3"<?,?> ⬅️ ::Size(%"val_2")
     3 |  # node_4
          %"val_4"<?,?> ⬅️ ::Constant() {value=TensorProtoTensor<INT64,[]>(array(0), name='')}
     4 |  # node_5
          %"val_5"<?,?> ⬅️ ::Equal(%"val_3", %"val_4")
     5 |  # node_6
          %"val_6"<FLOAT,[5,5]> ⬅️ ::ReduceMean(%"input_0", %"val_1") {keepdims=0, noop_with_empty_axes=0}
     6 |  # node_7
          %"val_7"<?,?> ⬅️ ::ReduceMean(%"input_0", %"val_1") {keepdims=1, noop_with_empty_axes=0}
     7 |  # node_8
          %"val_8"<?,?> ⬅️ ::Shape(%"input_0") {start=0}
     8 |  # node_9
          %"val_9"<?,?> ⬅️ ::Gather(%"val_8", %"val_1") {axis=0}
     9 |  # node_10
          %"val_10"<?,?> ⬅️ ::ReduceProd(%"val_9") {keepdims=0, noop_with_empty_axes=0}
    10 |  # node_11
          %"val_11"<?,?> ⬅️ ::Sub(%"input_0", %"val_7")
    11 |  # node_12
          %"val_12"<?,?> ⬅️ ::Mul(%"val_11", %"val_11")
    12 |  # node_13
          %"val_13"<?,?> ⬅️ ::ReduceMean(%"val_12", %"val_1") {keepdims=0, noop_with_empty_axes=0}
    13 |  # node_14
          %"val_14"<?,?> ⬅️ ::Cast(%"val_10") {to=1}
    14 |  # node_15
          %"val_15"<?,?> ⬅️ ::Mul(%"val_13", %"val_14")
    15 |  # node_16
          %"val_16"<?,?> ⬅️ ::Constant() {value=TensorProtoTensor<INT64,[]>(array(1), name='')}
    16 |  # node_17
          %"val_17"<?,?> ⬅️ ::Sub(%"val_10", %"val_16")
    17 |  # node_18
          %"val_18"<?,?> ⬅️ ::Cast(%"val_17") {to=1}
    18 |  # node_19
          %"val_19"<FLOAT,[5,5]> ⬅️ ::Div(%"val_15", %"val_18")
    return %"val_19"<FLOAT,[5,5]>, %"val_6"<FLOAT,[5,5]>
}

Transform the graph

Add an Identity after the second graph output and make its result the new output.

old_output = model.graph.outputs[-1]
producer = old_output.producer()
assert producer is not None

identity = ir.node("Identity", inputs=[old_output], name="output_identity")
identity.outputs[0].type = old_output.type
identity.outputs[0].shape = old_output.shape
model.graph.insert_after(producer, identity)
model.graph.outputs[-1] = identity.outputs[0]

The insertion preserves names, ownership, use-def relationships, and topological order, so no repair passes are needed. At this tutorial’s validation boundary, run the ONNX checker explicitly.

import onnx_ir.passes.common as common_passes

common_passes.CheckerPass(full_check=True)(model)

PassResult(model=Model(
    ir_version=8,
    opset_imports={'': 18},
    producer_name='pytorch',
    producer_version='2.0.0',
    domain=None,
    model_version=None,
    functions={},
    graph=Graph(
        name='torch_jit',
        inputs=(
            %"input_0"<FLOAT,[5,5,5]>
        ),
        outputs=(
            %"val_19"<FLOAT,[5,5]>,
            %"val_0"<FLOAT,[5,5]>
        ),
        len()=20
    )
), modified=False)

Save and reload

Use ir.save and ir.load at file boundaries. A temporary directory keeps this notebook from leaving generated files behind.

import os
import tempfile

with tempfile.TemporaryDirectory() as temp_dir:
    model_path = os.path.join(temp_dir, "model.onnx")
    ir.save(model, model_path)
    reloaded_model = ir.load(model_path)

assert reloaded_model.graph.outputs[-1].producer().op_type == "Identity"

Convert to or from ModelProto explicitly when integrating with protobuf-based APIs.

model_proto_back = ir.to_proto(model)
model_round_trip = ir.from_proto(model_proto_back)

Next steps