diff --git a/examples/benchmark.py b/examples/benchmark.py new file mode 100644 index 00000000..cc2257b4 --- /dev/null +++ b/examples/benchmark.py @@ -0,0 +1,63 @@ +# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com) +# +# MIT License +""" +File showcasing how to benchmark TensorRT engines. + +Demonstrates :func:`trtutils.benchmark_engine` for a single engine and +:func:`trtutils.benchmark_engines` for side-by-side comparison. Builds an +FP32 and FP16 variant of YOLOv8n and reports latency statistics for each. +""" + +from __future__ import annotations + +from pathlib import Path + +from trtutils import benchmark_engine, benchmark_engines, build_engine, set_log_level +from trtutils.download import download + + +def main() -> None: + onnx_path = Path("/tmp/yolov8n.onnx") # noqa: S108 + fp32_engine = Path("/tmp/yolov8n_fp32.engine") # noqa: S108 + fp16_engine = Path("/tmp/yolov8n_fp16.engine") # noqa: S108 + + if not onnx_path.exists(): + print("Downloading yolov8n ONNX model...") + download("yolov8n", onnx_path, imgsz=640, simplify=True) + + shapes = [("images", (1, 3, 640, 640))] + if not fp32_engine.exists(): + print("Building FP32 engine...") + build_engine(onnx_path, fp32_engine, shapes=shapes) + if not fp16_engine.exists(): + print("Building FP16 engine...") + build_engine(onnx_path, fp16_engine, fp16=True, shapes=shapes) + + print("\nSingle-engine benchmarks:") + for label, path in [("FP32", fp32_engine), ("FP16", fp16_engine)]: + result = benchmark_engine(path, iterations=200, warmup_iterations=20) + m = result.latency + print( + f" {label}: mean={m.mean:.3f} ms median={m.median:.3f} ms " + f"min={m.min:.3f} ms max={m.max:.3f} ms" + ) + + print("\nbenchmark_engines (serial):") + serial = benchmark_engines([fp32_engine, fp16_engine], iterations=200, warmup_iterations=20) + for label, result in zip(["FP32", "FP16"], serial): + print(f" {label}: mean={result.latency.mean:.3f} ms") + + print("\nbenchmark_engines (parallel, both engines run in lockstep):") + parallel = benchmark_engines( + [fp32_engine, fp16_engine], + iterations=200, + warmup_iterations=20, + parallel=True, + ) + print(f" combined: mean={parallel[0].latency.mean:.3f} ms") + + +if __name__ == "__main__": + set_log_level("ERROR") + main() diff --git a/examples/build.py b/examples/build.py new file mode 100644 index 00000000..0978ae56 --- /dev/null +++ b/examples/build.py @@ -0,0 +1,60 @@ +# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com) +# +# MIT License +""" +File showcasing how to build a TensorRT engine from an ONNX model. + +Demonstrates :func:`trtutils.builder.read_onnx` for peeking at the network +before building, then :func:`trtutils.build_engine` to produce a serialized +engine file. Bootstraps its own ONNX via :func:`trtutils.download.download`. +""" + +from __future__ import annotations + +import time +from pathlib import Path + +from trtutils import build_engine, set_log_level +from trtutils.builder import read_onnx +from trtutils.download import download + + +def main() -> None: + onnx_path = Path("/tmp/yolov8n.onnx") # noqa: S108 + engine_path = Path("/tmp/yolov8n.engine") # noqa: S108 + + if not onnx_path.exists(): + print("Downloading yolov8n ONNX model...") + download("yolov8n", onnx_path, imgsz=640, simplify=True) + + # peek at the parsed network before we build + network, _builder, _config, _parser = read_onnx(onnx_path) + print(f"ONNX network: {network.num_layers} layers") + for i in range(network.num_inputs): + t = network.get_input(i) + print(f" input {t.name}: shape={tuple(t.shape)}, dtype={t.dtype}") + for i in range(network.num_outputs): + t = network.get_output(i) + print(f" output {t.name}: shape={tuple(t.shape)}, dtype={t.dtype}") + # release the parser-side handles before building + del network, _builder, _config, _parser + + if engine_path.exists(): + engine_path.unlink() + + t0 = time.perf_counter() + build_engine( + onnx_path, + engine_path, + fp16=True, + shapes=[("images", (1, 3, 640, 640))], + ) + t1 = time.perf_counter() + + size_mb = engine_path.stat().st_size / (1024 * 1024) + print(f"Built FP16 engine in {t1 - t0:.2f} s -> {engine_path} ({size_mb:.2f} MB)") + + +if __name__ == "__main__": + set_log_level("ERROR") + main() diff --git a/examples/builder/dla.py b/examples/builder/dla.py new file mode 100644 index 00000000..2353afd6 --- /dev/null +++ b/examples/builder/dla.py @@ -0,0 +1,79 @@ +# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com) +# +# MIT License +""" +File showcasing how to build a TensorRT engine targeting the DLA. + +Demonstrates :func:`trtutils.builder.can_run_on_dla` to inspect which layers of +an ONNX model are DLA-compatible, and :func:`trtutils.builder.build_dla_engine` +to build a hybrid DLA/GPU engine. INT8 calibration is mandatory for DLA builds, +so we feed a :class:`trtutils.builder.SyntheticBatcher`. + +Exits cleanly when the system has no DLA hardware. +""" + +from __future__ import annotations + +import time +from pathlib import Path + +import numpy as np + +from trtutils import FLAGS, TRTEngine, set_log_level +from trtutils.builder import SyntheticBatcher, build_dla_engine, can_run_on_dla +from trtutils.download import download + + +def main() -> None: + if not FLAGS.HAS_DLA: + print(f"Skipping: no DLA cores available (NUM_DLA_CORES={FLAGS.NUM_DLA_CORES}).") + return + + onnx_path = Path("/tmp/yolov8n.onnx") # noqa: S108 + engine_path = Path("/tmp/yolov8n_dla.engine") # noqa: S108 + + if not onnx_path.exists(): + print("Downloading yolov8n ONNX model...") + download("yolov8n", onnx_path, imgsz=640, simplify=True) + + full_dla, chunks = can_run_on_dla(onnx_path) + print(f"Fully DLA-compatible: {full_dla}") + print(f"Found {len(chunks)} layer chunks:") + for i, (layers, start, end, on_dla) in enumerate(chunks): + target = "DLA" if on_dla else "GPU" + print(f" chunk {i}: layers [{start}-{end}] ({len(layers)} layers) -> {target}") + + # DLA builds need INT8 calibration data; use synthetic data for the demo + batcher = SyntheticBatcher( + shape=(640, 640, 3), + dtype=np.float32, + batch_size=1, + num_batches=8, + ) + + if engine_path.exists(): + engine_path.unlink() + + t0 = time.perf_counter() + build_dla_engine( + onnx_path, + engine_path, + data_batcher=batcher, + dla_core=0, + shapes=[("images", (1, 3, 640, 640))], + ) + t1 = time.perf_counter() + + size_mb = engine_path.stat().st_size / (1024 * 1024) + print(f"Built DLA engine in {t1 - t0:.2f} s -> {engine_path} ({size_mb:.2f} MB)") + + # confirm the engine is loadable; cuda_graph=False since DLA + graphs don't mix + engine = TRTEngine(engine_path, dla_core=0, warmup=True, cuda_graph=False) + engine.mock_execute() + print(f"Loaded {engine.name}, mock_execute OK") + del engine + + +if __name__ == "__main__": + set_log_level("ERROR") + main() diff --git a/examples/config.py b/examples/config.py new file mode 100644 index 00000000..5e895a8e --- /dev/null +++ b/examples/config.py @@ -0,0 +1,82 @@ +# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com) +# +# MIT License +""" +File showcasing trtutils configuration, logging, and profiling toggles. + +Demonstrates :data:`trtutils.FLAGS`, :data:`trtutils.CONFIG`, +:func:`trtutils.set_log_level`, the :data:`trtutils.NVTX` context manager, +and :func:`trtutils.register_jit` / :data:`trtutils.JIT` for Numba JIT +compilation. No model needed. +""" + +from __future__ import annotations + +import time + +import numpy as np + +from trtutils import ( + CONFIG, + FLAGS, + JIT, + NVTX, + register_jit, + set_log_level, +) + + +@register_jit(fastmath=True) +def sum_squares(arr: np.ndarray) -> float: + """Trivial numeric kernel — Numba JITs this when JIT is enabled.""" + total = 0.0 + for value in arr: + total += value * value + return float(total) + + +def main() -> None: + print("FLAGS:") + for attr in sorted( + a for a in dir(FLAGS) if not a.startswith("_") and not callable(getattr(FLAGS, a)) + ): + print(f" {attr}: {getattr(FLAGS, attr)}") + + print("\nCONFIG: loading TensorRT plugins (idempotent)...") + CONFIG.load_plugins() + print("CONFIG: plugins loaded.") + + print("\nLog level demo — toggle between INFO and ERROR:") + set_log_level("INFO") + print(" log level set to INFO (TensorRT messages would print here)") + set_log_level("ERROR") + print(" log level set back to ERROR") + + print("\nNVTX context manager — ranges are visible to Nsight Systems:") + with NVTX("example::demo"): + time.sleep(0.001) + print(f" NVTX_ENABLED after context exit: {FLAGS.NVTX_ENABLED}") + + print("\nJIT context manager — toggle Numba compilation around a hot loop:") + data = np.random.default_rng(0).standard_normal(100_000).astype(np.float32) + + # Warm up either path so we measure steady-state cost + sum_squares(data) + t0 = time.perf_counter() + sum_squares(data) + no_jit_ms = (time.perf_counter() - t0) * 1000.0 + print(f" baseline (JIT={FLAGS.JIT}): {no_jit_ms:.3f} ms") + + with JIT: + sum_squares(data) # one warmup so the JIT compile cost is excluded + t0 = time.perf_counter() + sum_squares(data) + with_jit_ms = (time.perf_counter() - t0) * 1000.0 + print( + f" inside JIT block (JIT={FLAGS.JIT}, Numba={FLAGS.FOUND_NUMBA}): {with_jit_ms:.3f} ms" + ) + + +if __name__ == "__main__": + set_log_level("ERROR") + main() diff --git a/examples/core.py b/examples/core.py new file mode 100644 index 00000000..0eb186af --- /dev/null +++ b/examples/core.py @@ -0,0 +1,97 @@ +# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com) +# +# MIT License +""" +File showcasing a high-level tour of the ``trtutils.core`` CUDA backend. + +Demonstrates :class:`trtutils.core.Device`, device introspection +(:func:`trtutils.core.get_device`, :func:`trtutils.core.get_device_name`, +:func:`trtutils.core.get_compute_capability`), stream lifecycle +(:func:`trtutils.core.create_stream` / :func:`trtutils.core.destroy_stream`, +:func:`trtutils.core.stream_synchronize`), explicit device memory +(:func:`trtutils.core.cuda_malloc`, :func:`trtutils.core.cuda_free`, +:func:`trtutils.core.memcpy_host_to_device` / +:func:`trtutils.core.memcpy_device_to_host`), and :class:`trtutils.core.CUDAGraph` +capture/replay wrapped around a synthetic stream sleep. + +The goal is to show that the ``core`` module exists and to give a guided +overview of its building blocks; it is not a deep dive into any one piece. +""" + +from __future__ import annotations + +import numpy as np + +from trtutils import set_log_level +from trtutils.core import ( + CUDAGraph, + Device, + create_stream, + cuda_free, + cuda_malloc, + destroy_stream, + get_compute_capability, + get_device, + get_device_count, + get_device_name, + get_num_dla_cores, + memcpy_device_to_host, + memcpy_host_to_device, + memcpy_host_to_device_async, + stream_synchronize, +) + + +def main() -> None: + print("Device info:") + print(f" current device index: {get_device()}") + print(f" device count: {get_device_count()}") + print(f" device name: {get_device_name()}") + print(f" compute capability: {get_compute_capability()}") + print(f" DLA cores: {get_num_dla_cores()}") + + # Device(idx) saves/restores the current device on exit; Device(None) is a no-op. + with Device(get_device()): + print(f" inside Device guard: {get_device()}") + + # synchronous memcpy roundtrip + host = np.arange(8, dtype=np.float32) + nbytes = host.nbytes + device_ptr = cuda_malloc(nbytes) + memcpy_host_to_device(device_ptr, host) + roundtrip = np.zeros_like(host) + memcpy_device_to_host(roundtrip, device_ptr) + cuda_free(device_ptr) + print(f"\nSync memcpy roundtrip: {host.tolist()} -> {roundtrip.tolist()}") + + # async memcpy through a stream + stream = create_stream() + device_ptr = cuda_malloc(nbytes) + memcpy_host_to_device_async(device_ptr, host, stream) + stream_synchronize(stream) + memcpy_device_to_host(roundtrip, device_ptr) + print(f"Async memcpy result: {roundtrip.tolist()}") + cuda_free(device_ptr) + + # CUDA graph capture/replay around an async memcpy + src = np.arange(16, dtype=np.float32) + dst = np.zeros_like(src) + device_ptr = cuda_malloc(src.nbytes) + graph = CUDAGraph(stream) + with graph: + memcpy_host_to_device_async(device_ptr, src, stream) + if graph.is_captured: + graph.launch() + stream_synchronize(stream) + memcpy_device_to_host(dst, device_ptr) + print(f"\nCUDA graph captured and replayed: dst[:4]={dst[:4].tolist()}") + else: + print("\nCUDA graph capture failed on this stream (skipping launch).") + graph.invalidate() + cuda_free(device_ptr) + destroy_stream(stream) + + +if __name__ == "__main__": + set_log_level("ERROR") + main() diff --git a/examples/classifier.py b/examples/image/classifier.py similarity index 89% rename from examples/classifier.py rename to examples/image/classifier.py index 5dfadf13..92118bae 100644 --- a/examples/classifier.py +++ b/examples/image/classifier.py @@ -15,12 +15,12 @@ def main() -> None: - engine_dir = Path(__file__).parent.parent / "data" / "engines" + engine_dir = Path(__file__).resolve().parent.parent.parent / "data" / "engines" engines = [ engine_dir / "resnet18.engine", ] - img_path = str(Path(__file__).parent.parent / "data" / "horse.jpg") + img_path = str(Path(__file__).resolve().parent.parent.parent / "data" / "horse.jpg") img = cv2.imread(img_path) if img is None: err_msg = f"Failed to load image from {img_path}" diff --git a/examples/depth_estimator.py b/examples/image/depth_estimator.py similarity index 96% rename from examples/depth_estimator.py rename to examples/image/depth_estimator.py index 32d2fbb1..80e072fe 100644 --- a/examples/depth_estimator.py +++ b/examples/image/depth_estimator.py @@ -15,7 +15,7 @@ from trtutils.download import download from trtutils.image import DepthEstimator -DATA_DIR = Path(__file__).resolve().parent.parent / "data" +DATA_DIR = Path(__file__).resolve().parent.parent.parent / "data" def main() -> None: diff --git a/examples/detector.py b/examples/image/detector.py similarity index 90% rename from examples/detector.py rename to examples/image/detector.py index a4f7d9ed..105069ce 100644 --- a/examples/detector.py +++ b/examples/image/detector.py @@ -15,7 +15,7 @@ def main() -> None: - engine_dir = Path(__file__).parent.parent / "data" / "engines" + engine_dir = Path(__file__).resolve().parent.parent.parent / "data" / "engines" engines = [ engine_dir / "trt_yolov7t.engine", engine_dir / "trt_yolov8n.engine", @@ -27,7 +27,7 @@ def main() -> None: engine_dir / "trt_yolov10n_dla.engine", ] - img_path = str(Path(__file__).parent.parent / "data" / "horse.jpg") + img_path = str(Path(__file__).resolve().parent.parent.parent / "data" / "horse.jpg") img = cv2.imread(img_path) if img is None: err_msg = f"Failed to load image from {img_path}" diff --git a/examples/sahi.py b/examples/image/sahi.py similarity index 92% rename from examples/sahi.py rename to examples/image/sahi.py index 06df494d..9b1a2971 100644 --- a/examples/sahi.py +++ b/examples/image/sahi.py @@ -30,10 +30,10 @@ def main() -> None: ) args = parser.parse_args() - engine_dir = Path(__file__).parent.parent / "data" / "engines" + engine_dir = Path(__file__).resolve().parent.parent.parent / "data" / "engines" engine_path = engine_dir / "trt_yolov10n.engine" - img_path = str(Path(__file__).parent.parent / "data" / "cars.jpeg") + img_path = str(Path(__file__).resolve().parent.parent.parent / "data" / "cars.jpeg") img = cv2.imread(img_path) if img is None: err_msg = f"Failed to load image from {img_path}" diff --git a/examples/inspect.py b/examples/inspect.py new file mode 100644 index 00000000..919ff14d --- /dev/null +++ b/examples/inspect.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com) +# +# MIT License +""" +File showcasing how to inspect TensorRT engines and ONNX models. + +Demonstrates :func:`trtutils.inspect_engine`, :func:`trtutils.inspect.get_engine_names`, +and :func:`trtutils.inspect.inspect_onnx_layers` — first peeking at a built +engine, then walking the source ONNX layer-by-layer to show the ONNX-to-TRT +fusion mapping. +""" + +from __future__ import annotations + +from pathlib import Path + +from trtutils import build_engine, inspect_engine, set_log_level +from trtutils.download import download +from trtutils.inspect import get_engine_names, inspect_onnx_layers + + +def main() -> None: + onnx_path = Path("/tmp/yolov8n.onnx") # noqa: S108 + engine_path = Path("/tmp/yolov8n.engine") # noqa: S108 + + if not onnx_path.exists(): + print("Downloading yolov8n ONNX model...") + download("yolov8n", onnx_path, imgsz=640, simplify=True) + if not engine_path.exists(): + print("Building yolov8n engine...") + build_engine(onnx_path, engine_path, fp16=True, shapes=[("images", (1, 3, 640, 640))]) + + # high-level engine summary + mem_size, batch_size, inputs, outputs = inspect_engine(engine_path) + print(f"Engine: {engine_path.name}") + print(f" device memory: {mem_size / (1024 * 1024):.2f} MB") + print(f" max batch size: {batch_size}") + print(f" inputs ({len(inputs)}):") + for name, shape, dtype, fmt in inputs: + print(f" {name}: shape={tuple(shape)} dtype={dtype} format={fmt}") + print(f" outputs ({len(outputs)}):") + for name, shape, dtype, fmt in outputs: + print(f" {name}: shape={tuple(shape)} dtype={dtype} format={fmt}") + + # input/output names in enumeration order + in_names, out_names = get_engine_names(engine_path) + print(f"Names: inputs={in_names}, outputs={out_names}") + + # walk the ONNX layers and show DLA compatibility per layer + onnx_layers = inspect_onnx_layers(onnx_path) + print(f"\nONNX layers: {len(onnx_layers)} total") + dla_count = sum(1 for layer in onnx_layers if layer.dla_compatible) + print(f" DLA-compatible: {dla_count}/{len(onnx_layers)}") + print("First 5 layers:") + for layer in onnx_layers[:5]: + tag = "DLA" if layer.dla_compatible else "GPU" + print( + f" [{layer.index:3d}] {layer.layer_type:<16} {layer.name} " + f"out={layer.output_tensor_size}B {tag}" + ) + + +if __name__ == "__main__": + set_log_level("ERROR") + main() diff --git a/examples/jetson.py b/examples/jetson.py new file mode 100644 index 00000000..2053ce58 --- /dev/null +++ b/examples/jetson.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com) +# +# MIT License +""" +File showcasing Jetson-specific benchmarking and profiling. + +Demonstrates :func:`trtutils.jetson.benchmark_engine` (latency + power + energy) +and :func:`trtutils.jetson.profile_engine` (per-layer energy contributions). +These mirror the standard :mod:`trtutils` benchmark / profile functions but +also sample VDD_TOTAL via tegrastats. + +Exits cleanly when not running on a Jetson device. ``cuda_graph=False`` is +required for engines with DLA layers because DLA does not support CUDA graphs. +""" + +from __future__ import annotations + +from pathlib import Path + +import tensorrt as trt + +from trtutils import FLAGS, build_engine, set_log_level +from trtutils.download import download + +# trtutils.jetson is only importable on Jetson; do the import lazily. +if FLAGS.IS_JETSON: + from trtutils import jetson + + +def main() -> None: + if not FLAGS.IS_JETSON: + print("Skipping: not running on a Jetson device.") + return + + onnx_path = Path("/tmp/yolov8n.onnx") # noqa: S108 + engine_path = Path("/tmp/yolov8n_detailed.engine") # noqa: S108 + + if not onnx_path.exists(): + print("Downloading yolov8n ONNX model...") + download("yolov8n", onnx_path, imgsz=640, simplify=True) + if not engine_path.exists(): + print("Building yolov8n engine with DETAILED profiling verbosity...") + build_engine( + onnx_path, + engine_path, + fp16=True, + shapes=[("images", (1, 3, 640, 640))], + profiling_verbosity=trt.ProfilingVerbosity.DETAILED, + ) + + print("Benchmarking on Jetson (measures latency + power + energy)...") + result = jetson.benchmark_engine( + engine_path, + iterations=500, + warmup_iterations=20, + cuda_graph=False, + ) + print(f" latency: mean={result.latency.mean:.3f} ms") + print(f" power: mean={result.power_draw.mean:.1f} mW") + print(f" energy: mean={result.energy.mean:.3f} mJ/iter") + + print("\nProfiling per-layer energy (this takes a while)...") + prof = jetson.profile_engine( + engine_path, + iterations=2000, + warmup_iterations=20, + cuda_graph=False, + ) + print(f" layers profiled: {len(prof.layers)}") + print(f" total power: {prof.power_draw.mean:.1f} mW") + print(f" total energy: {prof.energy.mean:.3f} mJ/iter") + + hottest = sorted(prof.layers, key=lambda layer: layer.energy, reverse=True)[:5] + print("\nTop 5 layers by energy:") + for layer in hottest: + print( + f" {layer.energy:7.3f} mJ ({layer.power:6.1f} mW for {layer.mean:5.3f} ms) " + f"{layer.name}" + ) + + +if __name__ == "__main__": + set_log_level("ERROR") + main() diff --git a/examples/kernel.py b/examples/kernel.py new file mode 100644 index 00000000..24c87439 --- /dev/null +++ b/examples/kernel.py @@ -0,0 +1,81 @@ +# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com) +# +# MIT License +""" +File showcasing the :class:`trtutils.core.Kernel` abstraction. + +Compiles a small CUDA kernel via NVRTC, allocates a device buffer, runs the +kernel on a stream, and copies the result back. ``Kernel`` reads CUDA source +from a ``.cu`` file and exposes ``create_args`` / ``__call__`` for ergonomic +launches. + +Self-contained: writes the kernel source to a temp file, no engine required. +""" + +from __future__ import annotations + +import tempfile +from pathlib import Path + +import numpy as np + +from trtutils import set_log_level +from trtutils.core import ( + Kernel, + create_stream, + cuda_free, + cuda_malloc, + destroy_stream, + memcpy_device_to_host, + stream_synchronize, +) + +ADD_ONE_KERNEL = """\ +extern "C" __global__ void add_one(float *out, int n) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + out[idx] = (float)idx + 1.0f; + } +} +""" + + +def main() -> None: + # Write the CUDA source to a temp file — Kernel takes a file path. + with tempfile.NamedTemporaryFile("w", suffix=".cu", delete=False) as f: + f.write(ADD_ONE_KERNEL) + cu_path = Path(f.name) + + kernel_name = "add_one" + kernel = Kernel(cu_path, name=kernel_name) + print(f"Compiled kernel '{kernel_name}' from {cu_path}") + + n = 32 + out_bytes = n * np.dtype(np.float32).itemsize + d_out = cuda_malloc(out_bytes) + stream = create_stream() + + # create_args boxes Python ints/floats and device pointers into a uint64 + # pointer array the CUDA launch ABI expects. + args = kernel.create_args(d_out, n) + + # 1 block of n threads + kernel((1, 1, 1), (n, 1, 1), stream, args) + stream_synchronize(stream) + + result = np.zeros(n, dtype=np.float32) + memcpy_device_to_host(result, d_out) + print(f"Kernel output (first 8 of {n}): {result[:8].tolist()}") + print( + f"Matches expected (idx + 1): {np.array_equal(result, np.arange(1, n + 1, dtype=np.float32))}" + ) + + kernel.free() + cuda_free(d_out) + destroy_stream(stream) + cu_path.unlink(missing_ok=True) + + +if __name__ == "__main__": + set_log_level("ERROR") + main() diff --git a/examples/profile.py b/examples/profile.py new file mode 100644 index 00000000..025fb42d --- /dev/null +++ b/examples/profile.py @@ -0,0 +1,68 @@ +# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com) +# +# MIT License +""" +File showcasing how to profile a TensorRT engine layer-by-layer. + +Demonstrates :func:`trtutils.profile_engine` for per-layer timing statistics +and :func:`trtutils.profiling.identify_quantize_speedups_by_layer` for +detecting layers that benefit from INT8 quantization. + +For useful per-layer names, the engine must be built with +``profiling_verbosity=trt.ProfilingVerbosity.DETAILED``. +""" + +from __future__ import annotations + +from pathlib import Path + +import tensorrt as trt + +from trtutils import build_engine, profile_engine, set_log_level +from trtutils.download import download +from trtutils.profiling import identify_quantize_speedups_by_layer + + +def main() -> None: + onnx_path = Path("/tmp/yolov8n.onnx") # noqa: S108 + engine_path = Path("/tmp/yolov8n_detailed.engine") # noqa: S108 + + if not onnx_path.exists(): + print("Downloading yolov8n ONNX model...") + download("yolov8n", onnx_path, imgsz=640, simplify=True) + + if not engine_path.exists(): + print("Building yolov8n engine with DETAILED profiling verbosity...") + build_engine( + onnx_path, + engine_path, + fp16=True, + shapes=[("images", (1, 3, 640, 640))], + profiling_verbosity=trt.ProfilingVerbosity.DETAILED, + ) + + result = profile_engine(engine_path, iterations=100, warmup_iterations=10) + print(f"Profiled {result.iterations} iterations across {len(result.layers)} layers") + print(f"Total per-iteration time: mean={result.total_time.mean:.3f} ms") + + top_n = 10 + slowest = sorted(result.layers, key=lambda layer: layer.mean, reverse=True)[:top_n] + print(f"\nTop {top_n} slowest layers:") + for layer in slowest: + print(f" {layer.mean:7.3f} ms {layer.name}") + + print("\nScanning for INT8 quantization speedups (this builds both FP16 + INT8 engines)...") + _fp16, _int8, speedups = identify_quantize_speedups_by_layer( + onnx_path, + iterations=50, + warmup_iterations=5, + ) + quantize_wins = sorted(speedups, key=lambda pair: pair[1], reverse=True)[:5] + print("Top 5 INT8 wins (positive % means INT8 faster):") + for name, speedup in quantize_wins: + print(f" {speedup:+6.2f}% {name}") + + +if __name__ == "__main__": + set_log_level("ERROR") + main() diff --git a/examples/trtexec.py b/examples/trtexec.py new file mode 100644 index 00000000..635e4a6b --- /dev/null +++ b/examples/trtexec.py @@ -0,0 +1,75 @@ +# Copyright (c) 2026 Justin Davis (davisjustin302@gmail.com) +# +# MIT License +""" +File showcasing the trtexec wrapper utilities. + +Demonstrates :func:`trtutils.find_trtexec` for locating the binary, +:func:`trtutils.run_trtexec` for running raw commands, and +:func:`trtutils.trtexec.build_engine` for building an engine through the +external tool. The resulting engine is then loaded with +:class:`trtutils.TRTEngine` to confirm round-trip compatibility. + +Exits cleanly when ``trtexec`` is not installed on the system. +""" + +from __future__ import annotations + +from pathlib import Path + +from trtutils import TRTEngine, find_trtexec, run_trtexec, set_log_level +from trtutils import trtexec as trtexec_mod +from trtutils.download import download + + +def main() -> None: + try: + trtexec_path = find_trtexec() + except FileNotFoundError as exc: + print(f"Skipping: {exc}") + return + + print(f"Found trtexec at: {trtexec_path}") + + # run a trivial command — the version banner — and print the first few lines + success, stdout, _stderr = run_trtexec("--help") + if not success: + print("trtexec --help did not exit cleanly; skipping rest.") + return + head = "\n".join(stdout.splitlines()[:3]) + print(f"Banner:\n{head}") + + onnx_path = Path("/tmp/yolov8n.onnx") # noqa: S108 + engine_path = Path("/tmp/yolov8n_trtexec.engine") # noqa: S108 + + if not onnx_path.exists(): + print("Downloading yolov8n ONNX model...") + download("yolov8n", onnx_path, imgsz=640, simplify=True) + + if engine_path.exists(): + engine_path.unlink() + + print("\nBuilding engine via trtexec.build_engine(fp16=True)...") + ok = trtexec_mod.build_engine( + onnx_path, + engine_path, + fp16=True, + shapes=[("images", (1, 3, 640, 640))], + ) + if not ok: + print("trtexec.build_engine reported failure.") + return + + size_mb = engine_path.stat().st_size / (1024 * 1024) + print(f"trtexec built engine: {engine_path} ({size_mb:.2f} MB)") + + # round-trip: load with TRTEngine + engine = TRTEngine(engine_path, warmup=True) + engine.mock_execute() + print(f"Loaded {engine.name} with TRTEngine, mock_execute OK") + del engine + + +if __name__ == "__main__": + set_log_level("ERROR") + main()