Benchmark streaming vs in-memory alignment of external data

This example benchmarks two ways to re-align tensor offsets in an ONNX model that stores its weights as external data:

1. Streaming approach — onnx_light.onnx.align_external_data_streaming() rewrites a (src.onnx, src.data) pair into a new (dst.onnx, dst.data) pair with every tensor offset aligned to a power-of-two boundary, without ever loading the weights in RAM. Only the proto metadata is parsed; tensor bytes are streamed file-to-file (via splice(2) on Linux for zero-copy in kernel space).

2. Usual in-memory approach — onnx_light.onnx.load() followed by SerializeToFile with onnx_light.onnx.SerializeOptions.alignment set to the same boundary. This loads the full set of weights into memory before re-writing them.

The example reports wall-clock time and peak process-resident memory usage for each approach. At equal alignment the two outputs are byte-equivalent for the tensor payloads — only the peak memory footprint differs.

Note

Peak memory is sampled from the process’s resident set size (RSS), not from tracemalloc. tracemalloc only tracks the Python allocator (pymalloc), so it misses the C++-owned buffers where onnx-light stores tensor raw_data — making both approaches appear similarly small and hiding the very overhead this example is meant to illustrate. RSS captures both Python and C++ allocations.

import gc
import os
import shutil
import threading
import time

import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
import pandas

import onnx_light.onnx as onnxl
import onnx_light.onnx.helper as oh
import onnx_light.onnx.numpy_helper as onh

Build a synthetic model

A handful of large initializers is enough to make the in-memory peak noticeably larger than the streaming peak. The example runs with a smaller model under UNITTEST_GOING=1 to keep CI fast.

N_INIT = 8 if os.environ.get("UNITTEST_GOING") == "1" else 24
DIM = 256 if os.environ.get("UNITTEST_GOING") == "1" else 1024
ALIGNMENT = 4096


def make_model(n_init: int = N_INIT, dim: int = DIM) -> onnxl.ModelProto:
    """Creates a synthetic ONNX model with *n_init* float32 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)


out_dir = "temp_plot_align_external_data_streaming"
os.makedirs(out_dir, exist_ok=True)

model = make_model()
total_weight_bytes = sum(int(np.prod(init.dims)) * 4 for init in model.graph.initializer)
print(f"Number of initializers : {len(model.graph.initializer)}")
print(f"Total weight bytes     : {total_weight_bytes / 2 ** 20:.2f} MB")
print(f"Alignment              : {ALIGNMENT} bytes")

Number of initializers : 24
Total weight bytes     : 96.00 MB
Alignment              : 4096 bytes

Save the source two-file model

Every benchmark below starts from the same on-disk source pair (src.onnx, src.data).

src_onnx = os.path.join(out_dir, "src.onnx")
src_data = src_onnx + ".data"


def _save_source(proto: onnxl.ModelProto) -> None:
    """Writes the synthetic model to a two-file ``(.onnx, .data)`` pair."""
    onnxl.save(proto, src_onnx, location=src_data)


_save_source(model)
print(f"Source two-file model : {src_onnx} + {src_data}")
print(f"Source weights size   : {os.path.getsize(src_data) / 2 ** 20:.2f} MB")

# Release the in-memory model now that the source files are on disk.
# Keeping ``model`` alive would inflate the RSS baseline for both benchmarks,
# making the streaming variant appear to use more memory than it actually does.
# After this point all benchmarks start from a clean low-memory baseline.
del model
gc.collect()

Source two-file model : temp_plot_align_external_data_streaming/src.onnx + temp_plot_align_external_data_streaming/src.onnx.data
Source weights size   : 96.00 MB

6

Helpers

Each benchmark is wrapped in a function that returns the elapsed wall time and the peak process-resident memory (RSS) observed while the function runs. RSS is sampled from /proc/self/statm (Linux) on a background thread; on platforms where /proc is unavailable the sampler falls back to resource.getrusage() (Unix) or reports 0. Unlike tracemalloc, RSS includes the C++-owned buffers where onnx-light stores tensor raw_data during an in-memory load — which is exactly the overhead the streaming variant avoids.

_PAGE_SIZE = os.sysconf("SC_PAGE_SIZE") if hasattr(os, "sysconf") else 4096


def _read_rss_bytes() -> int:
    """Returns the current process resident set size in bytes (0 if unknown)."""
    try:
        with open("/proc/self/statm", encoding="ascii") as stream:
            return int(stream.read().split()[1]) * _PAGE_SIZE
    except OSError:
        pass
    try:
        import resource

        ru = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        # Linux reports KiB, macOS reports bytes; both branches are >= our usage scale.
        return ru * 1024 if ru < 2**40 else ru
    except (ImportError, OSError):
        return 0


class _PeakRssSampler:
    """Samples process RSS on a background thread and records the peak observed."""

    def __init__(self, interval: float = 0.001):
        self._interval = interval
        self._stop = threading.Event()
        self._thread: threading.Thread | None = None
        self.baseline = 0
        self.peak = 0

    def __enter__(self):
        self.baseline = _read_rss_bytes()
        self.peak = self.baseline
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()
        return self

    def __exit__(self, *_exc) -> None:
        self._stop.set()
        if self._thread is not None:
            self._thread.join()
        cur = _read_rss_bytes()
        if cur > self.peak:
            self.peak = cur

    def _run(self) -> None:
        while not self._stop.wait(self._interval):
            cur = _read_rss_bytes()
            if cur > self.peak:
                self.peak = cur


def measure(name: str, fn) -> dict:
    """Runs *fn* and returns elapsed time (s) and peak RSS delta (MB)."""
    with _PeakRssSampler() as sampler:
        t0 = time.perf_counter()
        fn()
        elapsed = time.perf_counter() - t0
    peak_mb = max(sampler.peak - sampler.baseline, 0) / 2**20
    print(f"{name:<40} time={elapsed * 1e3:8.1f} ms   peak={peak_mb:7.2f} MB")
    return {"name": name, "time_s": elapsed, "peak_mb": peak_mb}


def _flush(path: str) -> None:
    """fsync a file so timings include write-back to disk."""
    with open(path, "r+b") as stream:
        stream.flush()
        os.fsync(stream.fileno())

Approach 1 — streaming alignment

The source .onnx is parsed with skip_raw_data=True, so only the metadata (a few KB) reaches RAM; tensor bytes are streamed file-to-file in chunks of chunk_size bytes, zero-padding to the requested alignment between tensors.

streaming_onnx = os.path.join(out_dir, "streaming.onnx")
streaming_data = os.path.join(out_dir, "streaming.data")


def run_streaming() -> None:
    onnxl.align_external_data_streaming(
        src_onnx_path=src_onnx,
        dst_onnx_path=streaming_onnx,
        dst_weights_path=streaming_data,
        alignment=ALIGNMENT,
        chunk_size=4 * 1024 * 1024,
    )
    _flush(streaming_data)
    _flush(streaming_onnx)

Approach 2 — in-memory load + re-save with alignment

The “usual” path: load the full model (weights and all) into RAM, then re-serialize it with SerializeOptions.alignment set to the same boundary.

inmem_onnx = os.path.join(out_dir, "inmem.onnx")
inmem_data = os.path.join(out_dir, "inmem.data")


def run_in_memory() -> None:
    loaded = onnxl.load(src_onnx, load_external_data=True, num_threads=1)
    opts = onnxl.SerializeOptions()
    opts.alignment = ALIGNMENT
    # Disable honouring existing external_data.location entries so every
    # tensor is written into the destination weights file passed below.
    opts.use_external_data_location = False
    loaded.SerializeToFile(inmem_onnx, opts, inmem_data)
    _flush(inmem_data)
    _flush(inmem_onnx)

Run the benchmarks

results = []
results.append(measure("align/streaming", run_streaming))
results.append(measure("align/in-memory", run_in_memory))

align/streaming                          time=    71.9 ms   peak=   3.88 MB
align/in-memory                          time=   134.1 ms   peak=  92.00 MB

Verify byte-equivalence of the aligned outputs

Both approaches must produce the same tensor payloads at the same aligned offsets. We compare each tensor’s bytes read from both destination weights files.

streaming_loaded = onnxl.load(streaming_onnx, load_external_data=False)
inmem_loaded = onnxl.load(inmem_onnx, load_external_data=False)

with open(streaming_data, "rb") as f_stream, open(inmem_data, "rb") as f_inmem:
    for s_init, m_init in zip(streaming_loaded.graph.initializer, inmem_loaded.graph.initializer):
        s_meta = {e.key: e.value for e in s_init.external_data}
        m_meta = {e.key: e.value for e in m_init.external_data}
        s_off, s_len = int(s_meta["offset"]), int(s_meta["length"])
        m_off, m_len = int(m_meta["offset"]), int(m_meta["length"])
        assert s_off % ALIGNMENT == 0, f"streaming offset {s_off} not aligned"
        assert m_off % ALIGNMENT == 0, f"in-memory offset {m_off} not aligned"
        assert s_len == m_len, f"length mismatch: {s_len} vs {m_len}"
        f_stream.seek(s_off)
        f_inmem.seek(m_off)
        assert f_stream.read(s_len) == f_inmem.read(m_len), f"payload mismatch on {s_init.name}"
print("Both approaches produced byte-equivalent aligned tensor payloads.")

Both approaches produced byte-equivalent aligned tensor payloads.

Results

df = pandas.DataFrame(results).set_index("name").sort_index()
print(df)

                   time_s    peak_mb
name
align/in-memory  0.134114  92.003906
align/streaming  0.071939   3.882812

Plot

The streaming variant trades a slightly different I/O pattern for an O(metadata) + chunk_size memory footprint, independent of the total weights size. The in-memory variant peaks proportionally to the model’s weight payload.

fig, (ax_time, ax_mem) = plt.subplots(1, 2, figsize=(12, 5))

row_names = df.index.tolist()
colors = ["darkorange" if "streaming" in n else "steelblue" for n in row_names]

ax_time.barh(row_names, df["time_s"], color=colors)
ax_time.set_title("Wall time (s, lower is better)")
ax_time.set_xlabel("seconds")
ax_time.grid(axis="x")

ax_mem.barh(row_names, df["peak_mb"], color=colors)
ax_mem.set_title("Peak RSS delta (MB, lower is better)")
ax_mem.set_xlabel("MB (resident set size)")
ax_mem.grid(axis="x")

fig.suptitle(
    f"Alignment={ALIGNMENT}B   weights={total_weight_bytes / 2 ** 20:.1f} MB"
    f"   #initializers={N_INIT}"
)
fig.legend(
    handles=[
        mpatches.Patch(color="darkorange", label="streaming"),
        mpatches.Patch(color="steelblue", label="in-memory"),
    ],
    loc="lower center",
    ncol=2,
)
fig.tight_layout(rect=(0, 0.05, 1, 1))
fig.savefig("plot_align_external_data_streaming.png")

Cleanup

shutil.rmtree(out_dir, ignore_errors=True)

Related examples

Number of threads used to load and save ONNX models

Number of threads used to load and save ONNX models

Measures loading and saving time for an ONNX model

Measures loading and saving time for an ONNX model

Profiles ONNX external-data save time

Profiles ONNX external-data save time