.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/benchmarks/plot_tree_ensemble_benchmark.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_benchmarks_plot_tree_ensemble_benchmark.py: 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. .. GENERATED FROM PYTHON SOURCE LINES 17-69 .. code-block:: Python 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) .. GENERATED FROM PYTHON SOURCE LINES 70-75 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. .. GENERATED FROM PYTHON SOURCE LINES 75-144 .. code-block:: Python 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 .. GENERATED FROM PYTHON SOURCE LINES 145-150 Measure both runtimes --------------------- ``measure`` returns the median duration of a callable, stopping early once ``--max-repeat-time`` seconds have been spent. .. GENERATED FROM PYTHON SOURCE LINES 150-227 .. code-block:: Python 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" ) .. rst-class:: sphx-glr-script-out .. code-block:: none 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 .. GENERATED FROM PYTHON SOURCE LINES 228-230 Plot the timings and the speedups --------------------------------- .. GENERATED FROM PYTHON SOURCE LINES 230-325 .. code-block:: Python 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() .. image-sg:: /auto_examples/benchmarks/images/sphx_glr_plot_tree_ensemble_benchmark_001.png :alt: TreeEnsemble: depth 4, float32, one output, 10 trees — CPU median time, 10 trees — speedup, 100 trees — CPU median time, 100 trees — speedup, 1,000 trees — CPU median time, 1,000 trees — speedup :srcset: /auto_examples/benchmarks/images/sphx_glr_plot_tree_ensemble_benchmark_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 3.923 seconds) .. _sphx_glr_download_auto_examples_benchmarks_plot_tree_ensemble_benchmark.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_tree_ensemble_benchmark.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_tree_ensemble_benchmark.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_tree_ensemble_benchmark.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_