{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Deploying DINOv3 as an INT8 TensorRT engine\n\nThe self-supervised foundation model [DINOv3](https://github.com/facebookresearch/dinov3) developed by Meta, is a\nuniversal vision backbone that a growing share of state-of-the-art models are\nnow built on top of. DINOv3's output is not boxes or masks, the features\nthemselves are the product.\n\n\n<img src=\"https://storage.googleapis.com/embedl-deploy-artifacts/pca_video.gif\" alt=\"Animated clip of a cockatoo next to its DINOv3 patch features,\">\n         foreground-segmented and colored by their top-3 PCA components \u2014 the\n         FP32 and Embedl INT8 columns track the moving subject in lockstep, one\n         block per 16\u00d716-px patch.\n   :width: 100%\n\n|\n\nIn this tutorial we step-by-step deploy DINOv3 with TensorRT using\n``embedl-deploy[tensorrt]`` on NVIDIA GPUs.\n\nPipeline:\n\n1. Load DINOv3 and fuse + INT8-quantize it with ``embedl_deploy``,\n   calibrating on imagenette (auto-downloaded).\n\n2. Export FP32 and INT8 (QDQ) ONNX graphs and build one TensorRT engine\n   from each.\n\n3. Verify the deployment: patch-similarity maps, PCA feature maps,\n   latency, and k-NN accuracy vs the FP32 reference.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>The engine build and everything after it require an NVIDIA GPU with\n    TensorRT 10.x (10.16 pip wheel recommended); tested on CUDA 12.8.\n    The first run downloads imagenette (~1.5 GB).</p></div>\n\n|\n\nInstall dependencies:\n\n```bash\npip install torch torchvision transformers onnx onnxscript matplotlib pillow\npip install \"tensorrt-cu12==10.16.*\"\npip install \"embedl-deploy[tensorrt]\"\n```\nThe DINOv3 checkpoint is gated: request access on its [Hugging Face model page](https://huggingface.co/facebook/dinov3-vitb16-pretrain-lvd1689m) (approval\nfrom Meta is usually quick) and authenticate before running:\n\n```bash\nhf auth login\n```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 0 \u2014 Configuration\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import sys\nimport time\nimport urllib.request\nfrom pathlib import Path\nfrom typing import Any\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\nfrom torch.utils.data import DataLoader\nfrom torchvision import datasets, transforms\n\nsys.setrecursionlimit(5000)\n\nARTIFACTS_PATH = Path(\"dinov3_artifacts\")\nARTIFACTS_PATH.mkdir(parents=True, exist_ok=True)\n\nMODEL_ID = \"facebook/dinov3-vitb16-pretrain-lvd1689m\"\nIMAGE_SIZE = 224\nPATCH_SIZE = 16\nPATCH_PREFIX = 5  # 1 CLS + 4 register tokens precede the patch tokens\nGRID = IMAGE_SIZE // PATCH_SIZE  # 14\u00d714 = 196 patch tokens\nN_CALIB = 512  # frames used for PTQ calibration\n\n# DINOv3 uses standard ImageNet normalization.\nIMAGENET_MEAN = [0.485, 0.456, 0.406]\nIMAGENET_STD = [0.229, 0.224, 0.225]\n\nDEMO_URL = \"http://images.cocodataset.org/val2017/000000039769.jpg\"\nDEMO_IMAGE = ARTIFACTS_PATH / \"demo_coco_cats.jpg\"\n\nONNX_FP32 = ARTIFACTS_PATH / \"dinov3-vitb16_original.onnx\"\nONNX_INT8 = ARTIFACTS_PATH / \"dinov3-vitb16_int8.onnx\"\nENGINE_FP16 = ARTIFACTS_PATH / \"dinov3.fp16.engine\"\nENGINE_INT8 = ARTIFACTS_PATH / \"dinov3.int8.engine\"\nTIMING_CACHE = ARTIFACTS_PATH / \"trt_timing.cache\"\nBUILDER_OPTIMIZATION_LEVEL = 3\n\nDEVICE = torch.device(\"cuda\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 1 \u2014 Load DINOv3\n\nA thin wrapper gives ``torch.export`` a positional tensor signature:\n``image [B, 3, 224, 224] \u2192 tokens [B, 201, 768]`` (1 CLS + 4 register +\n196 patch tokens). Eager attention keeps the exported graph free of\ndata-dependent SDPA branches. The untouched FP32 model is exported to\nONNX right away \u2014 it becomes the FP16 baseline engine.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from transformers import (  # type: ignore[import-not-found]\n    AutoConfig,\n    AutoModel,\n)\n\n\nclass Dinov3Encoder(nn.Module):\n    \"\"\"Wrap the backbone to expose an image-to-tokens forward signature.\"\"\"\n\n    def __init__(self, backbone: nn.Module) -> None:\n        \"\"\"Store the DINOv3 backbone.\"\"\"\n        super().__init__()\n        self.backbone = backbone\n\n    def forward(self, image: torch.Tensor) -> torch.Tensor:\n        \"\"\"Return the full token sequence.\"\"\"\n        return self.backbone(pixel_values=image).last_hidden_state\n\n\ndef load_dinov3() -> Dinov3Encoder:\n    \"\"\"Load a fresh, export-friendly DINOv3 model in fp32 eval mode.\"\"\"\n    cfg = AutoConfig.from_pretrained(MODEL_ID)\n    cfg.torch_dtype = torch.float32\n    if hasattr(cfg, \"_attn_implementation\"):\n        cfg._attn_implementation = \"eager\"  # noqa: SLF001  # SDPA branches break export\n    if hasattr(cfg, \"use_cache\"):\n        cfg.use_cache = False\n    backbone = AutoModel.from_pretrained(\n        MODEL_ID, config=cfg, torch_dtype=torch.float32\n    )\n    return Dinov3Encoder(backbone).eval().to(DEVICE)\n\n\nmodel = load_dinov3()\n\n\ndef export_onnx(module: nn.Module, path: Path) -> None:\n    \"\"\"Export a module to ONNX via the dynamo path.\"\"\"\n    dummy = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE, device=DEVICE)\n    with torch.no_grad():\n        torch.onnx.export(\n            module,\n            (dummy,),\n            str(path),\n            input_names=[\"image\"],\n            output_names=[\"last_hidden_state\"],\n            do_constant_folding=True,\n            dynamo=True,\n        )\n    print(f\"  {path.name}: {path.stat().st_size / 1e6:.0f} MB\")\n    del dummy\n    torch.cuda.empty_cache()\n\n\nexport_onnx(model, ONNX_FP32)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 2 \u2014 Calibration data\n\nPost-training quantization (PTQ) needs real images to calibrate activation\nranges. Imagenette (a 10-class ImageNet subset) auto-downloads via\ntorchvision; its train split provides calibration images (and later the\nk-NN index), the val split is the held-out test set.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "EVAL_TF = transforms.Compose(\n    [\n        transforms.Resize(256),\n        transforms.CenterCrop(IMAGE_SIZE),\n        transforms.ToTensor(),\n        transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),\n    ]\n)\n\n\ndef imagenette(split: str) -> datasets.Imagenette:\n    \"\"\"Load imagenette, downloading it on the first run.\"\"\"\n    try:\n        return datasets.Imagenette(\n            \".\", split=split, transform=EVAL_TF, download=True\n        )\n    except RuntimeError:\n        return datasets.Imagenette(\".\", split=split, transform=EVAL_TF)\n\n\ntrain_ds, val_ds = imagenette(\"train\"), imagenette(\"val\")\nloader_kw: dict[str, Any] = {\n    \"batch_size\": 64,\n    \"num_workers\": 4,\n    \"pin_memory\": True,\n}\ntrain_loader = DataLoader(train_ds, shuffle=False, **loader_kw)\nval_loader = DataLoader(val_ds, shuffle=False, **loader_kw)\n# The split is class-sorted \u2014 a shuffled loader (fixed seed) keeps a\n# partial calibration set class-balanced and reproducible.\ncalib_loader = DataLoader(\n    train_ds,\n    shuffle=True,\n    generator=torch.Generator().manual_seed(42),\n    **loader_kw,\n)\nprint(f\"  {len(train_ds)} train / {len(val_ds)} val images\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 3 \u2014 Fuse and INT8-quantize with Embedl Deploy\n\n``transform`` exports the model and applies the TensorRT fusion patterns;\n``quantize`` inserts INT8 quantizers and calibrates them. Three recipe\nchoices matter for ViT-family encoders:\n\n- **HISTOGRAM calibration for activations** \u2014 MINMAX lets a single\n  outlier in a ViT MLP blow up a whole tensor's quantization range.\n- **Skip the patch-embed stem** \u2014 the first 3-channel conv stays in\n  higher precision; quantizing the model boundary costs the most.\n- **Skip SmoothQuant for LayerNorm** \u2014 its learnable affine should not\n  be folded into quant scales.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from embedl_deploy import transform\nfrom embedl_deploy.backend import set_backend\nfrom embedl_deploy.quantize import (\n    CalibrationMethod,\n    ModulesToSkip,\n    Precision,\n    QuantConfig,\n    TensorQuantConfig,\n    quantize,\n)\nfrom embedl_deploy.tensorrt import TENSORRT_PATTERNS\n\nset_backend(\"tensorrt\")\n\ndummy = torch.randn(1, 3, IMAGE_SIZE, IMAGE_SIZE, device=DEVICE)\nwith torch.no_grad():\n    result = transform(model, (dummy,), TENSORRT_PATTERNS)\nprint(\n    f\"Patterns applied: {result.report['applied_count']}, \"\n    f\"skipped: {result.report['skipped_count']}\"\n)\nfused = result.model.eval().to(device=DEVICE, dtype=torch.float32)\n\n\ndef patch_embed_conv(module: nn.Module) -> nn.Conv2d | None:\n    \"\"\"Return the stem conv (3-channel input, patch-sized kernel) to skip.\"\"\"\n    for mod in module.modules():\n        if (\n            isinstance(mod, nn.Conv2d)\n            and mod.in_channels == 3\n            and max(mod.kernel_size) >= 7\n        ):\n            return mod\n    return None\n\n\nskip_mods: set[nn.Module] = set()\nif (stem := patch_embed_conv(fused)) is not None:\n    skip_mods.add(stem)\n\nconfig = QuantConfig(\n    activation=TensorQuantConfig(\n        precision=Precision.INT8,\n        symmetric=True,\n        per_channel=False,\n        calibration_method=CalibrationMethod.HISTOGRAM,\n    ),\n    weight=TensorQuantConfig(\n        precision=Precision.INT8,\n        symmetric=True,\n        per_channel=True,\n    ),\n    skip=ModulesToSkip(\n        stub=skip_mods,  # type: ignore[arg-type]\n        weight=skip_mods,  # type: ignore[arg-type]\n        smooth={nn.LayerNorm},\n    ),\n)\n\n\ndef calib_loop(m: nn.Module) -> None:\n    \"\"\"Feed calibration images one by one (the graph is batch-1).\"\"\"\n    m.eval()\n    seen = 0\n    with torch.no_grad():\n        for images, _ in calib_loader:\n            for i in range(images.shape[0]):\n                if seen >= N_CALIB:\n                    return\n                m(images[i : i + 1].to(DEVICE))\n                seen += 1\n\n\nprint(f\"Calibrating (HISTOGRAM) on {N_CALIB} images\")\nint8_model = quantize(\n    fused,\n    args=(dummy,),\n    config=config,\n    forward_loop=calib_loop,\n    freeze_weights=True,\n)\n\nexport_onnx(int8_model, ONNX_INT8)\n\ndel model, result, fused, int8_model, dummy\ntorch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 4 \u2014 Build the TensorRT engines\n\nOne FP16 baseline engine from the FP32 ONNX, one INT8+FP16 engine from\nthe QDQ ONNX. A shared timing cache makes the second build (and any\nrebuild) much faster.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import tensorrt as trt\n\n\ndef build_engine(onnx_path: Path, engine_path: Path, *, int8: bool) -> None:\n    \"\"\"Build a TRT engine. Always FP16; INT8 added if requested.\"\"\"\n    logger = trt.Logger(trt.Logger.WARNING)\n    trt.init_libnvinfer_plugins(logger, \"\")\n    builder = trt.Builder(logger)\n    network = builder.create_network(\n        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)\n    )\n    parser = trt.OnnxParser(network, logger)\n    if not parser.parse(onnx_path.read_bytes(), path=str(onnx_path)):\n        for i in range(parser.num_errors):\n            print(parser.get_error(i))\n        raise RuntimeError(\"ONNX parse failed\")\n\n    config = builder.create_builder_config()\n    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 * 1024**3)\n    config.builder_optimization_level = BUILDER_OPTIMIZATION_LEVEL\n    config.set_flag(trt.BuilderFlag.FP16)\n    if int8:\n        config.set_flag(trt.BuilderFlag.INT8)\n    cache = config.create_timing_cache(\n        TIMING_CACHE.read_bytes() if TIMING_CACHE.exists() else b\"\"\n    )\n    config.set_timing_cache(cache, ignore_mismatch=False)\n\n    tag = \"INT8+FP16\" if int8 else \"FP16\"\n    print(f\"Building {tag} engine from {onnx_path.name}\")\n    t0 = time.perf_counter()\n    plan = builder.build_serialized_network(network, config)\n    if plan is None:\n        raise RuntimeError(\"build_serialized_network returned None\")\n\n    TIMING_CACHE.write_bytes(bytes(config.get_timing_cache().serialize()))\n    engine_path.write_bytes(bytes(plan))\n    print(\n        f\"  built in {time.perf_counter() - t0:.0f} s \u2192 {engine_path.name} \"\n        f\"({len(bytes(plan)) / 1e6:.0f} MB)\"\n    )\n\n\nbuild_engine(ONNX_FP32, ENGINE_FP16, int8=False)\nbuild_engine(ONNX_INT8, ENGINE_INT8, int8=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Running the engines\n\nA small helper used for each step below. The engine works on fixed\nGPU memory locations, so its input and output tensors are allocated\nonce \u2014 each ``infer()`` just copies an image in, runs the engine, and\ncopies the result out.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "_TRT_TORCH = {\n    trt.float32: torch.float32,\n    trt.float16: torch.float16,\n    trt.int32: torch.int32,\n    trt.int64: torch.int64,\n    trt.int8: torch.int8,\n    trt.bool: torch.bool,\n}\n\n\nclass TrtRunner:\n    \"\"\"TRT engine wrapper backed by torch.cuda tensors as I/O buffers.\"\"\"\n\n    def __init__(self, engine_path: Path) -> None:\n        \"\"\"Deserialize the engine and bind I/O buffers.\"\"\"\n        logger = trt.Logger(trt.Logger.WARNING)\n        trt.init_libnvinfer_plugins(logger, \"\")\n        self.engine = trt.Runtime(logger).deserialize_cuda_engine(\n            engine_path.read_bytes()\n        )\n        self.ctx = self.engine.create_execution_context()\n        self.stream = torch.cuda.Stream()\n        self.bufs: dict[str, torch.Tensor] = {}\n        self.outs: list[str] = []\n        for i in range(self.engine.num_io_tensors):\n            name = self.engine.get_tensor_name(i)\n            shape = tuple(self.engine.get_tensor_shape(name))\n            dtype = _TRT_TORCH[self.engine.get_tensor_dtype(name)]\n            self.bufs[name] = torch.empty(shape, dtype=dtype, device=\"cuda\")\n            self.ctx.set_tensor_address(name, self.bufs[name].data_ptr())\n            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT:\n                self.outs.append(name)\n\n    def run(self) -> None:\n        \"\"\"Execute the engine on whatever is in the input buffers.\"\"\"\n        self.ctx.execute_async_v3(self.stream.cuda_stream)\n        self.stream.synchronize()\n\n    def infer(self, image: torch.Tensor) -> torch.Tensor:\n        \"\"\"Copy in \u2192 execute \u2192 return the token tensor on CPU, fp32.\"\"\"\n        buf = self.bufs[\"image\"]\n        buf.copy_(image.to(device=buf.device, dtype=buf.dtype))\n        self.run()\n        return self.bufs[self.outs[0]].detach().float().cpu()\n\n\nfp16_runner = TrtRunner(ENGINE_FP16)\nint8_runner = TrtRunner(ENGINE_INT8)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 5 \u2014 Feature inspection\n\nTwo views of the patch features, FP16 engine vs INT8 engine:\n\n- **Patch similarity** \u2014 pick an anchor patch (the \u2715), color every\n  patch by cosine similarity to it. Semantically matching regions light\n  up: anchoring on one cat highlights both cats.\n- **PCA features** \u2014 project the 768-dim patch features onto their top\n  3 principal components and render them as RGB.\n\n<img src=\"https://huggingface.co/datasets/embedl/documentation-images/resolve/main/dinov3-quantized-tensorrt/patch_similarity_cats.png\" alt=\"Anchor-patch cosine-similarity maps for the COCO cats image \u2014\">\n         TensorRT FP16 and Embedl INT8 produce near-identical heatmaps.\n   :width: 100%\n\n|\n\nPCA features rendered as RGB \u2014 matching colors mean the quantized\nmodel preserves the feature geometry:\n\n<img src=\"https://storage.googleapis.com/embedl-deploy-artifacts/pca_features.png\" alt=\"Cats-image patch features projected onto their top-3 PCA\">\n         components and rendered as RGB blocks; the DINOv3 FP32 and\n         Embedl INT8 columns are nearly identical.\n   :width: 100%\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt  # type: ignore[import-not-found]\nfrom PIL import Image  # type: ignore[import-not-found]\n\nANCHOR_YX = (5, 4)  # left cat's head, in the 14\u00d714 patch grid\n\n\ndef demo_image() -> Image.Image:\n    \"\"\"Download (once) and open the COCO cats demo image.\"\"\"\n    if not DEMO_IMAGE.exists():\n        urllib.request.urlretrieve(DEMO_URL, DEMO_IMAGE)\n    return Image.open(DEMO_IMAGE).convert(\"RGB\")\n\n\ndef patch_grid(tokens: torch.Tensor) -> torch.Tensor:\n    \"\"\"Reshape ``[1, 201, 768]`` tokens into a ``[14, 14, 768]`` patch grid.\"\"\"\n    return tokens[0, PATCH_PREFIX:].reshape(GRID, GRID, -1)\n\n\ndef similarity_map(patches: torch.Tensor, yx: tuple[int, int]) -> torch.Tensor:\n    \"\"\"Compute cosine similarity of every patch to the anchor patch.\"\"\"\n    anchor = patches[yx[0], yx[1]]\n    flat = patches.reshape(-1, patches.shape[-1])\n    return F.cosine_similarity(flat, anchor[None], dim=-1).reshape(GRID, GRID)\n\n\ndef pca_rgb(\n    patches: torch.Tensor, basis: torch.Tensor | None = None\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Project the patch features onto their top 3 principal components.\"\"\"\n    flat = patches.reshape(-1, patches.shape[-1])\n    flat = flat - flat.mean(dim=0)\n    if basis is None:\n        _, _, basis = torch.pca_lowrank(flat, q=3)\n    rgb = (flat @ basis).reshape(GRID, GRID, 3)\n    lo = rgb.amin(dim=(0, 1))\n    hi = rgb.amax(dim=(0, 1))\n    return (rgb - lo) / (hi - lo).clamp_min(1e-8), basis\n\n\nimage = demo_image()\ncrop = transforms.CenterCrop(IMAGE_SIZE)(transforms.Resize(256)(image))\nx = EVAL_TF(image).unsqueeze(0)\n\nfp16_patches = patch_grid(fp16_runner.infer(x))\nint8_patches = patch_grid(int8_runner.infer(x))\n\n\nfig, axes = plt.subplots(2, 3, figsize=(10.5, 7))\nfor row in axes:\n    row[0].imshow(crop)\nay, ax_ = ANCHOR_YX\naxes[0, 0].plot(\n    (ax_ + 0.5) * PATCH_SIZE, (ay + 0.5) * PATCH_SIZE, \"rx\", ms=12, mew=3\n)\n# Shared color scale so the two heatmaps are directly comparable.\nsim_kw = {\"cmap\": \"viridis\", \"vmin\": -0.2, \"vmax\": 1.0}\naxes[0, 1].imshow(similarity_map(fp16_patches, ANCHOR_YX), **sim_kw)\naxes[0, 2].imshow(similarity_map(int8_patches, ANCHOR_YX), **sim_kw)\nrgb_fp16, basis = pca_rgb(fp16_patches)\nrgb_int8, _ = pca_rgb(int8_patches, basis)\naxes[1, 1].imshow(rgb_fp16)\naxes[1, 2].imshow(rgb_int8)\naxes[0, 0].set_title(\"Input (\u2715 = anchor)\")\nfor col, name in [(1, \"TensorRT FP16\"), (2, \"Embedl INT8\")]:\n    axes[0, col].set_title(f\"{name} \u2014 similarity\")\n    axes[1, col].set_title(f\"{name} \u2014 PCA\")\nfor a in axes.ravel():\n    a.axis(\"off\")\nfig.tight_layout()\nfig.savefig(ARTIFACTS_PATH / \"feature_maps.png\", dpi=120, bbox_inches=\"tight\")\nprint(f\"  \u2192 {ARTIFACTS_PATH / 'feature_maps.png'}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Step 6 \u2014 Latency and k-NN accuracy\n\nBenchmark both engines: forward-pass latency via CUDA events, and\nk-NN classification accuracy on imagenette against the FP32 reference.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "N_WARMUP, N_ITERS = 20, 200\n\n\ndef bench(runner: TrtRunner) -> tuple[float, float]:\n    \"\"\"Compute mean and p95 engine latency in ms over ``N_ITERS`` runs.\"\"\"\n    for _ in range(N_WARMUP):\n        runner.run()\n    start = torch.cuda.Event(enable_timing=True)\n    end = torch.cuda.Event(enable_timing=True)\n    samples = []\n    for _ in range(N_ITERS):\n        start.record()\n        runner.run()\n        end.record()\n        torch.cuda.synchronize()\n        samples.append(start.elapsed_time(end))\n    samples.sort()\n    return sum(samples) / len(samples), samples[int(0.95 * len(samples)) - 1]\n\n\nfp16_mean, fp16_p95 = bench(fp16_runner)\nint8_mean, int8_p95 = bench(int8_runner)\nprint(\"Latency (DINOv3 ViT-B/16 @ 224x224, batch 1)\")\nprint(f\"  TensorRT FP16:  mean {fp16_mean:5.2f} ms  p95 {fp16_p95:5.2f} ms\")\nprint(\n    f\"  Embedl INT8:    mean {int8_mean:5.2f} ms  p95 {int8_p95:5.2f} ms  \"\n    f\"({fp16_mean / int8_mean:.2f}x vs FP16)\"\n)\n\nKNN_K, KNN_TAU = 20, 0.07\n\n\n@torch.no_grad()\ndef cls_features(\n    model_fn: nn.Module, loader: DataLoader[Any]\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Return L2-normalized CLS features (``[N, D]``) and labels (``[N]``).\"\"\"\n    model_fn.eval()\n    feats, labels = [], []\n    for images, targets in loader:\n        hidden = model_fn(images.to(DEVICE))\n        feats.append(F.normalize(hidden[:, 0, :].float(), dim=-1).cpu())\n        labels.append(targets)\n    return torch.cat(feats), torch.cat(labels)\n\n\n@torch.no_grad()\ndef cls_features_trt(\n    runner: TrtRunner, ds: datasets.Imagenette\n) -> torch.Tensor:\n    \"\"\"Return CLS features from the batch-1 engine for each image in `ds`.\"\"\"\n    feats = []\n    for image_t, _ in ds:\n        tokens = runner.infer(image_t.unsqueeze(0))\n        feats.append(F.normalize(tokens[:, 0, :], dim=-1))\n    return torch.cat(feats)\n\n\ndef knn_top1(\n    index_feats: torch.Tensor,\n    index_labels: torch.Tensor,\n    query_feats: torch.Tensor,\n    query_labels: torch.Tensor,\n    num_classes: int,\n) -> float:\n    \"\"\"Compute k-NN top-1 accuracy (%) via weighted exp(sim/tau) votes.\"\"\"\n    index_gpu = index_feats.to(DEVICE)\n    correct = 0\n    for start in range(0, len(query_feats), 256):\n        q = query_feats[start : start + 256].to(DEVICE)\n        top_sim, top_idx = (q @ index_gpu.T).topk(KNN_K, dim=1)\n        votes = torch.zeros(len(q), num_classes, device=DEVICE)\n        votes.scatter_add_(\n            1,\n            index_labels[top_idx.cpu()].to(DEVICE),\n            (top_sim / KNN_TAU).exp(),\n        )\n        pred = votes.argmax(dim=1).cpu()\n        correct += int(\n            (pred == query_labels[start : start + 256]).sum().item()\n        )\n    return 100.0 * correct / len(query_feats)\n\n\nprint(\"FP32 reference: extracting index (train) + query (val) features\")\nfp32 = load_dinov3()\nindex_feats, index_labels = cls_features(fp32, train_loader)\nval_feats, val_labels = cls_features(fp32, val_loader)\ndel fp32\ntorch.cuda.empty_cache()\n\nn_classes = len(train_ds.classes)\nfp32_acc = knn_top1(\n    index_feats, index_labels, val_feats, val_labels, n_classes\n)\nprint(f\"  FP32 knn_top1: {fp32_acc:.2f}%\")\n\nprint(\"INT8 engine: extracting query features\")\nint8_feats = cls_features_trt(int8_runner, val_ds)\nint8_acc = knn_top1(\n    index_feats, index_labels, int8_feats, val_labels, n_classes\n)\nprint(\n    f\"  INT8 knn_top1: {int8_acc:.2f}%  ({int8_acc - fp32_acc:+.2f} vs FP32)\"\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "On an NVIDIA L4 with TensorRT 10.16:\n\n```text\nLatency (DINOv3 ViT-B/16 @ 224x224, batch 1)\n  TensorRT FP16:  mean  1.50 ms  p95  1.52 ms\n  Embedl INT8:    mean  1.13 ms  p95  1.15 ms  (1.34x vs FP16)\nFP32 reference: extracting index (train) + query (val) features\n  FP32 knn_top1: 99.75%\nINT8 engine: extracting query features\n  INT8 knn_top1: 99.26%  (-0.48 vs FP32)\n```\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Summary\n\n::\n\n    HuggingFace \u2500\u25b6 Embedl Deploy fuse + INT8 PTQ \u2500\u25b6 ONNX (QDQ)\n                \u2500\u25b6 TensorRT engine (INT8+FP16)\n\nDINOv3 ViT-B/16 @ 224\u00d7224, batch 1, TensorRT 10.16:\n\n==============  ===========  ================  =====================\nPath            NVIDIA L4    Jetson AGX Orin   knn_top1 (imagenette)\n==============  ===========  ================  =====================\nTensorRT FP16   1.50 ms      2.71 ms           99.75 % (FP32 ref)\nEmbedl INT8     1.13 ms      2.25 ms           99.26 %\n==============  ===========  ================  =====================\n\nThe quantized model, engines and full evaluation protocol are published\nat [embedl/dinov3-quantized-tensorrt](https://huggingface.co/embedl/dinov3-quantized-tensorrt). For the\nsame recipe on a detection model, see the SAM3 tutorial.\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.15"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}