.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples_compute/plot_compute_information.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_compute_plot_compute_information.py: .. _l-example-plot-compute-information: Computing shape, tag, constant, release and in-place information ================================================================ *onnx-light* derives several pieces of information about a graph before it is executed: * **shape** — the element type and shape of every value (:mod:`onnx_light.onnx_core.shape_inference`). * **shape_tag** — a semantic tag (``shape``, ``axes``, ``weight`` or ``ambiguous``) attached to every value and node, guessing what role a tensor plays (e.g. a tensor holding a shape versus a numerical weight). * **constant** — whether a value's content is entirely known before inference starts (initializers, ``Constant`` outputs, and outputs of deterministic nodes whose inputs are all constant). * **release** — the last node after which a value is no longer needed and its buffer can be released. * **inplace** — which node outputs can reuse one of their input buffers instead of allocating a new one. This example builds a small model that exercises all five analyses, runs them through :class:`~onnx_light.onnx_core.shape_inference.ComputeContext` and the free functions in :mod:`onnx_light.onnx_core.shape_inference`, and shows how to retrieve the results either as in-memory objects or as ``metadata_props`` entries written directly on the model. The model computes:: two = Constant(2) # constant scalar const_prod = Mul(W, two) # constant (W is an initializer) added = Add(X, const_prod) # not constant (depends on X) relu_out = Relu(added) # may reuse "added" in place x_shape = Shape(X) # tagged "shape" Z = Reshape(relu_out, x_shape) # uses the "shape" tensor .. GENERATED FROM PYTHON SOURCE LINES 38-73 .. 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_core.shape_inference import ( INPLACE_REUSE_METADATA_KEY, RELEASE_AFTER_METADATA_KEY, RELEASE_AFTER_SHAPE_TAG_METADATA_KEY, ComputeContext, ShapesContext, apply_inferred_shapes_to_model, compute_inplace_reuse, compute_shape_model, write_constant_info_to_metadata, write_inplace_reuse_to_metadata, ) from onnx_light.tools import pretty_onnx # Not exposed as a Python constant yet; mirrors the C++ # ``onnx_compute::kNotUsedAfterMetadataKey`` used by # :func:`write_inplace_reuse_to_metadata`. NOT_USED_AFTER_METADATA_KEY = "onnx_light.not_used_after" def _metadata_value(node, key: str) -> str: """Returns the ``metadata_props`` value for ``key`` on ``node``, or ``""``.""" return next((entry.value for entry in node.metadata_props if entry.key == key), "") # Make sure the built-in operator schemas are registered before running # shape inference (the C++ dispatch table looks them up). defs.register_onnx_operator_set_schema() .. GENERATED FROM PYTHON SOURCE LINES 74-76 Build the model +++++++++++++++ .. GENERATED FROM PYTHON SOURCE LINES 76-103 .. code-block:: Python model = oh.make_model( oh.make_graph( [ oh.make_node( "Constant", [], ["two"], value=oh.make_tensor("two", onnxl.TensorProto.FLOAT, [], [2.0]), ), oh.make_node("Mul", ["W", "two"], ["const_prod"]), oh.make_node("Add", ["X", "const_prod"], ["added"]), oh.make_node("Relu", ["added"], ["relu_out"]), oh.make_node("Shape", ["X"], ["x_shape"]), oh.make_node("Reshape", ["relu_out", "x_shape"], ["Z"]), ], "compute_information_demo", inputs=[oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [3, 4])], outputs=[oh.make_tensor_value_info("Z", onnxl.TensorProto.FLOAT, None)], initializer=[oh.make_tensor("W", onnxl.TensorProto.FLOAT, [4], [1.0, 2.0, 3.0, 4.0])], ), opset_imports=[oh.make_opsetid("", 18)], ir_version=8, ) print(pretty_onnx(model)) .. rst-class:: sphx-glr-script-out .. code-block:: none opset: domain='' version=18 graph: name='compute_information_demo' input: float[3,4] X init: float[4] W 0: Constant() -> two 1: Mul(W, two) -> const_prod 2: Add(X, const_prod) -> added 3: Relu(added) -> relu_out 4: Shape(X) -> x_shape 5: Reshape(relu_out, x_shape) -> Z output: float[] Z .. GENERATED FROM PYTHON SOURCE LINES 104-111 1. shape ++++++++ :func:`compute_shape_model` seeds a :class:`ShapesContext` from the model's opset, initializers and inputs, then runs per-operator shape inference on every node. :func:`apply_inferred_shapes_to_model` writes the result back into ``model.graph.value_info`` and ``model.graph.output``. .. GENERATED FROM PYTHON SOURCE LINES 111-121 .. code-block:: Python shapes_ctx = ShapesContext() compute_shape_model(shapes_ctx, model) apply_inferred_shapes_to_model(shapes_ctx, model) print("Inferred shapes:") for name in ("two", "const_prod", "added", "relu_out", "x_shape", "Z"): descriptor = shapes_ctx.get(name) print(f" {name:<12} dtype={descriptor.dtype} shape={descriptor.shape.dims()}") .. rst-class:: sphx-glr-script-out .. code-block:: none Inferred shapes: two dtype=1 shape=[] const_prod dtype=1 shape=[4] added dtype=1 shape=[3, 4] relu_out dtype=1 shape=[3, 4] x_shape dtype=7 shape=[2] Z dtype=1 shape=[3, 4] .. GENERATED FROM PYTHON SOURCE LINES 122-128 2. shape_tag ++++++++++++ :meth:`ComputeContext.compute_value_and_node_tags` classifies every value and node. ``x_shape`` is tagged ``"shape"`` because it is produced by a ``Shape`` node and consumed as a shape argument by ``Reshape``. .. GENERATED FROM PYTHON SOURCE LINES 128-140 .. code-block:: Python compute_ctx = ComputeContext() value_tags, node_tags = compute_ctx.compute_value_and_node_tags(model.graph) print("\nValue tags:") for name, tag in sorted(value_tags.items()): print(f" {name:<12} {tag}") print("\nNode tags:") for node, tag in zip(model.graph.node, node_tags): print(f" {node.op_type:<10} outputs={list(node.output)!s:<16} tag={tag}") .. rst-class:: sphx-glr-script-out .. code-block:: none Value tags: W weight X weight Z weight added weight const_prod weight relu_out weight two weight x_shape shape Node tags: Constant outputs=['two'] tag=weight Mul outputs=['const_prod'] tag=weight Add outputs=['added'] tag=weight Relu outputs=['relu_out'] tag=weight Shape outputs=['x_shape'] tag=shape Reshape outputs=['Z'] tag=weight .. GENERATED FROM PYTHON SOURCE LINES 141-149 3. constant +++++++++++ ``two`` and ``const_prod`` only depend on the ``Constant`` node and the initializer ``W``, so both are constant. ``added``, ``relu_out``, ``x_shape`` and ``Z`` depend on the graph input ``X`` and are not. :func:`write_constant_info_to_metadata` records the result directly on the model, under the ``onnx_light.constant`` metadata key. .. GENERATED FROM PYTHON SOURCE LINES 149-171 .. code-block:: Python write_constant_info_to_metadata(model) def _is_constant(value_infos) -> dict[str, bool]: """Returns ``{name: is_constant}`` read from ``onnx_light.constant`` metadata.""" result = {} for value_info in value_infos: result[value_info.name] = any( entry.key == "onnx_light.constant" for entry in value_info.metadata_props ) return result constant_values = _is_constant(model.graph.value_info) constant_values.update(_is_constant(model.graph.initializer)) constant_values.update(_is_constant(model.graph.output)) print("\nConstant values:") for name in ("W", "two", "const_prod", "added", "relu_out", "x_shape", "Z"): print(f" {name:<12} {constant_values.get(name, False)}") .. rst-class:: sphx-glr-script-out .. code-block:: none Constant values: W True two True const_prod True added False relu_out False x_shape False Z False .. GENERATED FROM PYTHON SOURCE LINES 172-192 4. release / 5. inplace ++++++++++++++++++++++++ :func:`compute_inplace_reuse` returns, for every node, the list of ``InPlaceReuse`` opportunities (which output can reuse which input buffer). :func:`write_inplace_reuse_to_metadata` additionally records *release* information on ``metadata_props``: * ``onnx_light.inplace_reuse`` — the in-place opportunities, as ``output_index:input_index:kind`` triplets. * ``onnx_light.release_after`` — the values that are no longer needed once the node has run. * ``onnx_light.release_after_shape_tag`` — the subset of those released values that carry the ``"shape"`` tag (from ``value_tags`` above). * ``onnx_light.not_used_after`` — declared graph inputs or initializers that reach their last use at this node. Here ``relu_out = Relu(added)`` can overwrite the buffer of ``added`` in place (same element type and shape), and ``added`` is released right after ``Relu`` runs since nothing else reads it. .. GENERATED FROM PYTHON SOURCE LINES 192-212 .. code-block:: Python reuse = compute_inplace_reuse(shapes_ctx, model.graph) write_inplace_reuse_to_metadata(shapes_ctx, model.graph, value_tags) print("\nIn-place reuse and release information per node:") for node, node_reuse in zip(model.graph.node, reuse): reuse_desc = ", ".join( f"out{r.output_index}=in{r.input_index}({r.kind.name})" for r in node_reuse ) inplace_metadata = _metadata_value(node, INPLACE_REUSE_METADATA_KEY) release_after = _metadata_value(node, RELEASE_AFTER_METADATA_KEY) release_after_shape_tag = _metadata_value(node, RELEASE_AFTER_SHAPE_TAG_METADATA_KEY) not_used_after = _metadata_value(node, NOT_USED_AFTER_METADATA_KEY) print( f" {node.op_type:<10} outputs={list(node.output)!s:<16} " f"inplace=[{reuse_desc}] (metadata={inplace_metadata!r}) " f"release_after={release_after!r} " f"release_after_shape_tag={release_after_shape_tag!r} " f"not_used_after={not_used_after!r}" ) .. rst-class:: sphx-glr-script-out .. code-block:: none In-place reuse and release information per node: Constant outputs=['two'] inplace=[] (metadata='') release_after='' release_after_shape_tag='' not_used_after='' Mul outputs=['const_prod'] inplace=[] (metadata='') release_after='two' release_after_shape_tag='' not_used_after='W' Add outputs=['added'] inplace=[] (metadata='') release_after='const_prod' release_after_shape_tag='' not_used_after='' Relu outputs=['relu_out'] inplace=[out0=in0(kEqual)] (metadata='0:0:equal') release_after='added' release_after_shape_tag='' not_used_after='' Shape outputs=['x_shape'] inplace=[] (metadata='') release_after='' release_after_shape_tag='' not_used_after='X' Reshape outputs=['Z'] inplace=[out0=in0(kEqual)] (metadata='0:0:equal') release_after='relu_out;x_shape' release_after_shape_tag='x_shape' not_used_after='' .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.009 seconds) .. _sphx_glr_download_auto_examples_compute_plot_compute_information.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_compute_information.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_compute_information.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_compute_information.zip ` .. include:: plot_compute_information.recommendations .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_