.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples_optimization/plot_compute_context_memory.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_compute_context_memory.py: .. _l-example-plot-compute-context-memory: ComputeContext memory expressions ================================= :class:`~onnx_light.onnx_optim.shape_inference.ComputeContext` reports, for each node, how much memory is already live before execution, how much extra output allocation is still needed, and the resulting total. When input shapes are symbolic, these quantities stay symbolic as well. This example shows how to: 1. Build a small graph with one symbolic dimension ``N``. 2. Run shape inference and memory analysis. 3. Print a table with the symbolic memory expressions for every node. 4. Evaluate ``total_bytes`` for a few concrete values of ``N`` and plot the resulting curves. .. GENERATED FROM PYTHON SOURCE LINES 20-47 .. code-block:: Python from __future__ import annotations import matplotlib.pyplot as plt import onnx_light.onnx as onnxl import onnx_light.onnx.defs as defs import onnx_light.onnx.helper as helper from onnx_light.tools import pretty_onnx from onnx_light.onnx_optim.expressions import evaluate_expression from onnx_light.onnx_optim.shape_inference import ( ComputeContext, NODE_MEMORY_ALREADY_ALLOCATED_BYTES_KEY, NODE_MEMORY_INPUTS_KEY, NODE_MEMORY_INITIALIZERS_KEY, NODE_MEMORY_INTERMEDIATES_KEY, NODE_MEMORY_OUTPUT_ALLOCATION_BYTES_KEY, NODE_MEMORY_OUTPUTS_KEY, NODE_MEMORY_TOTAL_BYTES_KEY, ShapesContext, apply_inferred_shapes_to_model, compute_shape_model, ) # Built-in operator schemas must be registered before shape inference. defs.register_onnx_operator_set_schema() .. GENERATED FROM PYTHON SOURCE LINES 48-67 Build a graph with one symbolic dimension +++++++++++++++++++++++++++++++++++++++++ The graph keeps the rank fixed but leaves the leading dimension symbolic: .. code-block:: none X : float[N, 4] W : float[4, 4] M = MatMul(X, W) -> float[N, 4] C = Concat(M, X) -> float[2*N, 4] S = Shape(C) -> int64[2] A = Abs(C) -> float[2*N, 4] Z = Reshape(A, S) -> float[2*N, 4] ``W`` contributes constant initializer memory, ``Concat`` turns the symbolic leading dimension into ``2*N``, ``S`` is tagged as a shape tensor, and the last two nodes can reuse their input buffers in place. The memory table therefore mixes constant terms, symbolic terms, and zero-allocation steps. .. GENERATED FROM PYTHON SOURCE LINES 67-96 .. code-block:: Python model = helper.make_model( helper.make_graph( [ helper.make_node("MatMul", ["X", "W"], ["M"]), helper.make_node("Concat", ["M", "X"], ["C"], axis=0), helper.make_node("Shape", ["C"], ["S"]), helper.make_node("Abs", ["C"], ["A"]), helper.make_node("Reshape", ["A", "S"], ["Z"]), ], "compute_context_memory_demo", inputs=[helper.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, ["N", 4])], outputs=[helper.make_tensor_value_info("Z", onnxl.TensorProto.FLOAT, None)], initializer=[ helper.make_tensor( "W", onnxl.TensorProto.FLOAT, [4, 4], [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0], ) ], ), opset_imports=[helper.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_context_memory_demo' input: float[N,4] X init: float[4,4] W 0: MatMul(X, W) -> M 1: Concat(M, X) -> C 2: Shape(C) -> S 3: Abs(C) -> A 4: Reshape(A, S) -> Z output: float[] Z .. GENERATED FROM PYTHON SOURCE LINES 97-103 Run shape inference and memory analysis +++++++++++++++++++++++++++++++++++++++ ``ComputeContext`` consumes the symbolic shapes produced by :class:`ShapesContext`. Passing ``value_tags`` lets the per-source buckets keep semantic labels such as ``shape``. .. GENERATED FROM PYTHON SOURCE LINES 103-145 .. code-block:: Python shape_context = ShapesContext() compute_shape_model(shape_context, model) apply_inferred_shapes_to_model(shape_context, model) compute_context = ComputeContext() value_tags, _ = compute_context.compute_value_and_node_tags(model.graph) compute_context.compute_inplace_reuse_graph(model.graph, shape_context, value_tags=value_tags) memory_profiles = compute_context.memory MemoryScalar = int | str def evaluate_memory_scalar(value: MemoryScalar, assignment: dict[str, int]) -> int: """Evaluates *value* under *assignment*. Returns: The evaluated integer result. """ if isinstance(value, int): return value return evaluate_expression(value, assignment) def format_bucket(bucket: dict[str, MemoryScalar]) -> str: """Formats one tagged memory bucket. Returns: A stable string rendering of the bucket. """ if not bucket: return "-" parts = [] for tag, value in sorted(bucket.items(), key=lambda item: (item[0] != "", item[0])): label = "untagged" if tag == "" else tag parts.append(f"{label}={value}") return ", ".join(parts) .. GENERATED FROM PYTHON SOURCE LINES 146-157 Symbolic per-node memory table ++++++++++++++++++++++++++++++ The table below shows the symbolic profile computed for each node: * ``already_allocated`` is the live memory at node entry, * ``output_allocation`` is the fresh allocation still required for outputs, * ``total`` is their sum. The source buckets make it easy to see which bytes come from live inputs, initializers, intermediates, or newly allocated outputs. .. GENERATED FROM PYTHON SOURCE LINES 157-202 .. code-block:: Python rows = [] for node_index, node in enumerate(model.graph.node): profile = memory_profiles[node_index] rows.append( [ str(node_index), node.op_type, str(profile[NODE_MEMORY_ALREADY_ALLOCATED_BYTES_KEY]), str(profile[NODE_MEMORY_OUTPUT_ALLOCATION_BYTES_KEY]), format_bucket(profile[NODE_MEMORY_INPUTS_KEY]), format_bucket(profile[NODE_MEMORY_INITIALIZERS_KEY]), format_bucket(profile[NODE_MEMORY_INTERMEDIATES_KEY]), format_bucket(profile[NODE_MEMORY_OUTPUTS_KEY]), str(profile[NODE_MEMORY_TOTAL_BYTES_KEY]), ] ) headers = [ "node", "op", "already_allocated", "output_allocation", "inputs", "initializers", "intermediates", "outputs", "total", ] col_widths = [len(h) for h in headers] for row in rows: for i, cell in enumerate(row): if len(cell) > col_widths[i]: col_widths[i] = len(cell) separator = " " + " ".join("-" * w for w in col_widths) header_line = " " + " ".join(h.ljust(col_widths[i]) for i, h in enumerate(headers)) print("Symbolic ComputeContext.memory table:") print(separator) print(header_line) print(separator) for row in rows: print(" " + " ".join(cell.ljust(col_widths[i]) for i, cell in enumerate(row))) print(separator) .. rst-class:: sphx-glr-script-out .. code-block:: none Symbolic ComputeContext.memory table: ---- ------- ----------------- ----------------- ----------- ------------ --------------------- ----------- ------- node op already_allocated output_allocation inputs initializers intermediates outputs total ---- ------- ----------------- ----------------- ----------- ------------ --------------------- ----------- ------- 0 MatMul 16*N+64 16*N weight=16*N weight=64 - weight=16*N 32*N+64 1 Concat 32*N+64 32*N weight=16*N weight=64 weight=16*N weight=32*N 64*N+64 2 Shape 48*N+64 16 weight=16*N weight=64 weight=32*N shape=16 48*N+80 3 Abs 48*N+80 0 weight=16*N weight=64 shape=16, weight=32*N - 48*N+80 4 Reshape 48*N+80 0 weight=16*N weight=64 shape=16, weight=32*N - 48*N+80 ---- ------- ----------------- ----------------- ----------- ------------ --------------------- ----------- ------- .. GENERATED FROM PYTHON SOURCE LINES 203-209 Evaluate the symbolic expressions +++++++++++++++++++++++++++++++++ Once concrete values are chosen for ``N``, the symbolic totals become plain integers. Each line below evaluates the same node-wise ``total_bytes`` curve under a different assignment. .. GENERATED FROM PYTHON SOURCE LINES 209-237 .. code-block:: Python ASSIGNMENTS = [{"N": 1}, {"N": 8}, {"N": 32}, {"N": 128}] node_indices = list(range(len(memory_profiles))) print("\nEvaluated total_bytes per node:") print(f" {'N':>6} " + " ".join(f"node{i:>2}" for i in node_indices)) evaluated_totals: dict[int, list[int]] = {} for assignment in ASSIGNMENTS: n_value = assignment["N"] totals = [ evaluate_memory_scalar(profile[NODE_MEMORY_TOTAL_BYTES_KEY], assignment) for profile in memory_profiles ] evaluated_totals[n_value] = totals print(f" {n_value:>6} " + " ".join(f"{value:>6}" for value in totals)) fig, ax = plt.subplots(figsize=(8, 4.5)) for n_value, totals in evaluated_totals.items(): ax.plot(node_indices, totals, marker="o", linewidth=2, label=f"N={n_value}") ax.set_xticks(node_indices) ax.set_xticklabels([f"{i}:{node.op_type}" for i, node in enumerate(model.graph.node)]) ax.set_xlabel("node index") ax.set_ylabel("total bytes") ax.set_title("Evaluated ComputeContext.total_bytes") ax.grid(True, alpha=0.3) ax.legend() fig.tight_layout() .. image-sg:: /auto_examples_optimization/images/sphx_glr_plot_compute_context_memory_001.png :alt: Evaluated ComputeContext.total_bytes :srcset: /auto_examples_optimization/images/sphx_glr_plot_compute_context_memory_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none Evaluated total_bytes per node: N node 0 node 1 node 2 node 3 node 4 1 96 128 128 128 128 8 320 576 464 464 464 32 1088 2112 1616 1616 1616 128 4160 8256 6224 6224 6224 .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.154 seconds) .. _sphx_glr_download_auto_examples_optimization_plot_compute_context_memory.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_context_memory.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_compute_context_memory.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_compute_context_memory.zip ` .. include:: plot_compute_context_memory.recommendations .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_