.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples_optimization/plot_evaluate_shapes.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_evaluate_shapes.py: .. _l-example-plot-evaluate-shapes: Evaluating inferred shapes with concrete input dimensions ========================================================= ONNX models commonly use *symbolic* names — sometimes called *tokens* — for the dimensions of their inputs (``"batch"``, ``"seq_length"``, ``"d_model"``, ...). After running shape inference, every intermediate tensor's shape is expressed in terms of those input tokens, either as a plain dim_param (``"batch"``) or as a symbolic arithmetic expression (``"2*d_model"``, ``"batch*seq_length"``, ...). This example shows how to: 1. Build a small model whose inputs carry symbolic dimensions. 2. Run model-level shape inference with :func:`~onnx_light.onnx_optim.shape_inference.infer_shapes_model` to populate ``model.graph.value_info`` and ``model.graph.output`` with symbolic shapes. 3. Evaluate every inferred dimension to a concrete integer for several assignments of the input tokens, using :func:`~onnx_light.onnx_optim.expressions.evaluate_expression`. The same approach is useful in profiling and capacity-planning workflows where one wants to know the concrete tensor sizes a model produces for a given batch size, sequence length, hidden dimension, etc., without actually running the model. .. GENERATED FROM PYTHON SOURCE LINES 30-44 .. 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.expressions import evaluate_expression from onnx_light.onnx_optim.shape_inference import infer_shapes_model # Built-in operator schemas must be registered before shape inference # (the C++ dispatch table looks them up). defs.register_onnx_operator_set_schema() .. GENERATED FROM PYTHON SOURCE LINES 45-59 Build a model with symbolic input dimensions ++++++++++++++++++++++++++++++++++++++++++++ The graph mirrors a tiny transformer-style block: .. code-block:: none added = Add(X, Y) # [batch, seq, d_model] concat_out = Concat(added, X, axis=2) # [batch, seq, 2*d_model] Z = Reshape(concat_out, [0, 0, -1]) # [batch, seq, 2*d_model] Every input dimension is a symbolic *token*: ``"batch"``, ``"seq"`` and ``"d_model"``. None of them carries a concrete integer value at the time the model is built. .. GENERATED FROM PYTHON SOURCE LINES 59-82 .. code-block:: Python INPUT_TOKENS = ["batch", "seq", "d_model"] model = oh.make_model( oh.make_graph( [ oh.make_node("Add", ["X", "Y"], ["added"]), oh.make_node("Concat", ["added", "X"], ["concat_out"], axis=2), oh.make_node("Reshape", ["concat_out", "reshape_shape"], ["Z"]), ], "evaluate_shapes_demo", inputs=[ oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, INPUT_TOKENS), oh.make_tensor_value_info("Y", onnxl.TensorProto.FLOAT, INPUT_TOKENS), ], outputs=[oh.make_tensor_value_info("Z", onnxl.TensorProto.FLOAT, None)], initializer=[oh.make_tensor("reshape_shape", onnxl.TensorProto.INT64, [3], [0, 0, -1])], ), opset_imports=[oh.make_opsetid("", 18)], ir_version=8, ) .. GENERATED FROM PYTHON SOURCE LINES 83-92 Run shape inference +++++++++++++++++++ :func:`infer_shapes_model` mutates ``model`` in place and writes the inferred symbolic shapes into ``model.graph.value_info`` (intermediate tensors) and ``model.graph.output`` (graph outputs). Because the Reshape target ``[0, 0, -1]`` is propagated from the initializer, the last dimension of ``concat_out`` and ``Z`` is resolved symbolically to ``"2*d_model"``. .. GENERATED FROM PYTHON SOURCE LINES 92-118 .. code-block:: Python infer_shapes_model(model) def dims_of(value_info): """Returns the list of dim values/params for a ``ValueInfoProto``.""" return [ d.dim_value if d.dim_value else d.dim_param for d in value_info.type.tensor_type.shape.dim ] # Collect symbolic shapes for every named tensor in the graph: inputs, # intermediate value_info entries, and graph outputs. symbolic_shapes: dict[str, list] = {} for vi in model.graph.input: symbolic_shapes[vi.name] = dims_of(vi) for vi in model.graph.value_info: symbolic_shapes[vi.name] = dims_of(vi) for vi in model.graph.output: symbolic_shapes[vi.name] = dims_of(vi) print("Symbolic shapes after inference:") for name, shape in symbolic_shapes.items(): print(f" {name:<12} -> {shape}") .. rst-class:: sphx-glr-script-out .. code-block:: none Symbolic shapes after inference: X -> ['batch', 'seq', 'd_model'] Y -> ['batch', 'seq', 'd_model'] added -> ['batch', 'seq', 'd_model'] concat_out -> ['batch', 'seq', '2*d_model'] Z -> ['batch', 'seq', '2*d_model'] .. GENERATED FROM PYTHON SOURCE LINES 119-131 Evaluate every shape for concrete input values ++++++++++++++++++++++++++++++++++++++++++++++ Given an assignment of the input tokens to concrete integers, every dimension of every tensor can be evaluated: * Integer dims are returned unchanged. * String dims that match an input token are looked up directly. * String dims that are arithmetic expressions (e.g. ``"2*d_model"``) are reduced with :func:`~onnx_light.onnx_optim.expressions.evaluate_expression`, which accepts the mapping from token names to integer values. .. GENERATED FROM PYTHON SOURCE LINES 131-165 .. code-block:: Python def evaluate_shape(shape: list, assignment: "dict[str, int]") -> list[int]: """Returns *shape* with every symbolic dim replaced by its integer value.""" concrete: list[int] = [] for dim in shape: if isinstance(dim, int): concrete.append(dim) elif dim in assignment: concrete.append(assignment[dim]) else: # Symbolic expression such as ``"2*d_model"``. concrete.append(evaluate_expression(str(dim), assignment)) return concrete # Three different input-token assignments. The first one matches a tiny # debug configuration; the next two scale the batch and sequence length. ASSIGNMENTS = [ {"batch": 1, "seq": 8, "d_model": 16}, {"batch": 4, "seq": 128, "d_model": 64}, {"batch": 32, "seq": 512, "d_model": 768}, ] print("\nConcrete shapes for each input-token assignment:") header_assignments = [ "{" + ", ".join(f"{k}={v}" for k, v in a.items()) + "}" for a in ASSIGNMENTS ] print(f" {'tensor':<12} " + " ".join(f"{h:>28}" for h in header_assignments)) for name, shape in symbolic_shapes.items(): cells = [str(evaluate_shape(shape, a)) for a in ASSIGNMENTS] print(f" {name:<12} " + " ".join(f"{c:>28}" for c in cells)) .. rst-class:: sphx-glr-script-out .. code-block:: none Concrete shapes for each input-token assignment: tensor {batch=1, seq=8, d_model=16} {batch=4, seq=128, d_model=64} {batch=32, seq=512, d_model=768} X [1, 8, 16] [4, 128, 64] [32, 512, 768] Y [1, 8, 16] [4, 128, 64] [32, 512, 768] added [1, 8, 16] [4, 128, 64] [32, 512, 768] concat_out [1, 8, 32] [4, 128, 128] [32, 512, 1536] Z [1, 8, 32] [4, 128, 128] [32, 512, 1536] .. GENERATED FROM PYTHON SOURCE LINES 166-172 Per-tensor element count ++++++++++++++++++++++++ A common follow-up question is *"how many elements does each tensor hold for this configuration?"*. Once the shape has been evaluated to concrete integers the answer is just the product of the dimensions. .. GENERATED FROM PYTHON SOURCE LINES 172-180 .. code-block:: Python from math import prod # noqa: E402 print("\nElement counts per tensor and per assignment:") print(f" {'tensor':<12} " + " ".join(f"{h:>28}" for h in header_assignments)) for name, shape in symbolic_shapes.items(): cells = [str(prod(evaluate_shape(shape, a))) for a in ASSIGNMENTS] print(f" {name:<12} " + " ".join(f"{c:>28}" for c in cells)) .. rst-class:: sphx-glr-script-out .. code-block:: none Element counts per tensor and per assignment: tensor {batch=1, seq=8, d_model=16} {batch=4, seq=128, d_model=64} {batch=32, seq=512, d_model=768} X 128 32768 12582912 Y 128 32768 12582912 added 128 32768 12582912 concat_out 256 65536 25165824 Z 256 65536 25165824 .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.008 seconds) .. _sphx_glr_download_auto_examples_optimization_plot_evaluate_shapes.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_evaluate_shapes.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_evaluate_shapes.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_evaluate_shapes.zip ` .. include:: plot_evaluate_shapes.recommendations .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_