Inspect and edit nodes while parsing or serializing with a node callback

This example shows how to use onnx_light.onnx.ParseOptions.node_callback and onnx_light.onnx.SerializeOptions.node_callback to hook into every onnx_light.onnx.NodeProto of a model.

The callback receives each node by reference along with its parent onnx_light.onnx.GraphProto and may inspect or modify the node in place. The parent graph lets the callback locate the node’s surrounding graph — including the subgraphs nested inside control flow operators such as If, Loop and Scan.

import numpy as np

import onnx_light.onnx.helper as oh
import onnx_light.onnx.numpy_helper as onh
import onnx_light.onnx as onnxl

Build a small model with a subgraph

The main graph holds an Add node and an If node. The If node carries a then_branch subgraph with its own Identity node, so the callback visits nodes across two graphs.

arr = np.array([1.0, 2.0], dtype=np.float32)
add = oh.make_node("Add", ["X", "W"], ["Y"], name="add0")
sub_node = oh.make_node("Identity", ["cond"], ["Z"], name="id0")
then_graph = oh.make_graph([sub_node], "then_graph", [], [])
if_node = oh.make_node("If", ["cond"], ["Z"], name="if0", then_branch=then_graph)
graph = oh.make_graph([add, if_node], "main", [], [], initializer=[onh.from_array(arr, name="W")])
onnx_model = oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)], ir_version=9)

serialized = onnx_model.SerializeToString()

Inspect every node while parsing

node_callback fires once per node. The parent graph is passed as the second argument, so we can record which graph each node belongs to. The subgraph node is visited while the parser reads the enclosing If node, before the rest of the main graph.

parse_options = onnxl.ParseOptions()
visited = []


def on_node(node: onnxl.NodeProto, graph: onnxl.GraphProto):
    """Records the node op_type and the name of its parent graph."""
    visited.append((node.op_type, graph.name))


parse_options.node_callback = on_node

parsed_model = onnxl.ModelProto()
parsed_model.ParseFromString(serialized, parse_options)

for op_type, graph_name in visited:
    print(f"parsed node {op_type!r} in graph {graph_name!r}")

parsed node 'Identity' in graph 'then_graph'
parsed node 'Add' in graph 'main'
parsed node 'If' in graph 'main'

Edit nodes in place while parsing

Because the callback receives each node by reference, it can rewrite the node. Here we stamp a doc_string on every node as it is parsed.

edit_options = onnxl.ParseOptions()
edit_options.node_callback = lambda node, graph: setattr(node, "doc_string", "parsed")

edited = onnxl.ModelProto()
edited.ParseFromString(serialized, edit_options)
print(f"add0 doc_string: {edited.graph.node[0].doc_string!r}")
print(f"if0 doc_string:  {edited.graph.node[1].doc_string!r}")

add0 doc_string: 'parsed'
if0 doc_string:  'parsed'

Edit nodes while serializing

SerializeOptions.node_callback works the same way. The callback edits the nodes in place while the serialized bytes are produced, then onnx-light restores the original state, so edits never alter the model held by the caller. The stamped doc_string therefore appears only in the serialized bytes.

serialize_options = onnxl.SerializeOptions()
serialize_options.node_callback = lambda node, graph: setattr(node, "doc_string", "serialized")

stamped_bytes = onnx_model.SerializeToString(serialize_options)

reparsed = onnxl.ModelProto()
reparsed.ParseFromString(stamped_bytes)
print(f"serialized add0 doc_string:      {reparsed.graph.node[0].doc_string!r}")
sub = reparsed.graph.node[1].attribute[0].g
print(f"serialized subgraph doc_string:  {sub.node[0].doc_string!r}")

serialized add0 doc_string:      'serialized'
serialized subgraph doc_string:  'serialized'

The caller’s model is untouched: onnx-light restored every node the callback edited once the serialized bytes were produced.

print(f"original add0 doc_string still empty: {onnx_model.graph.node[0].doc_string!r}")

original add0 doc_string still empty: ''

Related examples

Track tensor weights while parsing with a raw_data callback

Track tensor weights while parsing with a raw_data callback

Load and save ONNX models with external data

Load and save ONNX models with external data

Number of threads used to load and save ONNX models

Number of threads used to load and save ONNX models

Example last updated

Date:

2026-08-21