.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples_optimization/plot_pretty_onnx.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_optimization_plot_pretty_onnx.py: .. _l-example-plot-pretty-onnx: pretty_onnx: shape info, shape tags, inplace and release annotations ====================================================================== :func:`~onnx_light.tools.pretty_onnx` renders any ONNX proto as a compact, human-readable text listing. Beyond the basic node-by-node view it supports four optional annotation layers: * **shape info** — run shape inference before rendering so every input, intermediate and output tensor shows its inferred dtype and shape. * **shape tags** — prefix nodes (and highlight tensors) whose role is identified as ``shape``, ``axes`` or ``weight`` by the semantic tag inference. * **inplace info** — annotate nodes whose output can safely reuse an input's buffer (``onnx_light.inplace_reuse`` metadata). * **release info** — annotate nodes after which a tensor is no longer needed and its memory can be freed (``onnx_light.release_after`` metadata). The example below builds a small graph, enriches it with all four kinds of metadata and then shows the output of :func:`~onnx_light.tools.pretty_onnx` at each level of verbosity. Graph structure ++++++++++++++++ .. code-block:: none X : float[N, 4] S = Shape(X) → int64[2] # S is a "shape" tensor A = Abs(X) → float[N, 4] B = Relu(A) → float[N, 4] # B can reuse A's buffer (inplace) Z = Reshape(B, S) → float[4, N] or float[N, 4] after inference .. GENERATED FROM PYTHON SOURCE LINES 37-56 .. code-block:: Python from __future__ import annotations import onnx_light.onnx as onnxl import onnx_light.onnx.defs as defs import onnx_light.onnx.helper as oh from onnx_light.onnx_optim.shape_inference import ( ShapesContext, apply_inferred_shapes_to_model, compute_shape_model, write_inplace_reuse_to_metadata, write_value_and_node_tags_to_metadata, ) from onnx_light.tools import pretty_onnx # Built-in operator schemas must be registered before shape inference. defs.register_onnx_operator_set_schema() .. GENERATED FROM PYTHON SOURCE LINES 57-62 Build the model +++++++++++++++ The graph is intentionally simple so that all four annotation layers are visible on just a handful of nodes. .. GENERATED FROM PYTHON SOURCE LINES 62-79 .. code-block:: Python model = oh.make_model( oh.make_graph( [ oh.make_node("Shape", ["X"], ["S"]), oh.make_node("Abs", ["X"], ["A"]), oh.make_node("Relu", ["A"], ["B"]), oh.make_node("Reshape", ["B", "S"], ["Z"]), ], "pretty_onnx_demo", inputs=[oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, ["N", 4])], outputs=[oh.make_tensor_value_info("Z", onnxl.TensorProto.FLOAT, None)], ), opset_imports=[oh.make_opsetid("", 18)], ir_version=8, ) .. GENERATED FROM PYTHON SOURCE LINES 80-81 Plain rendering — no annotations. .. GENERATED FROM PYTHON SOURCE LINES 81-86 .. code-block:: Python print("=== plain (no annotations) ===") print(pretty_onnx(model)) .. rst-class:: sphx-glr-script-out .. code-block:: none === plain (no annotations) === opset: domain='' version=18 graph: name='pretty_onnx_demo' input: float[N,4] X 0: Shape(X) -> S 1: Abs(X) -> A 2: Relu(A) -> B 3: Reshape(B, S) -> Z output: float[] Z .. GENERATED FROM PYTHON SOURCE LINES 87-99 Shape info ++++++++++ Pass ``shape_inference=True`` to run :mod:`onnx_light.onnx_optim.shape_inference` shape inference before rendering. After inference every value in ``model.graph.value_info`` and ``model.graph.output`` carries its inferred dtype and shape, which :func:`~onnx_light.tools.pretty_onnx` then shows next to every input/output/intermediate name. .. note:: ``shape_inference=True`` creates a *copy* of the model with the inferred shapes filled in; the original ``model`` object is not mutated. .. GENERATED FROM PYTHON SOURCE LINES 99-104 .. code-block:: Python print("\n=== shape_inference=True ===") print(pretty_onnx(model, shape_inference=True)) .. rst-class:: sphx-glr-script-out .. code-block:: none === shape_inference=True === opset: domain='' version=18 graph: name='pretty_onnx_demo' input: float[N,4] X 0: Shape(X) -> S 1: Abs(X) -> A 2: Relu(A) -> B 3: Reshape(B, S) -> Z output: float[] Z .. GENERATED FROM PYTHON SOURCE LINES 105-119 Shape tags ++++++++++ Shape-tag inference labels every tensor and node with a semantic role: * ``shape`` — a tensor whose value represents a shape or size (e.g. the output of a ``Shape`` node). * ``axes`` — a tensor whose value represents a set of axis indices. * ``weight`` — a trainable parameter (initializer). :func:`~onnx_light.onnx_optim.shape_inference.write_value_and_node_tags_to_metadata` writes these labels into ``metadata_props``; then ``include_node_tags=True`` renders them as ``[shape]``, ``[axes]`` or ``[weight]`` prefixes on the relevant node lines. .. GENERATED FROM PYTHON SOURCE LINES 119-126 .. code-block:: Python write_value_and_node_tags_to_metadata(model.graph) print("\n=== include_node_tags=True (after write_value_and_node_tags_to_metadata) ===") print(pretty_onnx(model, include_node_tags=True)) .. rst-class:: sphx-glr-script-out .. code-block:: none === include_node_tags=True (after write_value_and_node_tags_to_metadata) === opset: domain='' version=18 graph: name='pretty_onnx_demo' input: float[N,4] X 0: [shape] Shape(X) -> S 1: [weight] Abs(X) -> A 2: [weight] Relu(A) -> B 3: [weight] Reshape(B, S) -> Z output: float[] Z .. GENERATED FROM PYTHON SOURCE LINES 127-143 Inplace info ++++++++++++ Buffer-reuse analysis identifies pairs of node inputs and outputs that are *aliasable*: the output can overwrite the input's allocation when the shapes match and the input is not needed by any later node. :func:`~onnx_light.onnx_optim.shape_inference.write_inplace_reuse_to_metadata` records each reuse opportunity under the ``onnx_light.inplace_reuse`` metadata key as a ``out_idx:in_idx:kind`` triplet. With ``include_inplace=True`` these are rendered as ``inplace: out0=in0(equal)`` (or ``(greater)`` when the output buffer is strictly larger than the input). Shape inference must be run first so the analysis has concrete shapes to compare. .. GENERATED FROM PYTHON SOURCE LINES 143-153 .. code-block:: Python ctx = ShapesContext() compute_shape_model(ctx, model) apply_inferred_shapes_to_model(ctx, model) write_inplace_reuse_to_metadata(ctx, model.graph) print("\n=== include_inplace=True (after write_inplace_reuse_to_metadata) ===") print(pretty_onnx(model, include_inplace=True)) .. rst-class:: sphx-glr-script-out .. code-block:: none === include_inplace=True (after write_inplace_reuse_to_metadata) === opset: domain='' version=18 graph: name='pretty_onnx_demo' input: float[N,4] X 0: Shape(X) -> S 1: Abs(X) -> A 2: Relu(A) -> B inplace: out0=in0(equal) 3: Reshape(B, S) -> Z inplace: out0=in0(equal) output: float[N,4] Z .. GENERATED FROM PYTHON SOURCE LINES 154-163 Release info ++++++++++++ The same ``write_inplace_reuse_to_metadata`` call also records, for each node, the set of tensors whose last use ends at that node. A runtime can free those buffers immediately after the node executes. The names are stored under the ``onnx_light.release_after`` metadata key as a ``;``-separated list. With ``include_release=True`` they are rendered as ``release: A`` (or ``release: A, B`` when multiple tensors are freed). .. GENERATED FROM PYTHON SOURCE LINES 163-168 .. code-block:: Python print("\n=== include_release=True ===") print(pretty_onnx(model, include_release=True)) .. rst-class:: sphx-glr-script-out .. code-block:: none === include_release=True === opset: domain='' version=18 graph: name='pretty_onnx_demo' input: float[N,4] X 0: Shape(X) -> S 1: Abs(X) -> A 2: Relu(A) -> B release: A 3: Reshape(B, S) -> Z release: B, S output: float[N,4] Z .. GENERATED FROM PYTHON SOURCE LINES 169-175 All annotations combined +++++++++++++++++++++++++ The four flags are independent and compose freely. The listing below enables all of them at once so you can see the full picture in a single pass. .. GENERATED FROM PYTHON SOURCE LINES 175-178 .. code-block:: Python print("\n=== all annotations combined ===") print(pretty_onnx(model, include_node_tags=True, include_inplace=True, include_release=True)) .. rst-class:: sphx-glr-script-out .. code-block:: none === all annotations combined === opset: domain='' version=18 graph: name='pretty_onnx_demo' input: float[N,4] X 0: [shape] Shape(X) -> S 1: [weight] Abs(X) -> A 2: [weight] Relu(A) -> B inplace: out0=in0(equal) release: A 3: [weight] Reshape(B, S) -> Z inplace: out0=in0(equal) release: B, S output: float[N,4] Z .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.011 seconds) .. _sphx_glr_download_auto_examples_optimization_plot_pretty_onnx.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_pretty_onnx.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_pretty_onnx.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_pretty_onnx.zip ` .. include:: plot_pretty_onnx.recommendations .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_