Note
Go to the end to download the full example code.
Benchmark TreeEnsemble scheduling scenarios#
This example measures TreeEnsemble-5 regression forests while varying the
number of trees, input features, and batch size. Batch size 1 is represented
explicitly because ONNX Runtime uses a specialized execution path for that
case. The largest forest contains 10,000 trees and the widest input contains
4,096 features, both representative of production models. Those largest
configurations are enabled by --big, which defaults to True on machines
with more than 64 cores and can be disabled with --no-big.
The example is standalone: it builds every model, runs ONNX Runtime and onnx-light + onnx-light-cpu, and measures both without relying on any script from the repository.
import argparse
import gc
import os
import statistics
import time
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap, LogNorm, TwoSlopeNorm
import numpy as np
import onnxruntime
# ``onnx-light`` ships ``onnx_light.onnx`` as a drop-in replacement for the
# ``onnx`` package; use it to build the models so the example depends on
# onnx-light rather than onnx.
from onnx_light.onnx import TensorProto, checker, helper, numpy_helper
from onnx_light.onnx.reference import ReferenceEvaluator
from onnx_light_cpu import register_kernels
cpu_count = os.cpu_count() or 1
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("-r", "--repeat", type=int, default=10 * cpu_count)
parser.add_argument("-w", "--warmup", type=int, default=2 * cpu_count)
parser.add_argument("-t", "--max-repeat-time", type=float, default=1.0)
parser.add_argument(
"--big",
action=argparse.BooleanOptionalAction,
default=cpu_count > 64,
help=(
"include configurations with 10,000 trees or 4,096 features, "
"enabled by default on machines with more than 64 cores"
),
)
args, _ = parser.parse_known_args()
if args.repeat <= 0:
parser.error("--repeat must be greater than 0")
if args.warmup < 0:
parser.error("--warmup must be greater than or equal to 0")
if args.max_repeat_time <= 0:
parser.error("--max-repeat-time must be greater than 0")
if os.environ.get("UNITTEST_GOING"):
tree_counts = (10, 100)
feature_counts = (4, 64)
batch_sizes = (1, 32)
else:
tree_counts = (10, 100, 1000, 10000) if args.big else (10, 100, 1000)
feature_counts = (4, 16, 64, 256, 1024, 4096) if args.big else (4, 16, 64, 256, 1024)
batch_sizes = (1, 8, 32, 128)
Build one model per configuration#
Every configuration is a TreeEnsemble regression forest of depth 4 with a
single output, built directly here so the example runs on its own.
DEPTH = 4
def make_case(trees, features, rows, seed):
internal_count = (1 << DEPTH) - 1
leaf_count = 1 << DEPTH
nodes_featureids = []
nodes_splits = []
true_ids = []
false_ids = []
true_leafs = []
false_leafs = []
leaf_weights = []
for tree in range(trees):
node_offset = tree * internal_count
leaf_offset = tree * leaf_count
for local_node in range(internal_count):
nodes_featureids.append((tree + local_node) % features)
nodes_splits.append(((local_node % 7) - 3) * 0.25)
for child, ids, leafs in (
(2 * local_node + 1, true_ids, true_leafs),
(2 * local_node + 2, false_ids, false_leafs),
):
if child >= internal_count:
ids.append(leaf_offset + child - internal_count)
leafs.append(1)
else:
ids.append(node_offset + child)
leafs.append(0)
for leaf in range(leaf_count):
leaf_weights.append(((leaf + tree) % 11 - 5) / (max(trees, 1) * 8.0))
node = helper.make_node(
"TreeEnsemble",
["X"],
["scores"],
domain="ai.onnx.ml",
tree_roots=[tree * internal_count for tree in range(trees)],
nodes_featureids=nodes_featureids,
nodes_splits=numpy_helper.from_array(np.asarray(nodes_splits, dtype=np.float32)),
nodes_modes=numpy_helper.from_array(np.zeros(len(nodes_splits), dtype=np.uint8)),
nodes_truenodeids=true_ids,
nodes_falsenodeids=false_ids,
nodes_trueleafs=true_leafs,
nodes_falseleafs=false_leafs,
leaf_targetids=[0] * len(leaf_weights),
leaf_weights=numpy_helper.from_array(np.asarray(leaf_weights, dtype=np.float32)),
n_targets=1,
aggregate_function=1,
post_transform=0,
)
graph = helper.make_graph(
[node],
f"tree_grid_t{trees}_f{features}_b{rows}",
[helper.make_tensor_value_info("X", TensorProto.FLOAT, [rows, features])],
[helper.make_tensor_value_info("scores", TensorProto.FLOAT, [rows, 1])],
)
model = helper.make_model(
graph,
opset_imports=[helper.make_opsetid("", 18), helper.make_opsetid("ai.onnx.ml", 5)],
ir_version=13,
)
checker.check_model(model)
rng = np.random.default_rng(seed)
feeds = {"X": rng.standard_normal((rows, features)).astype(np.float32)}
return model.SerializeToString(), feeds
Measure both runtimes#
measure returns the median duration of a callable, stopping early once
--max-repeat-time seconds have been spent.
def measure(function, repeat, warmup, max_duration):
spent = 0.0
for _ in range(warmup):
start = time.perf_counter_ns()
function()
spent += (time.perf_counter_ns() - start) / 1e9
if spent >= max_duration:
break
timings = []
spent = 0.0
gc_enabled = gc.isenabled()
gc.disable()
try:
for _ in range(repeat):
start = time.perf_counter_ns()
function()
duration = (time.perf_counter_ns() - start) / 1e9
timings.append(duration)
spent += duration
if spent >= max_duration:
break
finally:
if gc_enabled:
gc.enable()
return statistics.median(timings)
register_kernels()
threads = min(4, cpu_count)
options = onnxruntime.SessionOptions()
options.intra_op_num_threads = threads
options.inter_op_num_threads = 1
options.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
records = []
for seed, (trees, features, batch) in enumerate(
(trees, features, batch)
for trees in tree_counts
for features in feature_counts
for batch in batch_sizes
):
model, feeds = make_case(trees, features, batch, seed)
cpu_session = ReferenceEvaluator(
model, cpu_execution={"num_threads": threads, "affinity_policy": "none"}
)
ort_session = onnxruntime.InferenceSession(
model, sess_options=options, providers=["CPUExecutionProvider"]
)
def cpu_run(session=cpu_session, current_feeds=feeds):
return session.run(None, current_feeds)
def ort_run(session=ort_session, current_feeds=feeds):
return session.run(None, current_feeds)
np.testing.assert_allclose(cpu_run()[0], ort_run()[0], rtol=2e-5, atol=2e-5)
cpu_median = measure(cpu_run, args.repeat, args.warmup, args.max_repeat_time)
ort_median = measure(ort_run, args.repeat, args.warmup, args.max_repeat_time)
records.append(
{
"trees": trees,
"features": features,
"rows": batch,
"cpu_median_seconds": cpu_median,
"ort_median_seconds": ort_median,
"speedup": ort_median / cpu_median,
}
)
print(
f"trees={trees:6d} features={features:5d} batch={batch:4d} "
f"onnx-light-cpu={cpu_median * 1e6:10.2f}us "
f"onnxruntime={ort_median * 1e6:10.2f}us "
f"speedup={ort_median / cpu_median:.2f}x"
)
trees= 10 features= 4 batch= 1 onnx-light-cpu= 5.38us onnxruntime= 6.20us speedup=1.15x
trees= 10 features= 4 batch= 8 onnx-light-cpu= 5.80us onnxruntime= 7.23us speedup=1.25x
trees= 10 features= 4 batch= 32 onnx-light-cpu= 6.87us onnxruntime= 9.41us speedup=1.37x
trees= 10 features= 4 batch= 128 onnx-light-cpu= 1137.79us onnxruntime= 14.00us speedup=0.01x
trees= 10 features= 16 batch= 1 onnx-light-cpu= 5.43us onnxruntime= 5.91us speedup=1.09x
trees= 10 features= 16 batch= 8 onnx-light-cpu= 5.80us onnxruntime= 6.64us speedup=1.15x
trees= 10 features= 16 batch= 32 onnx-light-cpu= 6.93us onnxruntime= 8.99us speedup=1.30x
trees= 10 features= 16 batch= 128 onnx-light-cpu= 859.85us onnxruntime= 13.94us speedup=0.02x
trees= 10 features= 64 batch= 1 onnx-light-cpu= 5.50us onnxruntime= 5.97us speedup=1.09x
trees= 10 features= 64 batch= 8 onnx-light-cpu= 4.17us onnxruntime= 4.42us speedup=1.06x
trees= 10 features= 64 batch= 32 onnx-light-cpu= 6.92us onnxruntime= 9.01us speedup=1.30x
trees= 10 features= 64 batch= 128 onnx-light-cpu= 568.69us onnxruntime= 14.11us speedup=0.02x
trees= 10 features= 256 batch= 1 onnx-light-cpu= 5.45us onnxruntime= 5.88us speedup=1.08x
trees= 10 features= 256 batch= 8 onnx-light-cpu= 5.78us onnxruntime= 6.75us speedup=1.17x
trees= 10 features= 256 batch= 32 onnx-light-cpu= 4.91us onnxruntime= 5.86us speedup=1.19x
trees= 10 features= 256 batch= 128 onnx-light-cpu= 1140.88us onnxruntime= 14.09us speedup=0.01x
trees= 10 features= 1024 batch= 1 onnx-light-cpu= 5.41us onnxruntime= 5.92us speedup=1.09x
trees= 10 features= 1024 batch= 8 onnx-light-cpu= 4.19us onnxruntime= 6.76us speedup=1.61x
trees= 10 features= 1024 batch= 32 onnx-light-cpu= 7.50us onnxruntime= 9.02us speedup=1.20x
trees= 10 features= 1024 batch= 128 onnx-light-cpu= 870.19us onnxruntime= 14.50us speedup=0.02x
trees= 100 features= 4 batch= 1 onnx-light-cpu= 6.93us onnxruntime= 7.32us speedup=1.06x
trees= 100 features= 4 batch= 8 onnx-light-cpu= 9.48us onnxruntime= 13.54us speedup=1.43x
trees= 100 features= 4 batch= 32 onnx-light-cpu= 19.84us onnxruntime= 36.68us speedup=1.85x
trees= 100 features= 4 batch= 128 onnx-light-cpu= 39.61us onnxruntime= 49.45us speedup=1.25x
trees= 100 features= 16 batch= 1 onnx-light-cpu= 562.19us onnxruntime= 7.65us speedup=0.01x
trees= 100 features= 16 batch= 8 onnx-light-cpu= 9.54us onnxruntime= 13.67us speedup=1.43x
trees= 100 features= 16 batch= 32 onnx-light-cpu= 19.94us onnxruntime= 37.36us speedup=1.87x
trees= 100 features= 16 batch= 128 onnx-light-cpu= 35.61us onnxruntime= 50.10us speedup=1.41x
trees= 100 features= 64 batch= 1 onnx-light-cpu= 7.23us onnxruntime= 7.70us speedup=1.07x
trees= 100 features= 64 batch= 8 onnx-light-cpu= 9.49us onnxruntime= 13.49us speedup=1.42x
trees= 100 features= 64 batch= 32 onnx-light-cpu= 19.70us onnxruntime= 36.56us speedup=1.86x
trees= 100 features= 64 batch= 128 onnx-light-cpu= 616.60us onnxruntime= 50.28us speedup=0.08x
trees= 100 features= 256 batch= 1 onnx-light-cpu= 562.51us onnxruntime= 7.65us speedup=0.01x
trees= 100 features= 256 batch= 8 onnx-light-cpu= 9.59us onnxruntime= 13.62us speedup=1.42x
trees= 100 features= 256 batch= 32 onnx-light-cpu= 20.38us onnxruntime= 36.60us speedup=1.80x
trees= 100 features= 256 batch= 128 onnx-light-cpu= 35.91us onnxruntime= 51.22us speedup=1.43x
trees= 100 features= 1024 batch= 1 onnx-light-cpu= 6.92us onnxruntime= 7.63us speedup=1.10x
trees= 100 features= 1024 batch= 8 onnx-light-cpu= 9.64us onnxruntime= 13.75us speedup=1.43x
trees= 100 features= 1024 batch= 32 onnx-light-cpu= 27.57us onnxruntime= 37.07us speedup=1.34x
trees= 100 features= 1024 batch= 128 onnx-light-cpu= 49.86us onnxruntime= 53.35us speedup=1.07x
trees= 1000 features= 4 batch= 1 onnx-light-cpu= 856.81us onnxruntime= 13.38us speedup=0.02x
trees= 1000 features= 4 batch= 8 onnx-light-cpu= 45.98us onnxruntime= 87.70us speedup=1.91x
trees= 1000 features= 4 batch= 32 onnx-light-cpu= 90.38us onnxruntime= 315.07us speedup=3.49x
trees= 1000 features= 4 batch= 128 onnx-light-cpu= 1616.15us onnxruntime= 410.86us speedup=0.25x
trees= 1000 features= 16 batch= 1 onnx-light-cpu= 1136.87us onnxruntime= 10.50us speedup=0.01x
trees= 1000 features= 16 batch= 8 onnx-light-cpu= 46.73us onnxruntime= 86.38us speedup=1.85x
trees= 1000 features= 16 batch= 32 onnx-light-cpu= 148.54us onnxruntime= 313.04us speedup=2.11x
trees= 1000 features= 16 batch= 128 onnx-light-cpu= 1468.75us onnxruntime= 414.93us speedup=0.28x
trees= 1000 features= 64 batch= 1 onnx-light-cpu= 1136.66us onnxruntime= 13.86us speedup=0.01x
trees= 1000 features= 64 batch= 8 onnx-light-cpu= 46.59us onnxruntime= 86.58us speedup=1.86x
trees= 1000 features= 64 batch= 32 onnx-light-cpu= 147.83us onnxruntime= 314.54us speedup=2.13x
trees= 1000 features= 64 batch= 128 onnx-light-cpu= 1580.88us onnxruntime= 429.12us speedup=0.27x
trees= 1000 features= 256 batch= 1 onnx-light-cpu= 14.03us onnxruntime= 10.51us speedup=0.75x
trees= 1000 features= 256 batch= 8 onnx-light-cpu= 46.70us onnxruntime= 87.59us speedup=1.88x
trees= 1000 features= 256 batch= 32 onnx-light-cpu= 151.41us onnxruntime= 317.31us speedup=2.10x
trees= 1000 features= 256 batch= 128 onnx-light-cpu= 1290.28us onnxruntime= 905.59us speedup=0.70x
trees= 1000 features= 1024 batch= 1 onnx-light-cpu= 581.07us onnxruntime= 13.22us speedup=0.02x
trees= 1000 features= 1024 batch= 8 onnx-light-cpu= 48.01us onnxruntime= 87.92us speedup=1.83x
trees= 1000 features= 1024 batch= 32 onnx-light-cpu= 631.56us onnxruntime= 663.67us speedup=1.05x
trees= 1000 features= 1024 batch= 128 onnx-light-cpu= 1716.84us onnxruntime= 944.83us speedup=0.55x
Plot the timings and the speedups#
by_dimensions = {(row["trees"], row["rows"], row["features"]): row for row in records}
timings = {
trees: np.array(
[
[
by_dimensions[(trees, batch, features)]["cpu_median_seconds"] * 1e6
for features in feature_counts
]
for batch in batch_sizes
]
)
for trees in tree_counts
}
speedups = {
trees: np.array(
[
[by_dimensions[(trees, batch, features)]["speedup"] for features in feature_counts]
for batch in batch_sizes
]
)
for trees in tree_counts
}
all_timings = np.concatenate([values.ravel() for values in timings.values()])
all_speedups = np.concatenate([values.ravel() for values in speedups.values()])
timing_norm = LogNorm(vmin=float(all_timings.min()), vmax=float(all_timings.max()))
speedup_norm = TwoSlopeNorm(
vmin=min(0.5, float(all_speedups.min())),
vcenter=1.0,
vmax=max(1.5, float(all_speedups.max())),
)
speedup_cmap = LinearSegmentedColormap.from_list(
"slower_neutral_faster", ["#b2182b", "#f7f7f7", "#1b7837"]
)
def speedup_annotation_color(speedup):
red, green, blue, _ = speedup_cmap(speedup_norm(speedup))
red, green, blue = (
channel / 12.92 if channel <= 0.04045 else ((channel + 0.055) / 1.055) ** 2.4
for channel in (red, green, blue)
)
luminance = 0.2126 * red + 0.7152 * green + 0.0722 * blue
return "black" if luminance > 0.5 else "white"
fig, axes = plt.subplots(
len(tree_counts),
2,
figsize=(12, 4 * len(tree_counts)),
squeeze=False,
layout="constrained",
)
for row, trees in enumerate(tree_counts):
timing_image = axes[row, 0].imshow(
timings[trees], aspect="auto", cmap="viridis", norm=timing_norm
)
speedup_image = axes[row, 1].imshow(
speedups[trees], aspect="auto", cmap=speedup_cmap, norm=speedup_norm
)
for row_index in range(len(batch_sizes)):
for feature_index in range(len(feature_counts)):
axes[row, 0].text(
feature_index,
row_index,
f"{timings[trees][row_index, feature_index]:.1f}",
ha="center",
va="center",
color="white",
)
axes[row, 1].text(
feature_index,
row_index,
f"{speedups[trees][row_index, feature_index]:.2f}x",
ha="center",
va="center",
color=speedup_annotation_color(speedups[trees][row_index, feature_index]),
)
axes[row, 0].set_title(f"{trees:,} trees — CPU median time")
axes[row, 1].set_title(f"{trees:,} trees — speedup")
for axis in axes.flat:
axis.set_xticks(range(len(feature_counts)), feature_counts, rotation=30)
axis.set_yticks(range(len(batch_sizes)), batch_sizes)
axis.set_xlabel("number of features")
axis.set_ylabel("batch size")
fig.colorbar(timing_image, ax=axes[:, 0], label="onnx-light CPU median time (us)")
fig.colorbar(
speedup_image,
ax=axes[:, 1],
label="speedup versus ONNX Runtime (red: slower, green: faster)",
)
fig.suptitle("TreeEnsemble: depth 4, float32, one output")
fig.savefig("plot_tree_ensemble_benchmark.png")
plt.show()
Total running time of the script: (0 minutes 3.923 seconds)