.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples_proto/plot_node_callback.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_examples_proto_plot_node_callback.py: .. _l-example-plot-node-callback: Inspect and edit nodes while parsing or serializing with a node callback ======================================================================== This example shows how to use :attr:`onnx_light.onnx.ParseOptions.node_callback` and :attr:`onnx_light.onnx.SerializeOptions.node_callback` to hook into every :class:`onnx_light.onnx.NodeProto` of a model. The callback receives each node by reference along with its parent :class:`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``. .. GENERATED FROM PYTHON SOURCE LINES 17-24 .. code-block:: Python 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 .. GENERATED FROM PYTHON SOURCE LINES 25-31 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. .. GENERATED FROM PYTHON SOURCE LINES 31-42 .. code-block:: Python 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() .. GENERATED FROM PYTHON SOURCE LINES 43-50 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. .. GENERATED FROM PYTHON SOURCE LINES 50-68 .. code-block:: Python 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}") .. rst-class:: sphx-glr-script-out .. code-block:: none parsed node 'Identity' in graph 'then_graph' parsed node 'Add' in graph 'main' parsed node 'If' in graph 'main' .. GENERATED FROM PYTHON SOURCE LINES 69-74 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. .. GENERATED FROM PYTHON SOURCE LINES 74-83 .. code-block:: Python 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}") .. rst-class:: sphx-glr-script-out .. code-block:: none add0 doc_string: 'parsed' if0 doc_string: 'parsed' .. GENERATED FROM PYTHON SOURCE LINES 84-91 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. .. GENERATED FROM PYTHON SOURCE LINES 91-103 .. code-block:: Python 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}") .. rst-class:: sphx-glr-script-out .. code-block:: none serialized add0 doc_string: 'serialized' serialized subgraph doc_string: 'serialized' .. GENERATED FROM PYTHON SOURCE LINES 104-106 The caller's model is untouched: onnx-light restored every node the callback edited once the serialized bytes were produced. .. GENERATED FROM PYTHON SOURCE LINES 106-108 .. code-block:: Python print(f"original add0 doc_string still empty: {onnx_model.graph.node[0].doc_string!r}") .. rst-class:: sphx-glr-script-out .. code-block:: none original add0 doc_string still empty: '' .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.002 seconds) .. _sphx_glr_download_auto_examples_proto_plot_node_callback.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_node_callback.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_node_callback.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_node_callback.zip ` .. include:: plot_node_callback.recommendations .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_