.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples_core/plot_threads_load_save.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_core_plot_threads_load_save.py: .. _l-example-plot-threads-load-save: Number of threads used to load and save ONNX models ===================================================== This example benchmarks :func:`onnx_light.onnx.load` and :func:`onnx_light.onnx.save` while varying the ``num_threads`` parameter across three on-disk layouts: * **single-file** – everything packed in one ``.onnx`` file (no external data). * **2 files** – one ``.onnx`` file plus one external weights file. * **multi-file** – one ``.onnx`` file plus several external weights files, produced by capping each file with ``max_external_file_size``. For each (layout, num_threads) combination the script measures the median wall-clock time over a few iterations and plots a summary chart so the effect of parallelism can be compared at a glance. .. GENERATED FROM PYTHON SOURCE LINES 20-123 .. code-block:: Python import os import platform import shutil import time import matplotlib.pyplot as plt import numpy as np import pandas import onnx_light.onnx.helper as oh import onnx_light.onnx.numpy_helper as onh import onnx_light.onnx as onnxl N_INIT = 8 if os.environ.get("UNITTEST_GOING") == "1" else 40 DIM = 128 if os.environ.get("UNITTEST_GOING") == "1" else 2048 N_ITER = 2 if os.environ.get("UNITTEST_GOING") == "1" else 5 THREAD_COUNTS = (1, 2) if os.environ.get("UNITTEST_GOING") == "1" else (1, 2, 3, 4, 8) def _detect_processor_name() -> str: """Returns a human-readable processor name across common platforms.""" # ``platform.processor()`` is often empty on Linux; fall back to # ``/proc/cpuinfo`` so the chart title shows something meaningful. name = platform.processor() or "" if not name and os.path.exists("/proc/cpuinfo"): try: with open("/proc/cpuinfo", encoding="utf-8") as f: for line in f: if line.startswith("model name"): name = line.split(":", 1)[1].strip() break except OSError: pass return name or platform.machine() or "unknown" def _detect_physical_cores() -> int: """Returns the number of physical CPU cores, or 0 if it cannot be determined.""" # On Linux, count unique (physical id, core id) pairs in ``/proc/cpuinfo``. if os.path.exists("/proc/cpuinfo"): try: cores: set[tuple[str, str]] = set() physical_id = "" core_id = "" with open("/proc/cpuinfo", encoding="utf-8") as f: for line in f: if line.startswith("physical id"): physical_id = line.split(":", 1)[1].strip() elif line.startswith("core id"): core_id = line.split(":", 1)[1].strip() elif line.strip() == "": if physical_id and core_id: cores.add((physical_id, core_id)) physical_id = "" core_id = "" if physical_id and core_id: cores.add((physical_id, core_id)) if cores: return len(cores) except OSError: pass return 0 CPU_COUNT = os.cpu_count() or 1 PHYSICAL_CORE_COUNT = _detect_physical_cores() PROCESSOR_NAME = _detect_processor_name() print(f"Processor: {PROCESSOR_NAME}") print(f"Logical cores: {CPU_COUNT}") print(f"Physical cores: {PHYSICAL_CORE_COUNT or 'unknown'}") def make_model(n_init: int = N_INIT, dim: int = DIM) -> onnxl.ModelProto: """Builds a synthetic ONNX model with *n_init* dense ``Gemm`` weights.""" initializers = [] nodes = [] inputs = [oh.make_tensor_value_info("X", onnxl.TensorProto.FLOAT, [None, dim])] prev = "X" for i in range(n_init): weight_name = f"W{i}" out_name = f"Y{i}" w = np.random.randn(dim, dim).astype(np.float32) initializers.append(onh.from_array(w, name=weight_name)) nodes.append(oh.make_node("Gemm", [prev, weight_name], [out_name], transB=1)) prev = out_name outputs = [oh.make_tensor_value_info(prev, onnxl.TensorProto.FLOAT, [None, dim])] graph = oh.make_graph(nodes, "bench_graph", inputs, outputs, initializer=initializers) return oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)], ir_version=9) def median_time(fn, n_iter: int = N_ITER) -> float: """Runs *fn* *n_iter* times and returns the median wall-clock time (seconds).""" timings = [] for _ in range(n_iter): begin = time.perf_counter() fn() timings.append(time.perf_counter() - begin) return float(np.median(timings)) .. rst-class:: sphx-glr-script-out .. code-block:: none Processor: x86_64 Logical cores: 4 Physical cores: 2 .. GENERATED FROM PYTHON SOURCE LINES 124-130 Build the model and prepare on-disk layouts -------------------------------------------- Three reference files are produced once and then reused for the load benchmarks: a single packed file, a model with one external data file, and a model whose external data is split into several files. .. GENERATED FROM PYTHON SOURCE LINES 130-165 .. code-block:: Python model = make_model() size_mb = model.ByteSize() / 2**20 print(f"Model size: {size_mb:.2f} MB ({N_INIT} initializers, dim={DIM})") out_dir = "temp_plot_threads_load_save" os.makedirs(out_dir, exist_ok=True) single_path = os.path.join(out_dir, "model_single.onnx") two_path = os.path.join(out_dir, "model_two.onnx") two_data = two_path + ".data" multi_path = os.path.join(out_dir, "model_multi.onnx") multi_data = multi_path + ".data" # Build each reference file from the same in-memory ModelProto. We load # the freshly written single-file model back through onnx_light so the # external-data writers operate on an ``onnxl.ModelProto``. onnxl.save(model, single_path) onnxl_model = onnxl.load(single_path) onnxl.save(onnxl_model, two_path, location=two_data) # Cap each external weights file at roughly a third of the total payload # so several files are produced regardless of the chosen DIM. total_bytes = sum(len(init.raw_data) if init.raw_data else 0 for init in model.graph.initializer) max_file = max(total_bytes // 3, 1) onnxl.save(onnxl_model, multi_path, location=multi_data, max_external_file_size=max_file) multi_files = sorted(p for p in os.listdir(out_dir) if p.startswith("model_multi.onnx.data")) print( f"Layouts ready: single ({os.path.getsize(single_path) / 2 ** 20:.2f} MB), " f"2-file (graph={os.path.getsize(two_path) / 1024:.1f} KB, " f"data={os.path.getsize(two_data) / 2 ** 20:.2f} MB), " f"multi-file ({len(multi_files)} data files)" ) .. rst-class:: sphx-glr-script-out .. code-block:: none Model size: 640.00 MB (40 initializers, dim=2048) Layouts ready: single (640.00 MB), 2-file (graph=5.0 KB, data=640.00 MB), multi-file (4 data files) .. GENERATED FROM PYTHON SOURCE LINES 166-172 Benchmark --------- For each layout and thread count we time both the load operation (reading from disk into an ``onnxl.ModelProto``) and the save operation (serializing the in-memory model back to a fresh path). .. GENERATED FROM PYTHON SOURCE LINES 172-232 .. code-block:: Python LAYOUTS = ( ("single-file", single_path, None), ("2 files", two_path, two_data), ("multi-file", multi_path, multi_data), ) rows = [] for layout_name, model_path, data_path in LAYOUTS: for num_threads in THREAD_COUNTS: def _load(path=model_path, threads=num_threads): onnxl.load(path, num_threads=threads, load_external_data=True) load_t = median_time(_load) save_out = os.path.join( out_dir, f"out_{layout_name.replace(' ', '_')}_t{num_threads}.onnx" ) save_data = save_out + ".data" if data_path is not None else None if data_path is None: def _save(out=save_out, threads=num_threads): onnxl.save(onnxl_model, out, num_threads=threads) elif layout_name == "multi-file": def _save(out=save_out, data=save_data, threads=num_threads, cap=max_file): onnxl.save( onnxl_model, out, location=data, num_threads=threads, max_external_file_size=cap, ) else: def _save(out=save_out, data=save_data, threads=num_threads): onnxl.save(onnxl_model, out, location=data, num_threads=threads) save_t = median_time(_save) rows.append( { "layout": layout_name, "num_threads": num_threads, "load_ms": load_t * 1e3, "save_ms": save_t * 1e3, } ) print( f"{layout_name:<11} threads={num_threads:<2} " f"load={load_t * 1e3:7.2f} ms save={save_t * 1e3:7.2f} ms" ) df = pandas.DataFrame(rows) print(df) .. rst-class:: sphx-glr-script-out .. code-block:: none single-file threads=1 load= 51.10 ms save=1602.07 ms single-file threads=2 load= 41.19 ms save= 222.20 ms single-file threads=3 load= 42.93 ms save= 213.61 ms single-file threads=4 load= 41.25 ms save= 380.08 ms single-file threads=8 load= 41.98 ms save= 438.32 ms 2 files threads=1 load= 55.39 ms save=1689.40 ms 2 files threads=2 load= 40.51 ms save= 190.55 ms 2 files threads=3 load= 41.07 ms save= 216.75 ms 2 files threads=4 load= 41.16 ms save= 768.13 ms 2 files threads=8 load= 41.63 ms save= 219.40 ms multi-file threads=1 load= 54.30 ms save=1747.01 ms multi-file threads=2 load= 39.71 ms save= 189.10 ms multi-file threads=3 load= 41.93 ms save= 199.80 ms multi-file threads=4 load= 41.63 ms save= 661.23 ms multi-file threads=8 load= 41.52 ms save= 192.91 ms layout num_threads load_ms save_ms 0 single-file 1 51.103396 1602.070349 1 single-file 2 41.190842 222.201847 2 single-file 3 42.927992 213.607195 3 single-file 4 41.248148 380.077316 4 single-file 8 41.981534 438.321364 5 2 files 1 55.391983 1689.399148 6 2 files 2 40.512599 190.548010 7 2 files 3 41.074210 216.749998 8 2 files 4 41.162067 768.133596 9 2 files 8 41.630794 219.396009 10 multi-file 1 54.304805 1747.013773 11 multi-file 2 39.710220 189.102314 12 multi-file 3 41.930013 199.801137 13 multi-file 4 41.625040 661.233498 14 multi-file 8 41.524388 192.906252 .. GENERATED FROM PYTHON SOURCE LINES 233-240 Conclusion plot --------------- The two panels below summarise the impact of ``num_threads`` on each layout. Lines that flatten quickly indicate operations that are already I/O bound and gain little from additional threads, while a steady downward slope reveals where parallelism pays off. .. GENERATED FROM PYTHON SOURCE LINES 240-268 .. code-block:: Python load_pivot = df.pivot(index="num_threads", columns="layout", values="load_ms") save_pivot = df.pivot(index="num_threads", columns="layout", values="save_ms") fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharex=True) for ax, pivot, title in ( (axes[0], load_pivot, "load time (ms)"), (axes[1], save_pivot, "save time (ms)"), ): for layout in pivot.columns: ax.plot(pivot.index, pivot[layout], marker="o", label=layout) ax.set_title(title) ax.set_xlabel("num_threads") ax.set_ylabel("time (ms, lower is better)") ax.set_xlim(left=0) ax.set_xticks(list(THREAD_COUNTS)) ax.grid(True, linestyle=":") ax.legend(title="layout") fig.suptitle( f"onnx_light load/save vs num_threads — model {size_mb:.1f} MB, " f"{N_INIT} initializers\n{PROCESSOR_NAME} " f"({CPU_COUNT} logical cores, " f"{PHYSICAL_CORE_COUNT or 'unknown'} physical cores)" ) fig.tight_layout() fig.savefig("plot_threads_load_save.png") .. image-sg:: /auto_examples_core/images/sphx_glr_plot_threads_load_save_001.png :alt: onnx_light load/save vs num_threads — model 640.0 MB, 40 initializers x86_64 (4 logical cores, 2 physical cores), load time (ms), save time (ms) :srcset: /auto_examples_core/images/sphx_glr_plot_threads_load_save_001.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 269-271 Cleanup ------- .. GENERATED FROM PYTHON SOURCE LINES 271-273 .. code-block:: Python shutil.rmtree(out_dir, ignore_errors=True) .. rst-class:: sphx-glr-timing **Total running time of the script:** (1 minutes 0.358 seconds) .. _sphx_glr_download_auto_examples_core_plot_threads_load_save.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_threads_load_save.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_threads_load_save.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_threads_load_save.zip ` .. include:: plot_threads_load_save.recommendations .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_