RF-DETR on NVIDIA GPUs — Export, Inference & Latency¶

Compares every export path that targets a CUDA GPU: eager PyTorch (fp32 and fp16+JIT) as the reference anchor, ONNX Runtime with the CUDA execution provider, and a TensorRT engine — each exported, run once for a correctness check, then benchmarked.

Format	Export	Inference
PyTorch	(no export — eager model)	predict() / inference(dtype=torch.float16)
ONNX (CUDA EP)	model.export(dynamic_batch=True) → .onnx	onnxruntime.InferenceSession
TensorRT	model.export(format="trt") → .trt	inference-models (selects and loads its own prebuilt TensorRT engine — see note below), and rfdetr's own TRTInference loading the .trt file directly

Every format is benchmarked at batch 1, 4, and 8 — see src/rfdetr/export/registry.py for which formats can hold one artifact across batch sizes and which cannot: ONNX and TensorRT both bake a dynamic batch axis (supports_dynamic_batch=True), so one exported artifact serves every row below for each of them; PyTorch needs no export step and trivially takes any batch as a (N, C, H, W) tensor. inference-models' managed TensorRT engine is a separate artifact (its engine build is not part of this repository) — its own batch handling is exercised through its documented multi-image __call__ signature (Union[np.ndarray, List[np.ndarray], torch.Tensor, List[torch.Tensor]], confirmed by reading inference_models/models/base/object_detection.py on GitHub), not verified against the package's source the way the two rfdetr-exported formats are.

Not covered here: CPU-only deployment (see the CPU cookbook), mobile/edge formats (see the mobile cookbook), or Apple Silicon (see the Apple cookbook). This notebook requires a CUDA GPU end to end — it will not run on a CPU-only or Apple Silicon runtime.

Two TensorRT artifacts, on purpose. The .trt file from the export step is a standalone artifact for raw TensorRT deployment, locked to this GPU + TensorRT version. inference-models selects and loads its own prebuilt TensorRT engine, so it does not load that file — it is the easier, portable inference path, benchmarked first below. The .trt file itself is then loaded and timed too, via rfdetr.export._tensorrt.inference.TRTInference — a reference runtime the package ships for parity testing, not a recommended production path (see its section below) — so both TensorRT artifacts get a number, not just the one inference-models loads. It is exported once with dynamic_batch=True, max_batch_size=8 so the same engine serves every batch size below, the same way the managed engine and ONNX do.

Which latency number is this? Every number in this notebook is a Tier C measurement: forward-only and end-to-end timings for each export format, on one L4, against a PyTorch anchor measured on that same L4. Its value is ratios within this machine — never set these milliseconds against the architecture headline on the Benchmarks page (T4, TensorRT FP16, model only, batch 1) or the number on Roboflow's RF-DETR model page (L4, Roboflow Inference, scope not stated, batch 1) without checking hardware, precision, batch size, and scope match. See Section 7 for where this notebook's own TensorRT number disagrees with Roboflow's model page.

1. Install¶

RFDETR_SPEC holds the install source and defaults to the develop branch, to pick up the newest export fixes; to install the 1.11.2 release from PyPI instead, swap in the commented >=1.11.2,<2 line. Colab ships mutually inconsistent preinstalled packages that otherwise crash import rfdetr, so two are aligned: torchaudio is uninstalled (RF-DETR never uses it, but transformers imports it when present and a torch/torchaudio CUDA-version mismatch then errors), and pillow is force-reinstalled to a clean version (a half-upgraded PIL breaks torchvision's import with cannot import name '_Ink'). The default PyPI onnxruntime-gpu wheel targets CUDA 11.8 and silently falls back to CPU on modern GPUs, so it is reinstalled from the CUDA-12 package index — --force-reinstall is required here: rfdetr[tensorrt] already pulls a plain onnxruntime-gpu wheel, and without it pip sees the requirement as already satisfied and skips the CUDA-12 build entirely.

Colab: select a GPU runtime (Runtime → Change runtime type → GPU), and after this cell Runtime → Restart session before running the next cell — Colab keeps old package versions loaded until a restart.

In [ ]:

RFDETR_SPEC = "rfdetr[onnx,tensorrt] @ https://github.com/roboflow/rf-detr/archive/refs/heads/develop.zip"
# RFDETR_SPEC = "rfdetr[onnx,tensorrt]>=1.11.2,<2"
!pip uninstall -q -y onnxruntime onnxruntime-gpu
!pip install -q "{RFDETR_SPEC}" "inference-models[trt10]" supervision pandas
!pip install -q --force-reinstall --no-deps onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/
!pip install -q --force-reinstall --no-deps "pillow==11.3.0"
!pip uninstall -q -y torchaudio

RFDETR_SPEC = "rfdetr[onnx,tensorrt] @ https://github.com/roboflow/rf-detr/archive/refs/heads/develop.zip" # RFDETR_SPEC = "rfdetr[onnx,tensorrt]>=1.11.2,<2" !pip uninstall -q -y onnxruntime onnxruntime-gpu !pip install -q "{RFDETR_SPEC}" "inference-models[trt10]" supervision pandas !pip install -q --force-reinstall --no-deps onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/ !pip install -q --force-reinstall --no-deps "pillow==11.3.0" !pip uninstall -q -y torchaudio

2. Setup and GPU check¶

Every format in this notebook requires CUDA — fail fast with a clear message if no GPU is visible. WARMUP_RUNS discards the first N inferences (GPU kernels are JIT-compiled on first use); MEASURE_RUNS then collects the steady-state timing distribution. BATCH_SIZES is the batch matrix every format section below measures — (1, 4, 8), not a single number, so each format contributes three rows to the final results table.

Every format section below shares a handful of small helpers that used to be copy-pasted into each per-hardware cookbook separately; they now live once, privately, in rfdetr.export._benchmark (and _decode_batch in particular in rfdetr.export._runtime.decode's companion): _artifact_size_mb reports an export artifact's size on disk, visualize_detections annotates and displays a supervision.Detections on the sample image, _tile_batch stacks the one sample image's preprocessed tensor into a (batch, C, H, W) array (all batch copies are the same image — this notebook measures throughput at a larger batch dimension, not batch diversity), _decode_batch runs decode_detections over every image in a batched raw-output pair, and measure_memory measures device-memory growth via torch.cuda.mem_get_info() — device-wide, so it also captures ONNX Runtime's CUDA execution provider and TensorRT's own cudaMalloc calls, not just PyTorch's own caching allocator.

In [ ]:

import importlib.metadata
from pathlib import Path

import numpy as np
import supervision as sv
import torch
from PIL import Image

from rfdetr.export._benchmark import (
    BenchmarkResult,
    _artifact_size_mb,
    _decode_batch,
    _enable_notebook_inline_matplotlib,
    _result_row,
    _tile_batch,
    measure_latency,
    measure_memory,
    visualize_detections,
)

if not torch.cuda.is_available():
    raise RuntimeError("This notebook requires a CUDA GPU; none is available.")
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"rfdetr: {importlib.metadata.version('rfdetr')}")

EXPORT_DIR = Path("export_cuda")
EXPORT_DIR.mkdir(exist_ok=True)
CONFIDENCE_THRESHOLD = 0.5
WARMUP_RUNS = 20
MEASURE_RUNS = 100
BATCH_SIZES = (1, 4, 8)

_enable_notebook_inline_matplotlib()

import importlib.metadata from pathlib import Path import numpy as np import supervision as sv import torch from PIL import Image from rfdetr.export._benchmark import ( BenchmarkResult, _artifact_size_mb, _decode_batch, _enable_notebook_inline_matplotlib, _result_row, _tile_batch, measure_latency, measure_memory, visualize_detections, ) if not torch.cuda.is_available(): raise RuntimeError("This notebook requires a CUDA GPU; none is available.") print(f"GPU: {torch.cuda.get_device_name(0)}") print(f"rfdetr: {importlib.metadata.version('rfdetr')}") EXPORT_DIR = Path("export_cuda") EXPORT_DIR.mkdir(exist_ok=True) CONFIDENCE_THRESHOLD = 0.5 WARMUP_RUNS = 20 MEASURE_RUNS = 100 BATCH_SIZES = (1, 4, 8) _enable_notebook_inline_matplotlib()

3. Sample image¶

A single street scene with several COCO classes (dog, bicycle, car) is enough to verify detections. The image is downloaded once and reused for every format below.

In [ ]:

import urllib.request

IMAGE_URL = "https://media.roboflow.com/notebooks/examples/dog.jpeg"
IMAGE_PATH = EXPORT_DIR / "sample.jpg"
if not IMAGE_PATH.exists():
    urllib.request.urlretrieve(IMAGE_URL, IMAGE_PATH)

image = Image.open(IMAGE_PATH).convert("RGB")
print(f"Sample image: {image.size[0]}×{image.size[1]}")

import urllib.request IMAGE_URL = "https://media.roboflow.com/notebooks/examples/dog.jpeg" IMAGE_PATH = EXPORT_DIR / "sample.jpg" if not IMAGE_PATH.exists(): urllib.request.urlretrieve(IMAGE_URL, IMAGE_PATH) image = Image.open(IMAGE_PATH).convert("RGB") print(f"Sample image: {image.size[0]}×{image.size[1]}")

4. PyTorch baseline — predict() / inference()¶

No export needed — this is the reference every other format in this notebook is compared against. predict() is the unoptimized fp32 baseline and accepts a list of images, batching them internally in one forward call, so the batch matrix below feeds it [image] * batch. inference(dtype=torch.float16, batch_size=batch) fuses layers with torch.jit.script, halves arithmetic precision, and compiles the scripted module for exactly batch_size images — calling predict() with a different batch afterward raises ValueError: Batch size mismatch (confirmed by an earlier run of this notebook), so the batch matrix needs one recompile per batch size, same as the export formats below; remove_optimized_model() after the loop reverts the in-place optimization so model stays reusable for the export calls that follow.

Both calls include preprocessing and postprocessing — there is no separate "forward-only" path through eager predict(), so only an end-to-end number is reported for the PyTorch baseline.

In [ ]:

from functools import partial

from rfdetr import RFDETRSmall

with measure_memory(device="cuda") as mem:
    model = RFDETRSmall()
    baseline_detections = model.predict(image, threshold=CONFIDENCE_THRESHOLD)
pytorch_fp32_memory_mb = mem.delta_mb
print(f"PyTorch fp32 baseline: {len(baseline_detections)} detections above {CONFIDENCE_THRESHOLD}")
visualize_detections(baseline_detections, image)

pytorch_fp32_results: dict[int, BenchmarkResult] = {}
for batch in BATCH_SIZES:
    images_batch = [image] * batch
    pytorch_fp32_results[batch] = measure_latency(
        partial(model.predict, images_batch),
        label=f"PyTorch predict() fp32 batch={batch}",
        device="cuda",
        warmup=WARMUP_RUNS,
        runs=MEASURE_RUNS,
    )

pytorch_fp16_jit_results: dict[int, BenchmarkResult] = {}
pytorch_fp16_jit_memory_mb: dict[int, float] = {}
for batch in BATCH_SIZES:
    images_batch = [image] * batch
    with measure_memory(device="cuda") as mem:
        model.inference(dtype=torch.float16, batch_size=batch)
        _ = model.predict(images_batch)
    pytorch_fp16_jit_memory_mb[batch] = mem.delta_mb
    pytorch_fp16_jit_results[batch] = measure_latency(
        partial(model.predict, images_batch),
        label=f"PyTorch inference() fp16+JIT batch={batch}",
        device="cuda",
        warmup=WARMUP_RUNS,
        runs=MEASURE_RUNS,
    )
model.remove_optimized_model()

for batch in BATCH_SIZES:
    r32, r16 = pytorch_fp32_results[batch], pytorch_fp16_jit_results[batch]
    print(f"  batch={batch}  fp32  {r32.mean_ms:6.2f} ms ± {r32.std_ms:5.2f}   ({batch * r32.fps:6.1f} FPS)")
    print(
        f"  batch={batch}  fp16+JIT  {r16.mean_ms:6.2f} ms ± {r16.std_ms:5.2f}   ({batch * r16.fps:6.1f} FPS)"
        f"   +{pytorch_fp16_jit_memory_mb[batch]:.1f} MB"
    )

from functools import partial from rfdetr import RFDETRSmall with measure_memory(device="cuda") as mem: model = RFDETRSmall() baseline_detections = model.predict(image, threshold=CONFIDENCE_THRESHOLD) pytorch_fp32_memory_mb = mem.delta_mb print(f"PyTorch fp32 baseline: {len(baseline_detections)} detections above {CONFIDENCE_THRESHOLD}") visualize_detections(baseline_detections, image) pytorch_fp32_results: dict[int, BenchmarkResult] = {} for batch in BATCH_SIZES: images_batch = [image] * batch pytorch_fp32_results[batch] = measure_latency( partial(model.predict, images_batch), label=f"PyTorch predict() fp32 batch={batch}", device="cuda", warmup=WARMUP_RUNS, runs=MEASURE_RUNS, ) pytorch_fp16_jit_results: dict[int, BenchmarkResult] = {} pytorch_fp16_jit_memory_mb: dict[int, float] = {} for batch in BATCH_SIZES: images_batch = [image] * batch with measure_memory(device="cuda") as mem: model.inference(dtype=torch.float16, batch_size=batch) _ = model.predict(images_batch) pytorch_fp16_jit_memory_mb[batch] = mem.delta_mb pytorch_fp16_jit_results[batch] = measure_latency( partial(model.predict, images_batch), label=f"PyTorch inference() fp16+JIT batch={batch}", device="cuda", warmup=WARMUP_RUNS, runs=MEASURE_RUNS, ) model.remove_optimized_model() for batch in BATCH_SIZES: r32, r16 = pytorch_fp32_results[batch], pytorch_fp16_jit_results[batch] print(f" batch={batch} fp32 {r32.mean_ms:6.2f} ms ± {r32.std_ms:5.2f} ({batch * r32.fps:6.1f} FPS)") print( f" batch={batch} fp16+JIT {r16.mean_ms:6.2f} ms ± {r16.std_ms:5.2f} ({batch * r16.fps:6.1f} FPS)" f" +{pytorch_fp16_jit_memory_mb[batch]:.1f} MB" )

5. ONNX (CUDA execution provider)¶

What it is. ONNX (Open Neural Network Exchange) is a portable model-graph format most inference runtimes understand — ONNX Runtime, TensorRT, and OpenVINO can all consume the same .onnx file. Good for avoiding a single-vendor runtime lock-in: export once, then run it through whichever of those runtimes fits the target machine without re-exporting; on a CUDA GPU specifically, ONNX Runtime's CUDA execution provider is a faster-to-set-up alternative to a full TensorRT engine build. See the ONNX export docs.

Export¶

dynamic_batch=True bakes a dynamic batch axis into the graph (registry.py marks ONNX supports_dynamic_batch=True) — one exported .onnx file serves every batch size below.

In [ ]:

onnx_path = model.export(output_dir=str(EXPORT_DIR), dynamic_batch=True)
onnx_size_mb = _artifact_size_mb(onnx_path)
print(f"ONNX model: {onnx_path}  ({onnx_size_mb:.1f} MB)")

onnx_path = model.export(output_dir=str(EXPORT_DIR), dynamic_batch=True) onnx_size_mb = _artifact_size_mb(onnx_path) print(f"ONNX model: {onnx_path} ({onnx_size_mb:.1f} MB)")

Inference¶

_create_onnx_session (the same session-construction helper RFDETR.predict()'s reference decoder uses internally) is asked for CUDAExecutionProvider with a CPUExecutionProvider fallback; it raises if the CUDA provider silently fails to activate.

In [ ]:

from rfdetr.export._onnx.inference import _create_onnx_session, _run_inference

with measure_memory(device="cuda") as mem:
    onnx_session = _create_onnx_session(onnx_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
    active_provider = onnx_session.get_providers()[0]
    if active_provider != "CUDAExecutionProvider":
        raise RuntimeError(f"CUDAExecutionProvider not active (got {active_provider!r}) — see the install cell above.")
    onnx_detections, _ = _run_inference(onnx_session, IMAGE_PATH, threshold=CONFIDENCE_THRESHOLD)
onnx_memory_mb = mem.delta_mb
print(f"ONNX (CUDA): {len(onnx_detections)} detections above {CONFIDENCE_THRESHOLD}")
visualize_detections(onnx_detections, image, EXPORT_DIR / "annotated_onnx.jpg")

from rfdetr.export._onnx.inference import _create_onnx_session, _run_inference with measure_memory(device="cuda") as mem: onnx_session = _create_onnx_session(onnx_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"]) active_provider = onnx_session.get_providers()[0] if active_provider != "CUDAExecutionProvider": raise RuntimeError(f"CUDAExecutionProvider not active (got {active_provider!r}) — see the install cell above.") onnx_detections, _ = _run_inference(onnx_session, IMAGE_PATH, threshold=CONFIDENCE_THRESHOLD) onnx_memory_mb = mem.delta_mb print(f"ONNX (CUDA): {len(onnx_detections)} detections above {CONFIDENCE_THRESHOLD}") visualize_detections(onnx_detections, image, EXPORT_DIR / "annotated_onnx.jpg")

Benchmark — batch 1, 4, and 8¶

Two scopes per batch: forward_ms times only session.run with preprocessing done once outside the loop; end2end_ms times preprocessing + session.run + decoding every image in the batch, the same scope predict() above was timed at. Both scopes use synchronized host wall-clock timing; PyTorch CUDA events may miss work submitted on ONNX Runtime's stream. Memory [MB] is only meaningful once per shared session: the batch-1 row reuses onnx_memory_mb (session construction, from the Inference cell above); the batch-4 and batch-8 rows measure the incremental arena growth of the same session's first call at that larger batch, not a second construction.

In [ ]:

from rfdetr.export._runtime.decode import decode_detections
from rfdetr.export._runtime.preprocess import preprocess_to_nchw

_onnx_input_name = onnx_session.get_inputs()[0].name
_, _onnx_channels, _onnx_height, _onnx_width = onnx_session.get_inputs()[0].shape
_onnx_output_names = [out.name for out in onnx_session.get_outputs()]
_onnx_boxes_idx = next(i for i, name in enumerate(_onnx_output_names) if "dets" in name)
_onnx_logits_idx = next(i for i, name in enumerate(_onnx_output_names) if "labels" in name)


def _onnx_single_input() -> np.ndarray:
    return preprocess_to_nchw(image, _onnx_height, _onnx_width, _onnx_channels)


def _make_onnx_end2end(batch: int):
    def _run() -> None:
        feed = {_onnx_input_name: _tile_batch(_onnx_single_input(), batch)}
        raw = onnx_session.run(None, feed)
        _decode_batch(raw[_onnx_boxes_idx], raw[_onnx_logits_idx], image.size, batch, threshold=CONFIDENCE_THRESHOLD)

    return _run


onnx_forward_results: dict[int, BenchmarkResult] = {}
onnx_end2end_results: dict[int, BenchmarkResult] = {}
onnx_memory_by_batch: dict[int, float] = {BATCH_SIZES[0]: onnx_memory_mb}
for batch in BATCH_SIZES:
    feed = {_onnx_input_name: _tile_batch(_onnx_single_input(), batch)}
    if batch != BATCH_SIZES[0]:
        with measure_memory(device="cuda") as mem:
            onnx_session.run(None, feed)
        onnx_memory_by_batch[batch] = mem.delta_mb
    onnx_forward_results[batch] = measure_latency(
        partial(onnx_session.run, None, feed),
        label=f"ONNX (CUDA) forward batch={batch}",
        device="cpu",
        warmup=WARMUP_RUNS,
        runs=MEASURE_RUNS,
    )
    onnx_end2end_results[batch] = measure_latency(
        _make_onnx_end2end(batch),
        label=f"ONNX (CUDA) end2end batch={batch}",
        device="cpu",
        warmup=WARMUP_RUNS,
        runs=MEASURE_RUNS,
    )
    fwd, e2e = onnx_forward_results[batch], onnx_end2end_results[batch]
    print(f"  batch={batch}  forward  {fwd.mean_ms:6.2f} ms ± {fwd.std_ms:5.2f}")
    print(f"  batch={batch}  end2end  {e2e.mean_ms:6.2f} ms ± {e2e.std_ms:5.2f}   ({batch * e2e.fps:6.1f} FPS)")

from rfdetr.export._runtime.decode import decode_detections from rfdetr.export._runtime.preprocess import preprocess_to_nchw _onnx_input_name = onnx_session.get_inputs()[0].name _, _onnx_channels, _onnx_height, _onnx_width = onnx_session.get_inputs()[0].shape _onnx_output_names = [out.name for out in onnx_session.get_outputs()] _onnx_boxes_idx = next(i for i, name in enumerate(_onnx_output_names) if "dets" in name) _onnx_logits_idx = next(i for i, name in enumerate(_onnx_output_names) if "labels" in name) def _onnx_single_input() -> np.ndarray: return preprocess_to_nchw(image, _onnx_height, _onnx_width, _onnx_channels) def _make_onnx_end2end(batch: int): def _run() -> None: feed = {_onnx_input_name: _tile_batch(_onnx_single_input(), batch)} raw = onnx_session.run(None, feed) _decode_batch(raw[_onnx_boxes_idx], raw[_onnx_logits_idx], image.size, batch, threshold=CONFIDENCE_THRESHOLD) return _run onnx_forward_results: dict[int, BenchmarkResult] = {} onnx_end2end_results: dict[int, BenchmarkResult] = {} onnx_memory_by_batch: dict[int, float] = {BATCH_SIZES[0]: onnx_memory_mb} for batch in BATCH_SIZES: feed = {_onnx_input_name: _tile_batch(_onnx_single_input(), batch)} if batch != BATCH_SIZES[0]: with measure_memory(device="cuda") as mem: onnx_session.run(None, feed) onnx_memory_by_batch[batch] = mem.delta_mb onnx_forward_results[batch] = measure_latency( partial(onnx_session.run, None, feed), label=f"ONNX (CUDA) forward batch={batch}", device="cpu", warmup=WARMUP_RUNS, runs=MEASURE_RUNS, ) onnx_end2end_results[batch] = measure_latency( _make_onnx_end2end(batch), label=f"ONNX (CUDA) end2end batch={batch}", device="cpu", warmup=WARMUP_RUNS, runs=MEASURE_RUNS, ) fwd, e2e = onnx_forward_results[batch], onnx_end2end_results[batch] print(f" batch={batch} forward {fwd.mean_ms:6.2f} ms ± {fwd.std_ms:5.2f}") print(f" batch={batch} end2end {e2e.mean_ms:6.2f} ms ± {e2e.std_ms:5.2f} ({batch * e2e.fps:6.1f} FPS)")

6. TensorRT¶

What it is. TensorRT is NVIDIA's own inference optimizer and runtime — it compiles a model into an engine specialized for the exact GPU architecture and TensorRT version that built it. Good for squeezing the lowest possible latency out of an NVIDIA GPU in production, at the cost of that engine being locked to the machine (and TensorRT version) it was built on — convert on the same GPU family you plan to deploy to. See the TensorRT export docs.

Export¶

format="trt" compiles a standalone FP16 engine (tens of seconds to a few minutes).

In [ ]:

trt_engine_path = model.export(format="trt", output_dir=str(EXPORT_DIR))
trt_size_mb = _artifact_size_mb(trt_engine_path)
print(f"TensorRT engine: {trt_engine_path}  ({trt_size_mb:.1f} MB)")

trt_engine_path = model.export(format="trt", output_dir=str(EXPORT_DIR)) trt_size_mb = _artifact_size_mb(trt_engine_path) print(f"TensorRT engine: {trt_engine_path} ({trt_size_mb:.1f} MB)")

Inference¶

Inference goes through inference-models, which selects and loads its own prebuilt TensorRT engine (the managed engine rows below) — the easier, portable path that also handles preprocessing, postprocessing, and class names — rather than loading the .trt file exported above. Its precision is controlled by AutoModel.from_pretrained(..., quantization=...), independent of the fp16 kwarg used for the standalone .trt export: quantization=None (default) picks automatically based on device capabilities, and "fp32"/"fp16"/"bf16"/"int8" request one explicitly. verbose=True prints inference-models' own package-selection log line, which names the backend and the quantization it actually picked for quantization=None below. It is constructed once and fed every batch size below through its documented List[np.ndarray] call signature — what it does per batch internally is not part of this repository and not something this notebook can verify, only measure.

In [ ]:

from inference_models import AutoModel, BackendType

with measure_memory(device="cuda") as mem:
    trt_model = AutoModel.from_pretrained("rfdetr-small", backend=BackendType.TRT, verbose=True)
    trt_predictions = trt_model(np.array(image))
trt_memory_mb: dict[int, float] = {BATCH_SIZES[0]: mem.delta_mb}

trt_detections = trt_predictions[0].to_supervision()
trt_detections = trt_detections[trt_detections.confidence > CONFIDENCE_THRESHOLD]
print(f"TensorRT: {len(trt_detections)} detections above {CONFIDENCE_THRESHOLD}")
visualize_detections(trt_detections, image, EXPORT_DIR / "annotated_tensorrt.jpg")

from inference_models import AutoModel, BackendType with measure_memory(device="cuda") as mem: trt_model = AutoModel.from_pretrained("rfdetr-small", backend=BackendType.TRT, verbose=True) trt_predictions = trt_model(np.array(image)) trt_memory_mb: dict[int, float] = {BATCH_SIZES[0]: mem.delta_mb} trt_detections = trt_predictions[0].to_supervision() trt_detections = trt_detections[trt_detections.confidence > CONFIDENCE_THRESHOLD] print(f"TensorRT: {len(trt_detections)} detections above {CONFIDENCE_THRESHOLD}") visualize_detections(trt_detections, image, EXPORT_DIR / "annotated_tensorrt.jpg")

Benchmark — batch 1, 4, and 8¶

inference-models does not expose a forward-only hook separate from its own preprocessing, so only an end-to-end number is reported here — the same limitation as the PyTorch baseline above. Memory [MB] for batch 1 reuses the construction cost from the Inference cell above; batch 4 and 8 measure the incremental cost of this same model's first call at that batch.

In [ ]:

trt_end2end_results: dict[int, BenchmarkResult] = {}
for batch in BATCH_SIZES:
    images_batch = [np.array(image)] * batch
    if batch != BATCH_SIZES[0]:
        with measure_memory(device="cuda") as mem:
            trt_model(images_batch)
        trt_memory_mb[batch] = mem.delta_mb
    trt_end2end_results[batch] = measure_latency(
        partial(trt_model, images_batch),
        label=f"TensorRT (inference-models, managed engine) auto end2end batch={batch}",
        device="cuda",
        warmup=WARMUP_RUNS,
        runs=MEASURE_RUNS,
    )
    r = trt_end2end_results[batch]
    print(f"  batch={batch}  {r.mean_ms:6.2f} ms ± {r.std_ms:5.2f}   ({batch * r.fps:6.1f} FPS)")

trt_end2end_results: dict[int, BenchmarkResult] = {} for batch in BATCH_SIZES: images_batch = [np.array(image)] * batch if batch != BATCH_SIZES[0]: with measure_memory(device="cuda") as mem: trt_model(images_batch) trt_memory_mb[batch] = mem.delta_mb trt_end2end_results[batch] = measure_latency( partial(trt_model, images_batch), label=f"TensorRT (inference-models, managed engine) auto end2end batch={batch}", device="cuda", warmup=WARMUP_RUNS, runs=MEASURE_RUNS, ) r = trt_end2end_results[batch] print(f" batch={batch} {r.mean_ms:6.2f} ms ± {r.std_ms:5.2f} ({batch * r.fps:6.1f} FPS)")

TensorRT (fp32, explicit)¶

quantization="fp32" asks inference-models for its full-precision prebuilt engine — the numeric-parity anchor against the auto-selected (typically fp16 on a capable GPU) row above.

In [ ]:

with measure_memory(device="cuda") as mem:
    trt_fp32_model = AutoModel.from_pretrained(
        "rfdetr-small", backend=BackendType.TRT, quantization="fp32", verbose=True
    )
    trt_fp32_predictions = trt_fp32_model(np.array(image))
trt_fp32_memory_mb: dict[int, float] = {BATCH_SIZES[0]: mem.delta_mb}

trt_fp32_detections = trt_fp32_predictions[0].to_supervision()
trt_fp32_detections = trt_fp32_detections[trt_fp32_detections.confidence > CONFIDENCE_THRESHOLD]
print(f"TensorRT fp32: {len(trt_fp32_detections)} detections above {CONFIDENCE_THRESHOLD}")
visualize_detections(trt_fp32_detections, image, EXPORT_DIR / "annotated_tensorrt_fp32.jpg")

trt_fp32_end2end_results: dict[int, BenchmarkResult] = {}
for batch in BATCH_SIZES:
    images_batch = [np.array(image)] * batch
    if batch != BATCH_SIZES[0]:
        with measure_memory(device="cuda") as mem:
            trt_fp32_model(images_batch)
        trt_fp32_memory_mb[batch] = mem.delta_mb
    trt_fp32_end2end_results[batch] = measure_latency(
        partial(trt_fp32_model, images_batch),
        label=f"TensorRT (inference-models, managed engine) fp32 end2end batch={batch}",
        device="cuda",
        warmup=WARMUP_RUNS,
        runs=MEASURE_RUNS,
    )
    r = trt_fp32_end2end_results[batch]
    print(f"  batch={batch}  {r.mean_ms:6.2f} ms ± {r.std_ms:5.2f}   ({batch * r.fps:6.1f} FPS)")

with measure_memory(device="cuda") as mem: trt_fp32_model = AutoModel.from_pretrained( "rfdetr-small", backend=BackendType.TRT, quantization="fp32", verbose=True ) trt_fp32_predictions = trt_fp32_model(np.array(image)) trt_fp32_memory_mb: dict[int, float] = {BATCH_SIZES[0]: mem.delta_mb} trt_fp32_detections = trt_fp32_predictions[0].to_supervision() trt_fp32_detections = trt_fp32_detections[trt_fp32_detections.confidence > CONFIDENCE_THRESHOLD] print(f"TensorRT fp32: {len(trt_fp32_detections)} detections above {CONFIDENCE_THRESHOLD}") visualize_detections(trt_fp32_detections, image, EXPORT_DIR / "annotated_tensorrt_fp32.jpg") trt_fp32_end2end_results: dict[int, BenchmarkResult] = {} for batch in BATCH_SIZES: images_batch = [np.array(image)] * batch if batch != BATCH_SIZES[0]: with measure_memory(device="cuda") as mem: trt_fp32_model(images_batch) trt_fp32_memory_mb[batch] = mem.delta_mb trt_fp32_end2end_results[batch] = measure_latency( partial(trt_fp32_model, images_batch), label=f"TensorRT (inference-models, managed engine) fp32 end2end batch={batch}", device="cuda", warmup=WARMUP_RUNS, runs=MEASURE_RUNS, ) r = trt_fp32_end2end_results[batch] print(f" batch={batch} {r.mean_ms:6.2f} ms ± {r.std_ms:5.2f} ({batch * r.fps:6.1f} FPS)")

TensorRT (raw .trt engine, TRTInference)¶

The rows above never touch the .trt file model.export(format="trt") wrote — this section loads that file directly, with rfdetr.export._tensorrt.inference.TRTInference, the package's own reference runtime for a raw engine. It exists for export/parity testing, not as a recommended production path — inference-models above is that path — but it is the only way to time the .trt artifact itself rather than a separately-built managed engine. sync_mode=True avoids the pycuda dependency [tensorrt-bench] would add for the async/stream path; device="cuda" pins it to the current device without a second resolve. dynamic_batch=True, max_batch_size=8 on the export bakes an optimization profile spanning batch 1–8 into one engine, so this single TRTInference instance serves every batch size below; TRTInference declares each call's input shape on the execution context automatically, no separate export per batch needed here (unlike CoreML/Core AI/ExecuTorch on Apple Silicon, or LiteRT/ExecuTorch-XNNPACK on mobile, which do bake a fixed batch shape and need one artifact per batch size — see those cookbooks). Nothing casts the input on the caller's behalf, so the tensor is built in the engine's own input dtype (fp16 unless the export forced fp32 I/O) and made contiguous before the call — a hidden cast here would be timed as part of the engine.

In [ ]:

trt_raw_engine_path = model.export(
    format="trt", dynamic_batch=True, max_batch_size=max(BATCH_SIZES), output_dir=str(EXPORT_DIR)
)
trt_raw_engine_size_mb = _artifact_size_mb(trt_raw_engine_path)
print(f"TensorRT engine (dynamic batch 1-{max(BATCH_SIZES)}): {trt_raw_engine_path}  ({trt_raw_engine_size_mb:.1f} MB)")

trt_raw_engine_path = model.export( format="trt", dynamic_batch=True, max_batch_size=max(BATCH_SIZES), output_dir=str(EXPORT_DIR) ) trt_raw_engine_size_mb = _artifact_size_mb(trt_raw_engine_path) print(f"TensorRT engine (dynamic batch 1-{max(BATCH_SIZES)}): {trt_raw_engine_path} ({trt_raw_engine_size_mb:.1f} MB)")

In [ ]:

from rfdetr.export._tensorrt.inference import TRTInference


def _trt_raw_preprocess(batch: int) -> torch.Tensor:
    single = preprocess_to_nchw(image, _trt_raw_height, _trt_raw_width, _trt_raw_channels)
    tiled = _tile_batch(single, batch)
    return torch.from_numpy(tiled).to(device="cuda", dtype=_trt_raw_input_dtype).contiguous()


def _make_trt_raw_end2end(batch: int):
    def _run() -> None:
        blob = {_trt_raw_input_name: _trt_raw_preprocess(batch)}
        outputs = trt_raw_runtime(blob)
        boxes = outputs[_trt_raw_boxes_name].detach().float().cpu().numpy()
        logits = outputs[_trt_raw_logits_name].detach().float().cpu().numpy()
        _decode_batch(boxes, logits, image.size, batch, threshold=CONFIDENCE_THRESHOLD)

    return _run

from rfdetr.export._tensorrt.inference import TRTInference def _trt_raw_preprocess(batch: int) -> torch.Tensor: single = preprocess_to_nchw(image, _trt_raw_height, _trt_raw_width, _trt_raw_channels) tiled = _tile_batch(single, batch) return torch.from_numpy(tiled).to(device="cuda", dtype=_trt_raw_input_dtype).contiguous() def _make_trt_raw_end2end(batch: int): def _run() -> None: blob = {_trt_raw_input_name: _trt_raw_preprocess(batch)} outputs = trt_raw_runtime(blob) boxes = outputs[_trt_raw_boxes_name].detach().float().cpu().numpy() logits = outputs[_trt_raw_logits_name].detach().float().cpu().numpy() _decode_batch(boxes, logits, image.size, batch, threshold=CONFIDENCE_THRESHOLD) return _run

Inference¶

One correctness call before benchmarking, same as every other format above. Memory [MB] in the results table below covers constructing TRTInference (engine deserialize + output-buffer cudaMalloc, sized for the profile maximum, batch 8) plus this first call -- both happen inside the same measure_memory window, same scope as every other format's memory number. Building the runtime outside this window would read back ~0 MB, since its first call reuses buffers TRTInference already allocated at construction rather than allocating anything new.

In [ ]:

with measure_memory(device="cuda") as mem:
    trt_raw_runtime = TRTInference(trt_raw_engine_path, device="cuda", sync_mode=True)

    _trt_raw_input_name = trt_raw_runtime.input_names[0]
    _trt_raw_input_binding = trt_raw_runtime.bindings[_trt_raw_input_name]
    _, _trt_raw_channels, _trt_raw_height, _trt_raw_width = _trt_raw_input_binding.shape
    _trt_raw_input_dtype = torch.from_numpy(np.empty(0, dtype=_trt_raw_input_binding.dtype)).dtype
    _trt_raw_boxes_name = next(name for name in trt_raw_runtime.output_names if "dets" in name)
    _trt_raw_logits_name = next(name for name in trt_raw_runtime.output_names if "labels" in name)

    _trt_raw_input_b1 = {_trt_raw_input_name: _trt_raw_preprocess(BATCH_SIZES[0])}
    trt_raw_outputs = trt_raw_runtime(_trt_raw_input_b1)
trt_raw_memory_mb: dict[int, float] = {BATCH_SIZES[0]: mem.delta_mb}

trt_raw_boxes = trt_raw_outputs[_trt_raw_boxes_name].detach().float().cpu().numpy()[0]
trt_raw_logits = trt_raw_outputs[_trt_raw_logits_name].detach().float().cpu().numpy()[0]
trt_raw_decoded = decode_detections(trt_raw_boxes, trt_raw_logits, image.size, threshold=CONFIDENCE_THRESHOLD)
trt_raw_detections = sv.Detections(
    xyxy=trt_raw_decoded.xyxy,
    confidence=trt_raw_decoded.confidence,
    class_id=trt_raw_decoded.class_id.astype(int),
)
print(f"TensorRT (raw engine): {len(trt_raw_detections)} detections above {CONFIDENCE_THRESHOLD}")
visualize_detections(trt_raw_detections, image, EXPORT_DIR / "annotated_tensorrt_raw.jpg")

with measure_memory(device="cuda") as mem: trt_raw_runtime = TRTInference(trt_raw_engine_path, device="cuda", sync_mode=True) _trt_raw_input_name = trt_raw_runtime.input_names[0] _trt_raw_input_binding = trt_raw_runtime.bindings[_trt_raw_input_name] _, _trt_raw_channels, _trt_raw_height, _trt_raw_width = _trt_raw_input_binding.shape _trt_raw_input_dtype = torch.from_numpy(np.empty(0, dtype=_trt_raw_input_binding.dtype)).dtype _trt_raw_boxes_name = next(name for name in trt_raw_runtime.output_names if "dets" in name) _trt_raw_logits_name = next(name for name in trt_raw_runtime.output_names if "labels" in name) _trt_raw_input_b1 = {_trt_raw_input_name: _trt_raw_preprocess(BATCH_SIZES[0])} trt_raw_outputs = trt_raw_runtime(_trt_raw_input_b1) trt_raw_memory_mb: dict[int, float] = {BATCH_SIZES[0]: mem.delta_mb} trt_raw_boxes = trt_raw_outputs[_trt_raw_boxes_name].detach().float().cpu().numpy()[0] trt_raw_logits = trt_raw_outputs[_trt_raw_logits_name].detach().float().cpu().numpy()[0] trt_raw_decoded = decode_detections(trt_raw_boxes, trt_raw_logits, image.size, threshold=CONFIDENCE_THRESHOLD) trt_raw_detections = sv.Detections( xyxy=trt_raw_decoded.xyxy, confidence=trt_raw_decoded.confidence, class_id=trt_raw_decoded.class_id.astype(int), ) print(f"TensorRT (raw engine): {len(trt_raw_detections)} detections above {CONFIDENCE_THRESHOLD}") visualize_detections(trt_raw_detections, image, EXPORT_DIR / "annotated_tensorrt_raw.jpg")

Benchmark — batch 1, 4, and 8¶

forward times only the TRTInference call, input already on-device and in the right dtype; end2end adds preprocessing and decoding — the same two scopes ONNX above was split into. Batch 4 and 8 measure the incremental memory of this same engine's first call at that batch, the same convention as ONNX and the managed engine above.

In [ ]:

trt_raw_forward_results: dict[int, BenchmarkResult] = {}
trt_raw_end2end_results: dict[int, BenchmarkResult] = {}
for batch in BATCH_SIZES:
    trt_raw_input = {_trt_raw_input_name: _trt_raw_preprocess(batch)}
    if batch != BATCH_SIZES[0]:
        with measure_memory(device="cuda") as mem:
            trt_raw_runtime(trt_raw_input)
        trt_raw_memory_mb[batch] = mem.delta_mb
    trt_raw_forward_results[batch] = measure_latency(
        partial(trt_raw_runtime, trt_raw_input),
        label=f"TensorRT (raw .trt engine) forward batch={batch}",
        device="cuda",
        warmup=WARMUP_RUNS,
        runs=MEASURE_RUNS,
    )
    trt_raw_end2end_results[batch] = measure_latency(
        _make_trt_raw_end2end(batch),
        label=f"TensorRT (raw .trt engine) end2end batch={batch}",
        device="cuda",
        warmup=WARMUP_RUNS,
        runs=MEASURE_RUNS,
    )
    fwd, e2e = trt_raw_forward_results[batch], trt_raw_end2end_results[batch]
    print(f"  batch={batch}  forward  {fwd.mean_ms:6.2f} ms ± {fwd.std_ms:5.2f}")
    print(f"  batch={batch}  end2end  {e2e.mean_ms:6.2f} ms ± {e2e.std_ms:5.2f}   ({batch * e2e.fps:6.1f} FPS)")

trt_raw_forward_results: dict[int, BenchmarkResult] = {} trt_raw_end2end_results: dict[int, BenchmarkResult] = {} for batch in BATCH_SIZES: trt_raw_input = {_trt_raw_input_name: _trt_raw_preprocess(batch)} if batch != BATCH_SIZES[0]: with measure_memory(device="cuda") as mem: trt_raw_runtime(trt_raw_input) trt_raw_memory_mb[batch] = mem.delta_mb trt_raw_forward_results[batch] = measure_latency( partial(trt_raw_runtime, trt_raw_input), label=f"TensorRT (raw .trt engine) forward batch={batch}", device="cuda", warmup=WARMUP_RUNS, runs=MEASURE_RUNS, ) trt_raw_end2end_results[batch] = measure_latency( _make_trt_raw_end2end(batch), label=f"TensorRT (raw .trt engine) end2end batch={batch}", device="cuda", warmup=WARMUP_RUNS, runs=MEASURE_RUNS, ) fwd, e2e = trt_raw_forward_results[batch], trt_raw_end2end_results[batch] print(f" batch={batch} forward {fwd.mean_ms:6.2f} ms ± {fwd.std_ms:5.2f}") print(f" batch={batch} end2end {e2e.mean_ms:6.2f} ms ± {e2e.std_ms:5.2f} ({batch * e2e.fps:6.1f} FPS)")

7. Results¶

Measured on an NVIDIA L4 (compute capability 8.9, 22 GB), Colab GPU runtime, rfdetr <VERSION>, batches 1/4/8, 20 warmup + 100 timed runs per row, RFDETRSmall. Numbers vary by GPU model, TensorRT version, and driver — rerun the cells below to get yours. Config is the precision/EP used for that row; — marks a scope this format doesn't have. forward [ms] / end2end [ms] are per-call latency for the whole batch (a batch-8 call takes longer wall time than a batch-1 call because it does 8× the work); FPS [img/s] is batch × 1000 / end2end_ms, images processed per second, not calls per second — the column to compare across batch sizes. Memory [MB] is device-memory growth (torch.cuda.mem_get_info()) — full construction-plus-first-call cost at batch 1, incremental first-larger-batch-call growth at batch 4 and 8 (see each benchmark cell's own note above for which applies to that row) — device-wide, so it also captures ONNX Runtime's CUDA execution provider and TensorRT's own allocations, not just PyTorch's caching allocator.

Format	Batch	Config	forward [ms]	end2end [ms]	FPS [img/s] (end2end)	Memory [MB]
PyTorch predict()	1	fp32	—	19.88 ± 0.47	50.3	—
PyTorch inference()	1	fp16+JIT	—	11.37 ± 0.75	88.0	161.5
ONNX	1	CUDA EP	7.35 ± 0.07	13.78 ± 0.55	72.6	291.5
TensorRT (inference-models)	1	managed engine, auto	—	3.41 ± 0.06	293.5	140.5
TensorRT (inference-models)	1	managed engine, fp32	—	6.62 ± 0.05	151.0	228.6
TensorRT (raw .trt engine)	1	TRTInference, fp16	2.16 ± 0.03	9.82 ± 0.34	101.8	182.5
PyTorch predict()	4	fp32	—	32.78 ± 0.46	122.0	—
PyTorch inference()	4	fp16+JIT	—	17.78 ± 0.40	225.0	381.7
ONNX	4	CUDA EP	28.80 ± 0.49	35.61 ± 0.56	112.3	268.4
TensorRT (inference-models)	4	managed engine, auto	—	10.75 ± 0.30	372.2	0.0
TensorRT (inference-models)	4	managed engine, fp32	—	24.82 ± 0.52	161.2	0.0
TensorRT (raw .trt engine)	4	TRTInference, fp16	7.40 ± 0.13	20.94 ± 0.24	191.0	0.0
PyTorch predict()	8	fp32	—	67.43 ± 0.67	118.6	—
PyTorch inference()	8	fp16+JIT	—	41.24 ± 0.83	194.0	908.1
ONNX	8	CUDA EP	64.11 ± 0.51	76.97 ± 0.48	103.9	536.9
TensorRT (inference-models)	8	managed engine, auto	—	19.95 ± 0.21	401.1	0.0
TensorRT (inference-models)	8	managed engine, fp32	—	49.18 ± 0.71	162.7	0.0
TensorRT (raw .trt engine)	8	TRTInference, fp16	14.83 ± 0.17	36.03 ± 0.66	222.0	0.0

Measured on an NVIDIA L4, Colab GPU runtime, 100 timed + 20 warmup runs per row. Every batch-4/8 TensorRT row reads 0.0 MB, on both the managed engine and the raw engine, and this is expected, not a bug: the raw engine is exported dynamic_batch=True, max_batch_size=8, so TRTInference sizes every output buffer for the batch-8 profile maximum once, at construction — batch 1's row pays that full cost (182.5 MB), batch 4 and 8 then allocate nothing further, because the buffers they need already exist. The managed engine's 0.0 at batch 4/8 cannot be verified the same way (inference-models' engine build is not part of this repository), but the same explanation — headroom already reserved at construction — is the consistent one.

TensorRT's auto-selected precision is the fastest row at every batch and scales the best: 293.5 → 372.2 → 401.1 FPS across batch 1/4/8, a 1.37× throughput gain from batching alone, well ahead of forcing fp32 (151.0 → 161.2 → 162.7 FPS, batching barely moves it). The raw .trt engine's forward scales close to linearly with batch — 2.16 → 7.40 → 14.83 ms is 3.43× at batch 4 and 6.87× at batch 8, both under the batch multiplier, so per-image cost keeps dropping as the batch grows. ONNX's forward/end2end gap is ~6.5 ms at batch 1–4 (6.43 ms, 6.81 ms) and roughly doubles at batch 8 (12.86 ms) — the CPU-side pre/postprocessing grows far more slowly than the GPU-side forward pass, so it becomes a smaller fraction of the total at larger batches.

Open gap: these numbers vs Roboflow's product page. This notebook, L4, batch 1: 3.41 ms end-to-end for the managed engine, 2.16 ms forward for the raw .trt engine, 7.35 ms forward for ONNX. Roboflow's RF-DETR model page, L4, Roboflow Inference, batch 1: 8.3 ms TensorRT and 12.9 ms ONNX. The model page does not state what its timing covers; the two harnesses have not been run side by side.

In [ ]:

import pandas as pd
from IPython.display import display

_rows = []
for batch in BATCH_SIZES:
    _rows += [
        _result_row("PyTorch predict()", batch, "fp32", None, pytorch_fp32_results[batch], None),
        _result_row(
            "PyTorch inference()",
            batch,
            "fp16+JIT",
            None,
            pytorch_fp16_jit_results[batch],
            pytorch_fp16_jit_memory_mb[batch],
        ),
        _result_row(
            "ONNX",
            batch,
            "CUDA EP",
            onnx_forward_results[batch],
            onnx_end2end_results[batch],
            onnx_memory_by_batch.get(batch),
        ),
        _result_row(
            "TensorRT (inference-models)",
            batch,
            "managed engine, auto",
            None,
            trt_end2end_results[batch],
            trt_memory_mb.get(batch),
        ),
        _result_row(
            "TensorRT (inference-models)",
            batch,
            "managed engine, fp32",
            None,
            trt_fp32_end2end_results[batch],
            trt_fp32_memory_mb.get(batch),
        ),
        _result_row(
            "TensorRT (raw .trt engine)",
            batch,
            "TRTInference, fp16",
            trt_raw_forward_results[batch],
            trt_raw_end2end_results[batch],
            trt_raw_memory_mb.get(batch),
        ),
    ]

summary = pd.DataFrame(_rows).set_index(["Format", "Batch"])
display(summary)
print(
    f"\n{MEASURE_RUNS} timed + {WARMUP_RUNS} warmup runs per row, batches {BATCH_SIZES}, GPU: {torch.cuda.get_device_name(0)}."
)

import pandas as pd from IPython.display import display _rows = [] for batch in BATCH_SIZES: _rows += [ _result_row("PyTorch predict()", batch, "fp32", None, pytorch_fp32_results[batch], None), _result_row( "PyTorch inference()", batch, "fp16+JIT", None, pytorch_fp16_jit_results[batch], pytorch_fp16_jit_memory_mb[batch], ), _result_row( "ONNX", batch, "CUDA EP", onnx_forward_results[batch], onnx_end2end_results[batch], onnx_memory_by_batch.get(batch), ), _result_row( "TensorRT (inference-models)", batch, "managed engine, auto", None, trt_end2end_results[batch], trt_memory_mb.get(batch), ), _result_row( "TensorRT (inference-models)", batch, "managed engine, fp32", None, trt_fp32_end2end_results[batch], trt_fp32_memory_mb.get(batch), ), _result_row( "TensorRT (raw .trt engine)", batch, "TRTInference, fp16", trt_raw_forward_results[batch], trt_raw_end2end_results[batch], trt_raw_memory_mb.get(batch), ), ] summary = pd.DataFrame(_rows).set_index(["Format", "Batch"]) display(summary) print( f"\n{MEASURE_RUNS} timed + {WARMUP_RUNS} warmup runs per row, batches {BATCH_SIZES}, GPU: {torch.cuda.get_device_name(0)}." )

Next steps¶

• Fine-tuned weights — pass pretrain_weights="<path/to/checkpoint.pth>" when constructing the model.
• Not on a CUDA GPU? — see the CPU cookbook, the mobile/edge cookbook, or the Apple Silicon cookbook.
• See the Export documentation for every format and option.