Note
Go to the end to download the full example code.
Profiles ONNX external-data save time#
This example profiles how long it takes to save a model with external data
using onnx and onnx_light.onnx.
It follows the same benchmark style as Measures loading and saving time for an ONNX model but focuses only on the external-data save scenario.
import os
import shutil
import cProfile
import pstats
import matplotlib.patches as mpatches
import numpy as np
import onnx
import onnx.helper as oh
import onnx.numpy_helper as onh
import pandas
import onnx_light.onnx as onnxl
N_INIT = 40
DIM = 256 if os.environ.get("UNITTEST_GOING") == "1" else 3072
def make_model(n_init: int = N_INIT, dim: int = DIM) -> onnx.ModelProto:
"""Creates a synthetic ONNX model with large initializers."""
initializers = []
nodes = []
inputs = [oh.make_tensor_value_info("X", onnx.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, onnx.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 profile_call(name: str, fn, repeat=1) -> dict:
"""Profiles the given callable with cProfile.
Args:
name: Benchmark name used in printed output and the result row.
fn: Callable to execute under cProfile.
Returns:
A dictionary with the benchmark name and total profiled time in seconds.
"""
profiler = cProfile.Profile()
for _ in range(repeat):
profiler.runcall(fn)
profile_stats = pstats.Stats(profiler).sort_stats("cumulative")
print(f"\n{name}\n{'-' * len(name)}")
profile_stats.print_stats(20)
return {"name": name, "total": float(profile_stats.total_tt)}
def _flush_file(path: str) -> None:
"""Flushes one file descriptor so benchmark timing includes write-back."""
with open(path, "r+b") as stream:
stream.flush()
os.fsync(stream.fileno())
model = make_model()
size_bytes = model.ByteSize()
print(f"Model size: {size_bytes / 2 ** 20:.3f} MB")
out_dir = "temp_plot_save_external_data_time"
os.makedirs(out_dir, exist_ok=True)
onnx_model = model
onnx_input_path = os.path.join(out_dir, "bench.onnx")
onnx.save(onnx_model, onnx_input_path)
onnx_light_model = onnxl.load(onnx_input_path)
results = []
# ``onnx.save_model(..., save_as_external_data=True)`` mutates the in-memory
# model by replacing ``raw_data`` with external-data metadata. Benchmark it as a
# single-shot operation so the row reflects the full conversion + write cost
# instead of re-saving an already externalized model on later iterations.
# Both saved files are explicitly ``fsync``-ed so this row includes descriptor
# flush/write-back overhead, matching the ``onnxlight`` row.
onnx_external_path = os.path.join(out_dir, "out_onnx_ext.onnx")
onnx_external_location = "out_onnx_ext.data"
onnx_external_data_path = os.path.join(out_dir, onnx_external_location)
def _save_onnx_external_with_flush() -> None:
onnx.save_model(
onnx_model,
onnx_external_path,
save_as_external_data=True,
all_tensors_to_one_file=True,
location=onnx_external_location,
)
_flush_file(onnx_external_data_path)
_flush_file(onnx_external_path)
results.append(profile_call("save/2filex1/onnx", _save_onnx_external_with_flush, repeat=1))
print(f"{results[-1]['name']:<35} total={results[-1]['total'] * 1e3:.1f} ms")
# ``onnx_light.onnx.save`` restores the in-memory model after the write, but we
# keep the benchmark single-shot so the rows stay directly comparable.
onnx_light_external_path = os.path.join(out_dir, "out_onnxlight_ext.onnx")
onnx_light_external_data = onnx_light_external_path + ".data"
def _save_onnxlight_external_with_flush() -> None:
onnxl.save(onnx_light_model, onnx_light_external_path, location=onnx_light_external_data)
_flush_file(onnx_light_external_data)
_flush_file(onnx_light_external_path)
results.append(
profile_call("save/2filex1/onnxlight", _save_onnxlight_external_with_flush, repeat=1)
)
print(f"{results[-1]['name']:<35} total={results[-1]['total'] * 1e3:.1f} ms")
onnx_light_external_x4_path = os.path.join(out_dir, "out_onnxlight_ext_x4.onnx")
onnx_light_external_x4_data = onnx_light_external_x4_path + ".data"
results.append(
profile_call(
"save/2filex4/onnxlight",
lambda: onnxl.save(
onnx_light_model,
onnx_light_external_x4_path,
location=onnx_light_external_x4_data,
num_threads=4,
),
repeat=1,
)
)
print(f"{results[-1]['name']:<35} total={results[-1]['total'] * 1e3:.1f} ms")
Model size: 1440.002 MB
save/2filex1/onnx
-----------------
2761 function calls in 5.727 seconds
Ordered by: cumulative time
List reduced from 155 to 20 due to restriction <20>
ncalls tottime percall cumtime percall filename:lineno(function)
1 0.000 0.000 5.727 5.727 ~/github/onnx-light/docs/examples/core/plot_save_external_data_time.py:103(_save_onnx_external_with_flush)
1 0.000 0.000 3.275 3.275 ~/github/onnx/onnx/__init__.py:299(save_model)
2 0.000 0.000 2.452 1.226 ~/github/onnx-light/docs/examples/core/plot_save_external_data_time.py:71(_flush_file)
2 2.451 1.226 2.451 1.226 {built-in method posix.fsync}
1 0.001 0.001 2.295 2.295 ~/github/onnx/onnx/external_data_helper.py:393(write_external_data_tensors)
40 1.097 0.027 2.294 0.057 ~/github/onnx/onnx/external_data_helper.py:278(save_external_data)
40 1.181 0.030 1.181 0.030 {method 'write' of '_io.BufferedRandom' objects}
1 0.961 0.961 0.979 0.979 ~/github/onnx/onnx/external_data_helper.py:204(convert_model_to_external_data)
1 0.000 0.000 0.015 0.015 /usr/lib/python3.12/uuid.py:674(uuid1)
1 0.000 0.000 0.015 0.015 /usr/lib/python3.12/uuid.py:650(getnode)
1 0.000 0.000 0.015 0.015 /usr/lib/python3.12/uuid.py:521(_ip_getnode)
1 0.000 0.000 0.015 0.015 /usr/lib/python3.12/uuid.py:410(_find_mac_near_keyword)
1 0.000 0.000 0.015 0.015 /usr/lib/python3.12/uuid.py:360(_get_command_stdout)
1 0.000 0.000 0.011 0.011 /usr/lib/python3.12/subprocess.py:1165(communicate)
1 0.011 0.011 0.011 0.011 {method 'read' of '_io.BufferedReader' objects}
43 0.005 0.000 0.005 0.000 {method '__exit__' of '_io._IOBase' objects}
80 0.004 0.000 0.005 0.000 ~/github/onnx/onnx/external_data_helper.py:176(set_external_data)
40 0.005 0.000 0.005 0.000 ~/github/onnx/onnx/external_data_helper.py:28(_open_external_data_fd)
1 0.000 0.000 0.003 0.003 /usr/lib/python3.12/subprocess.py:807(__init__)
1 0.000 0.000 0.003 0.003 /usr/lib/python3.12/subprocess.py:1791(_execute_child)
save/2filex1/onnx total=5727.1 ms
save/2filex1/onnxlight
----------------------
17 function calls in 2.390 seconds
Ordered by: cumulative time
ncalls tottime percall cumtime percall filename:lineno(function)
1 0.000 0.000 2.389 2.389 ~/github/onnx-light/docs/examples/core/plot_save_external_data_time.py:124(_save_onnxlight_external_with_flush)
2 0.000 0.000 1.297 0.648 ~/github/onnx-light/docs/examples/core/plot_save_external_data_time.py:71(_flush_file)
2 1.296 0.648 1.296 0.648 {built-in method posix.fsync}
1 1.093 1.093 1.093 1.093 ~/github/onnx-light/onnx_light/onnx_lib/io_helper.py:49(save)
2 0.000 0.000 0.000 0.000 {built-in method _io.open}
2 0.000 0.000 0.000 0.000 {method '__exit__' of '_io._IOBase' objects}
1 0.000 0.000 0.000 0.000 {method 'disable' of '_lsprof.Profiler' objects}
2 0.000 0.000 0.000 0.000 {method 'flush' of '_io.BufferedRandom' objects}
2 0.000 0.000 0.000 0.000 {method 'fileno' of '_io.BufferedRandom' objects}
2 0.000 0.000 0.000 0.000 {built-in method builtins.isinstance}
save/2filex1/onnxlight total=2389.6 ms
save/2filex4/onnxlight
----------------------
5 function calls in 1.662 seconds
Ordered by: cumulative time
ncalls tottime percall cumtime percall filename:lineno(function)
1 0.000 0.000 1.662 1.662 ~/github/onnx-light/docs/examples/core/plot_save_external_data_time.py:140(<lambda>)
1 1.662 1.662 1.662 1.662 ~/github/onnx-light/onnx_light/onnx_lib/io_helper.py:49(save)
1 0.000 0.000 0.000 0.000 {method 'disable' of '_lsprof.Profiler' objects}
2 0.000 0.000 0.000 0.000 {built-in method builtins.isinstance}
save/2filex4/onnxlight total=1661.7 ms
Results#
df = pandas.DataFrame(results).set_index("name").sort_index()
print(df)
total
name
save/2filex1/onnx 5.727098
save/2filex1/onnxlight 2.389581
save/2filex4/onnxlight 1.661712
Plot#
ax = df[["total"]].plot.barh(
title=f"size={size_bytes / 2 ** 20:.2f} MB\nexternal-data save (s)\nlower is better",
xlabel="seconds",
legend=False,
figsize=(12, 6),
)
row_names = df.index.tolist()
for container in ax.containers:
for bar, name in zip(container, row_names):
bar.set_facecolor("darkorange" if "onnxlight" in name else "steelblue")
ax.legend(
handles=[
mpatches.Patch(color="steelblue", label="onnx"),
mpatches.Patch(color="darkorange", label="onnxlight"),
]
)
ax.grid(axis="x")
ax.figure.tight_layout()
ax.figure.savefig("plot_save_external_data_time.png")
Cleanup#
shutil.rmtree(out_dir, ignore_errors=True)
Total running time of the script: (0 minutes 41.486 seconds)