.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples_runtime/plot_runtime_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_runtime_plot_runtime_memory.py: .. _l-example-plot-runtime-memory: Profile the runtime memory of a model with the event log ======================================================== When a :class:`RuntimeContext` runs a model with event logging enabled, it records a :class:`RuntimeEvent` for every tensor map mutation and every node dispatch. Attaching a :class:`SimpleRawBufferAllocator` to the context makes each event additionally carry the allocator's *live* footprint (``allocated_bytes``, the total size of every buffer alive at that moment) and its *peak* footprint (``peak_bytes``, the maximum ever reached). Together with the wall-clock ``duration_ns`` of each node this turns the event log into a per-node time-and-memory profile. This example: * builds a small graph that materialises two intermediate activations (``Mul`` then ``Add``) so the allocator's live memory grows and shrinks as nodes run, * attaches a :class:`SimpleRawBufferAllocator` to a :class:`~onnx_light.reference.ReferenceEvaluator` and runs it with ``events_enabled=True`` and ``release_intermediates=True``, * prints :meth:`RuntimeEvent.summary` for every recorded event — a one-line recap including the live and peak memory, * isolates the per-node ``run_node`` events to read the duration and the allocator memory observed right after each node, * renders a compact table of the memory profile. .. GENERATED FROM PYTHON SOURCE LINES 30-40 .. code-block:: Python from __future__ import annotations import numpy as np from onnx_light.onnx_lib import parser from onnx_light.onnx.reference import ReferenceEvaluator from onnx_light.onnx_py._onnxpykernels import runtime as _runtime from onnx_light.tools import pretty_onnx .. GENERATED FROM PYTHON SOURCE LINES 41-47 Build a small ONNX model ++++++++++++++++++++++++ The graph computes ``w = (x * two) + x``. Both ``z = x * two`` and ``w = z + x`` are intermediate activations the runtime has to hold in memory while the graph runs. .. GENERATED FROM PYTHON SOURCE LINES 47-59 .. code-block:: Python model = parser.parse_model( '' "agraph (float[N] x) => (float[N] w) " "" "{" " z = Mul(x, two)" " w = Add(z, x)" "}" ) print(pretty_onnx(model)) .. rst-class:: sphx-glr-script-out .. code-block:: none opset: domain='' version=18 graph: name='agraph' input: float[N] x init: float[] two 0: Mul(x, two) -> z 1: Add(z, x) -> w output: float[N] w .. GENERATED FROM PYTHON SOURCE LINES 60-71 Attach an allocator and run with event logging ++++++++++++++++++++++++++++++++++++++++++++++ :class:`SimpleRawBufferAllocator` is a fixed-capacity pool: ``capacity`` is the maximum number of buffers that may be alive at the same time (a generous upper bound is fine — empty slots are cheap). Passing it to :class:`~onnx_light.reference.ReferenceEvaluator` routes the runtime's buffer storage through it, so every recorded event reports the live and peak memory. ``release_intermediates=True`` frees each intermediate as soon as its last consumer has run, which is exactly what makes the live footprint go down again in the log. .. GENERATED FROM PYTHON SOURCE LINES 71-81 .. code-block:: Python allocator = _runtime.SimpleRawBufferAllocator(64) sess = ReferenceEvaluator( model, events_enabled=True, release_intermediates=True, allocator=allocator ) x = np.arange(1024, dtype=np.float32) (w,) = sess.run(None, {"x": x}) print(f"w[:4] = {w[:4]}") .. rst-class:: sphx-glr-script-out .. code-block:: none w[:4] = [0. 3. 6. 9.] .. GENERATED FROM PYTHON SOURCE LINES 82-88 Read the event log ++++++++++++++++++ :meth:`RuntimeEvent.summary` renders each event as a single line. The ``mem=`` field is the allocator's live footprint right after the action and ``peak=`` is the largest footprint reached so far. .. GENERATED FROM PYTHON SOURCE LINES 88-94 .. code-block:: Python events = sess.events() print(f"Recorded {len(events)} event(s):") for ev in events: print(f" {ev.summary()}") .. rst-class:: sphx-glr-script-out .. code-block:: none Recorded 7 event(s): [add/input] 'x' mem=0B peak=0B [add/initializer] 'two' mem=4B peak=4B [add/intermediate] 'z' node#0 mem=4100B peak=4100B [run_node/unknown] ai.onnx::Mul(x, two) node#0 took 22484ns executor#4/2 mem=4100B peak=4100B [add/intermediate] 'w' node#1 mem=4100B peak=4100B [run_node/unknown] ai.onnx::Add(z, x) node#1 took 5328ns executor#4/2 mem=4100B peak=4100B [remove/unknown] 'z' mem=4B peak=4100B .. GENERATED FROM PYTHON SOURCE LINES 95-101 Isolate the per-node memory profile +++++++++++++++++++++++++++++++++++ The ``run_node`` events summarise the dispatch of each node: its ``op_type``, the ``duration_ns`` it took and the allocator memory observed right after it ran. Filtering the log by action recovers just those entries. .. GENERATED FROM PYTHON SOURCE LINES 101-111 .. code-block:: Python run_nodes = [ev for ev in events if ev.action == _runtime.RuntimeEventAction.kRunNode] print("Per-node time and memory profile:") for ev in run_nodes: d = ev.as_dict() print( f" {d['op_type']:<8s} took {d['duration_ns']:>8d} ns " f"live={d['allocated_bytes']:>6d} B peak={d['peak_bytes']:>6d} B" ) .. rst-class:: sphx-glr-script-out .. code-block:: none Per-node time and memory profile: Mul took 22484 ns live= 4100 B peak= 4100 B Add took 5328 ns live= 4100 B peak= 4100 B .. GENERATED FROM PYTHON SOURCE LINES 112-117 The allocator also exposes the peak directly +++++++++++++++++++++++++++++++++++++++++++++ The same peak is available on the allocator object itself, independently of the event log. .. GENERATED FROM PYTHON SOURCE LINES 117-120 .. code-block:: Python print(f"Allocator peak footprint: {allocator.peak_allocated_size} bytes") .. rst-class:: sphx-glr-script-out .. code-block:: none Allocator peak footprint: 4100 bytes .. GENERATED FROM PYTHON SOURCE LINES 121-127 Render the memory profile as a table ++++++++++++++++++++++++++++++++++++ A simple matplotlib figure doubles as the sphinx-gallery thumbnail and a compact recap of the recorded events, showing the live and peak memory next to each action. .. GENERATED FROM PYTHON SOURCE LINES 127-159 .. code-block:: Python import matplotlib.pyplot as plt # noqa: E402 rows = [] for ev in events: d = ev.as_dict() label = d["op_type"] if d["action"] == "run_node" else d["name"] rows.append( [ d["action"], label, str(d["duration_ns"]) if d["action"] == "run_node" else "", str(d["allocated_bytes"]), str(d["peak_bytes"]), ] ) fig, ax = plt.subplots(figsize=(8, 1.6 + 0.3 * len(rows))) ax.set_axis_off() table = ax.table( cellText=rows, colLabels=["action", "name / op", "duration (ns)", "live (B)", "peak (B)"], loc="center", cellLoc="left", colLoc="left", ) table.auto_set_font_size(False) table.set_fontsize(9) table.scale(1.0, 1.3) ax.set_title("RuntimeContext memory profile") fig.tight_layout() fig.savefig("plot_runtime_memory.png") .. image-sg:: /auto_examples_runtime/images/sphx_glr_plot_runtime_memory_001.png :alt: RuntimeContext memory profile :srcset: /auto_examples_runtime/images/sphx_glr_plot_runtime_memory_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 0.116 seconds) .. _sphx_glr_download_auto_examples_runtime_plot_runtime_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_runtime_memory.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_runtime_memory.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_runtime_memory.zip ` .. include:: plot_runtime_memory.recommendations .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_