roboflow--rf-detr/docs/cookbooks/inference-latency-benchmark.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5720dea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------------------------------------------------------------\n",
    "# RF-DETR\n",
    "# Copyright (c) 2025 Roboflow. All Rights Reserved.\n",
    "# Licensed under the Apache License, Version 2.0 [see LICENSE for details]\n",
    "# ------------------------------------------------------------------------"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66c62889",
   "metadata": {},
   "source": [
    "# RF-DETR Inference Latency Benchmark\n",
    "\n",
    "Measures inference latency for three RF-DETR families across three configs:\n",
    "\n",
    "| Config | Description |\n",
    "|--------|-------------|\n",
    "| **FP32** | `predict()` — unoptimized baseline |\n",
    "| **FP16+JIT** | `optimize_for_inference(dtype=torch.float16)` |\n",
    "| **ONNX** | exported `.onnx` via `onnxruntime-gpu` |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1568b82",
   "metadata": {},
   "source": [
    "## 1. Install\n",
    "\n",
    "We need `onnxruntime-gpu` built against CUDA 12 — the default PyPI wheel targets CUDA 11.8 and silently\n",
    "falls back to CPU on modern GPUs.  The Microsoft CUDA-12 package index ships the correct build.\n",
    "\n",
    "> **Colab**: after running this cell, go to **Runtime → Restart session**, then run from the next cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a116013",
   "metadata": {
    "title": "[bash]"
   },
   "outputs": [],
   "source": [
    "!pip uninstall -y onnxruntime onnxruntime-gpu\n",
    "!pip install -q \"rfdetr[onnx]\" pillow pandas\n",
    "!pip install -q onnxruntime-gpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-12/pypi/simple/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ca0838b",
   "metadata": {},
   "source": [
    "## 2. Config\n",
    "\n",
    "`WARMUP_RUNS` discards the first N inferences — GPU kernels are JIT-compiled on first use, so early timings\n",
    "are outliers.  `MEASURE_RUNS` then collects the steady-state distribution.  20 + 100 is a reasonable balance\n",
    "between statistical stability and total wall-clock time per model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e47b983",
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections.abc import Callable\n",
    "from pathlib import Path\n",
    "from typing import Any, NamedTuple\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "\n",
    "WARMUP_RUNS = 20\n",
    "MEASURE_RUNS = 100\n",
    "EXPORT_DIR = Path(\"benchmark_output\")\n",
    "EXPORT_DIR.mkdir(exist_ok=True)\n",
    "\n",
    "if not torch.cuda.is_available():\n",
    "    raise RuntimeError(\"This benchmark requires a CUDA GPU.\")\n",
    "print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
    "\n",
    "import onnxruntime as ort\n",
    "\n",
    "_ort_providers = ort.get_available_providers()\n",
    "print(f\"ORT {ort.__version__}, providers: {_ort_providers}\")\n",
    "if \"CUDAExecutionProvider\" not in _ort_providers:\n",
    "    raise RuntimeError(\n",
    "        f\"onnxruntime-gpu with CUDA support required. Available providers: {_ort_providers}. \"\n",
    "        \"Fix: reinstall from the CUDA-12 index (see install cell) then restart runtime.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ccdc87e",
   "metadata": {},
   "source": [
    "## 3. Sample images\n",
    "\n",
    "Latency depends on resolution, not pixel content, so synthetic noise images are equivalent to real photos\n",
    "for benchmarking purposes.  Using a fixed seed makes results reproducible across runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06a16bb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(42)\n",
    "images: list[Image.Image] = [Image.fromarray(rng.integers(0, 256, (640, 640, 3), dtype=np.uint8)) for _ in range(10)]\n",
    "print(f\"Generated {len(images)} synthetic 640×640 RGB images\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db0d3d7c",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 4. Latency helpers\n",
    "\n",
    "GPU kernels execute asynchronously — `time.perf_counter()` returns before the GPU finishes, giving\n",
    "misleadingly low numbers.  CUDA events are inserted directly into the GPU command stream and timestamped\n",
    "on the device, so `elapsed_time()` measures actual kernel execution.  `torch.cuda.synchronize()` after\n",
    "each run flushes the stream and ensures the event fires before we read it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f928319",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BenchmarkResult(NamedTuple):\n",
    "    \"\"\"Single benchmark measurement.\"\"\"\n",
    "\n",
    "    label: str\n",
    "    mean_ms: float\n",
    "    std_ms: float\n",
    "\n",
    "    @property\n",
    "    def fps(self) -> float:\n",
    "        \"\"\"Frames per second.\"\"\"\n",
    "        return 1000.0 / self.mean_ms\n",
    "\n",
    "\n",
    "def measure_latency_gpu(\n",
    "    fn: Callable[[], object],\n",
    "    warmup: int = WARMUP_RUNS,\n",
    "    runs: int = MEASURE_RUNS,\n",
    ") -> tuple[float, float]:\n",
    "    \"\"\"Return (mean_ms, std_ms) using CUDA events.\"\"\"\n",
    "    for _ in range(warmup):\n",
    "        fn()\n",
    "    torch.cuda.synchronize()\n",
    "    start = torch.cuda.Event(enable_timing=True)\n",
    "    end = torch.cuda.Event(enable_timing=True)\n",
    "    timings: list[float] = []\n",
    "    for _ in range(runs):\n",
    "        start.record()\n",
    "        fn()\n",
    "        end.record()\n",
    "        torch.cuda.synchronize()\n",
    "        timings.append(start.elapsed_time(end))\n",
    "    arr = np.array(timings)\n",
    "    return float(arr.mean()), float(arr.std())\n",
    "\n",
    "\n",
    "_measure = measure_latency_gpu"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f7ee4f5",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 5. Per-config benchmark functions\n",
    "\n",
    "Three inference paths are compared:\n",
    "\n",
    "- **FP32** — `predict()` as shipped.  Full 32-bit arithmetic on GPU.  Baseline.\n",
    "- **FP16+JIT** — `optimize_for_inference(dtype=torch.float16)` fuses layers with `torch.jit.script` and\n",
    "  halves the arithmetic precision.  Typically 1.5–2× faster than FP32 with negligible accuracy loss on\n",
    "  modern tensor cores.\n",
    "- **ONNX** — the model is exported to the Open Neural Network Exchange format and run through ONNX Runtime,\n",
    "  bypassing PyTorch entirely.  ORT applies its own graph optimisations and can use the TensorRT execution\n",
    "  provider for additional speedup.  Benchmarked separately on CPU and GPU to show the provider impact."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "533f037b",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "from rfdetr.export._onnx.inference import _onnx_runtime\n",
    "\n",
    "\n",
    "def _predict_fp32(model: Any, image: Image.Image) -> BenchmarkResult:\n",
    "    \"\"\"Baseline FP32 predict() latency.\"\"\"\n",
    "    mean, std = _measure(lambda: model.predict(image))\n",
    "    return BenchmarkResult(\"predict() FP32\", mean, std)\n",
    "\n",
    "\n",
    "def _predict_fp16(model: Any, image: Image.Image) -> BenchmarkResult:\n",
    "    \"\"\"FP16+JIT latency — applies and removes optimize_for_inference.\"\"\"\n",
    "    model.optimize_for_inference(dtype=torch.float16)\n",
    "    mean, std = _measure(lambda: model.predict(image))\n",
    "    model.remove_optimized_model()\n",
    "    return BenchmarkResult(\"predict() FP16+JIT\", mean, std)\n",
    "\n",
    "\n",
    "def _export_onnx(model: Any, export_dir: Path) -> Path:\n",
    "    \"\"\"Export model to ONNX and return the path.\"\"\"\n",
    "    return Path(model.export(output_dir=str(export_dir)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f6bc92a",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## 6. Model benchmark runner\n",
    "\n",
    "Each model is loaded fresh, exported once, and then each inference config is timed independently.\n",
    "The ONNX path reuses the same exported file for both CPU and CUDA providers so export cost is not\n",
    "counted in latency.  CUDA ONNX is skipped gracefully when no GPU is available; missing `onnxruntime-gpu`\n",
    "raises immediately so misconfigured environments are caught early."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8afd179a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_model_benchmark(\n",
    "    model_cls: type,\n",
    "    model_name: str,\n",
    "    images: list[Image.Image],\n",
    ") -> list[BenchmarkResult]:\n",
    "    \"\"\"Run FP32 / FP16 / ONNX benchmarks for one model and print results.\"\"\"\n",
    "    print(f\"\\n{'=' * 62}\")\n",
    "    print(f\"  {model_name}\")\n",
    "    print(\"=\" * 62)\n",
    "\n",
    "    model: Any = model_cls()\n",
    "    export_dir = EXPORT_DIR / model_name.split()[0]\n",
    "    export_dir.mkdir(exist_ok=True)\n",
    "    image = images[0]\n",
    "\n",
    "    fp32 = _predict_fp32(model, image)\n",
    "    fp16 = _predict_fp16(model, image)\n",
    "    results: list[BenchmarkResult] = [fp32, fp16]\n",
    "\n",
    "    onnx_path = _export_onnx(model, export_dir)\n",
    "    for providers in ([\"CPUExecutionProvider\"], [\"CUDAExecutionProvider\", \"CPUExecutionProvider\"]):\n",
    "        if providers[0] == \"CUDAExecutionProvider\" and not torch.cuda.is_available():\n",
    "            print(\"  ⚠ ONNX (CUDA) skipped — no CUDA GPU\")\n",
    "            continue\n",
    "        mean_ms, std_ms, label = _onnx_runtime(onnx_path, image, providers, WARMUP_RUNS, MEASURE_RUNS)\n",
    "        results.append(BenchmarkResult(f\"ONNX ({label})\", mean_ms, std_ms))\n",
    "\n",
    "    for r in results:\n",
    "        print(f\"  {r.label:<30}  {r.mean_ms:6.2f} ms ± {r.std_ms:5.2f}   ({r.fps:6.1f} FPS)\")\n",
    "\n",
    "    onnx_results = [r for r in results if r.label.startswith(\"ONNX\")]\n",
    "    speedups = [f\"FP16 {fp32.mean_ms / fp16.mean_ms:.1f}×\"]\n",
    "    if onnx_results:\n",
    "        speedups.append(f\"ONNX {fp32.mean_ms / onnx_results[0].mean_ms:.1f}×\")\n",
    "    print(f\"  Speedup vs FP32: {' | '.join(speedups)}\")\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7ccd41c",
   "metadata": {},
   "source": [
    "## 7. Benchmark loop — detection · segmentation · keypoint\n",
    "\n",
    "Three model families are benchmarked to show how task complexity affects latency.  Detection (`RFDETRMedium`)\n",
    "outputs boxes only; segmentation (`RFDETRSegSmall`) additionally predicts per-object masks, which adds\n",
    "decoder cost; keypoint (`RFDETRKeypointPreview`) predicts skeleton joints and is typically the lightest\n",
    "of the three at smaller resolutions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ab7ffbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rfdetr import RFDETRKeypointPreview, RFDETRMedium, RFDETRSegSmall\n",
    "\n",
    "MODELS: list[tuple[type, str]] = [\n",
    "    (RFDETRMedium, \"RFDETRMedium — detection\"),\n",
    "    (RFDETRSegSmall, \"RFDETRSegSmall — segmentation\"),\n",
    "    (RFDETRKeypointPreview, \"RFDETRKeypointPreview — keypoint\"),\n",
    "]\n",
    "\n",
    "all_results: dict[str, list[BenchmarkResult]] = {}\n",
    "for _model_cls, _model_name in MODELS:\n",
    "    all_results[_model_name] = run_model_benchmark(_model_cls, _model_name, images)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebbba970",
   "metadata": {},
   "source": [
    "## 8. Summary\n",
    "\n",
    "The table shows FPS (frames per second) for each model × config combination.  Higher is better.\n",
    "Compare columns to see which model fits your latency budget; compare rows to choose the right\n",
    "inference backend for your deployment target (Python server, edge CPU, or ONNX Runtime service)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d8eddd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "summary = {\n",
    "    model_name.split()[0]: {r.label: round(r.fps, 1) for r in results} for model_name, results in all_results.items()\n",
    "}\n",
    "df = pd.DataFrame(summary)\n",
    "df.index.name = \"Config \\\\ Model\"\n",
    "print(df.to_string())\n",
    "print(f\"\\nFPS — {MEASURE_RUNS} timed + {WARMUP_RUNS} warmup runs, batch 1, GPU: {torch.cuda.get_device_name(0)}.\")"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "title,-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}