.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_tutorials/thor_fp8_vit.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_tutorials_thor_fp8_vit.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_tutorials_thor_fp8_vit.py:


FP8 ViT on NVIDIA Jetson AGX Thor
=================================

Thor's Blackwell GPU runs FP8 GEMMs, so a Vision Transformer whose attention
and MLP matmuls are quantized to FP8 beats the FP16 model on latency without
touching the rest of the network. This tutorial takes the HuggingFace
``google/vit-base-patch16-224`` classifier through the whole path:

1. ``transform()`` fuses it with the TensorRT patterns under the
   ``tensorrt_fp8`` backend.
2. ``quantize()`` inserts FP8 Q/DQ stubs and calibrates them on ImageNette.
3. ``convert_float_to_float16()`` turns everything that is not FP8 into
   float16 *in PyTorch*, so ``torch.onnx.export`` writes the final graph.
4. ``trtexec --stronglyTyped`` builds an engine that keeps every tensor at the
   precision the ONNX says, and the engine is profiled and scored on the
   device.

The FP16 baseline goes through exactly the same steps minus ``quantize()``, so
the two engines differ only in the Q/DQ nodes. There is no ONNX-level rewriting
anywhere: what ``torch.onnx.export`` produces is what TensorRT builds.

The script runs end to end wherever ``trtexec`` is found, so the simplest way
to use it is to run the whole file on the Thor: the PyTorch half (steps 1--3)
takes a few minutes there and writes two ONNX files under ``artifacts/``, and
step 4 builds, profiles and scores them in place. On a workstation without
TensorRT the script stops after the export.

Measured on a Jetson AGX Thor with TensorRT 10.13.3 (batch 1, GPU compute time
with CUDA graphs, ImageNette validation, 3925 images):

.. code-block:: text

    engine                          median ms    top-1
    fp16, weakly typed --fp16           1.399    86.0 %
    fp16, strongly typed                1.402    86.0 %
    fp8 + fp16, strongly typed          1.065    86.2 %

FP8 is 1.32x faster than the FP16 engine at the same accuracy.

Install dependencies:

.. code-block:: bash

    pip install "embedl-deploy[tensorrt]" torchvision transformers onnx onnxscript

On the Thor itself, the PyPI ``torch`` wheel runs on the GPU as is (``torch
2.14, cu130``), and the ``tensorrt`` Python module comes with JetPack -- create
the virtualenv with ``--system-site-packages`` to see it.

.. GENERATED FROM PYTHON SOURCE LINES 56-64

Setup: the model and ImageNette
--------------------------------

The HuggingFace classifier takes ``pixel_values`` and returns a
``logits``-carrying output; a thin wrapper gives it the plain
``tensor -> tensor`` signature every step below expects. Attention is set
to the ``eager`` implementation so the fuser sees the individual matmuls
and softmax rather than a fused SDPA call.

.. GENERATED FROM PYTHON SOURCE LINES 64-165

.. code-block:: Python


    import copy
    import itertools
    import json
    import re
    import shutil
    import subprocess
    import tarfile
    import urllib.request
    from pathlib import Path

    import torch
    import torchvision
    from torch import fx, nn
    from torchvision import transforms
    from transformers import (  # type: ignore[import-not-found]
        AutoModelForImageClassification,
    )

    MODEL_ID = "google/vit-base-patch16-224"
    IMAGENETTE_URL = (
        "https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz"
    )
    DATA_DIR = Path("artifacts/data")
    IMAGENETTE_DIR = DATA_DIR / "imagenette2-320"
    OUT_DIR = Path("artifacts/thor_fp8_vit")
    OUT_DIR.mkdir(parents=True, exist_ok=True)

    #: ImageNette's ten classes, as ImageNet-1k indices.
    IMAGENETTE_TO_IMAGENET = [0, 217, 482, 491, 497, 566, 569, 571, 574, 701]
    CALIBRATION_IMAGES = 128
    BATCH_SIZE = 1

    torch.manual_seed(0)


    class ImageClassifier(nn.Module):
        """``pixel_values -> logits`` view of a HuggingFace image classifier."""

        def __init__(self, model_id: str) -> None:
            """Load the pretrained classifier `model_id` with eager attention."""
            super().__init__()
            self.model = AutoModelForImageClassification.from_pretrained(
                model_id, attn_implementation="eager"
            )

        def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
            """Return the ImageNet-1k logits for `pixel_values`."""
            logits: torch.Tensor = self.model(pixel_values=pixel_values).logits
            return logits


    def download_imagenette() -> None:
        """Download and extract ImageNette if not already present."""
        if IMAGENETTE_DIR.exists():
            return
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        tgz_path = DATA_DIR / "imagenette2-320.tgz"
        print(f"Downloading ImageNette to {tgz_path} ...")
        urllib.request.urlretrieve(IMAGENETTE_URL, str(tgz_path))
        with tarfile.open(tgz_path) as tar:
            tar.extractall(DATA_DIR, filter="data")
        tgz_path.unlink()


    download_imagenette()

    # ViT-B/16 was trained with 0.5 mean/std and a straight resize to 224.
    preprocess = transforms.Compose(
        [
            transforms.Resize(224),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ]
    )


    def remap(label: int) -> int:
        """Map an ImageNette class index to its ImageNet class index."""
        return IMAGENETTE_TO_IMAGENET[label]


    train_set = torchvision.datasets.ImageFolder(
        str(IMAGENETTE_DIR / "train"), transform=preprocess, target_transform=remap
    )
    val_set = torchvision.datasets.ImageFolder(
        str(IMAGENETTE_DIR / "val"), transform=preprocess, target_transform=remap
    )
    val_loader = torch.utils.data.DataLoader(val_set, batch_size=BATCH_SIZE)

    calibration_loader = torch.utils.data.DataLoader(
        train_set, batch_size=BATCH_SIZE, shuffle=True
    )
    calibration_batches = [
        images
        for images, _ in itertools.islice(calibration_loader, CALIBRATION_IMAGES)
    ]

    example_input = torch.randn(1, 3, 224, 224)


.. GENERATED FROM PYTHON SOURCE LINES 166-175

Step 1: fuse for TensorRT with the FP8 policy
----------------------------------------------

``tensorrt_fp8`` is the TensorRT backend with the FP8 precision policy
measured on Thor: the same conversions and fusions as ``tensorrt``, plus
the knowledge of which kernels take FP8 profitably and which stay in
floating point. ``transform()`` rewrites the model into fused modules
around those kernels -- the attention projections, the MLP linears -- and
leaves everything else as it is.

.. GENERATED FROM PYTHON SOURCE LINES 175-188

.. code-block:: Python


    from embedl_deploy import transform
    from embedl_deploy.backend import set_backend

    set_backend("tensorrt_fp8")

    model = ImageClassifier(MODEL_ID).eval()
    fused = transform(model, (example_input,)).model

    with torch.no_grad():
        drift = (model(example_input) - fused(example_input)).abs().max().item()
    print(f"max |logit drift| after fusion: {drift:.2e}")


.. GENERATED FROM PYTHON SOURCE LINES 189-197

Step 2: quantize to FP8
------------------------

Activations and weights both get E4M3 FP8 Q/DQ. Calibration runs the 128
training images through the model to pick each stub's scale. Layers the
FP8 policy rules out -- the three-channel patch-embedding stem, for one --
are left in floating point on its own. ``freeze_weights=True`` bakes the
weight scales in, which the export in step 3 requires.

.. GENERATED FROM PYTHON SOURCE LINES 197-241

.. code-block:: Python


    from embedl_deploy.quantize import (
        Precision,
        QuantConfig,
        TensorQuantConfig,
        quantize,
    )


    def forward_loop(m: nn.Module) -> None:
        """Feed the calibration images through `m`."""
        with torch.no_grad():
            for images in calibration_batches:
                m(images)


    fp8 = TensorQuantConfig(Precision.FP8)
    # ``quantize()`` works in place, so the baseline keeps its own copy.
    quantized = quantize(
        copy.deepcopy(fused),
        (example_input,),
        config=QuantConfig(activation=fp8, weight=fp8),
        forward_loop=forward_loop,
        freeze_weights=True,
    )


    def accuracy(m: nn.Module, limit: int | None = None) -> float:
        """Top-1 accuracy of `m` on the validation split (or its first `limit`)."""
        correct = total = 0
        with torch.no_grad():
            for images, labels in val_loader:
                correct += int((m(images).argmax(dim=1) == labels).sum())
                total += labels.numel()
                if limit is not None and total >= limit:
                    break
        return correct / total


    # A quick sanity check in PyTorch: the FP8 fake-quantized model on the first
    # 256 validation images. The full-split numbers come from the engines.
    print(f"FP32 top-1 (256 images):     {accuracy(fused, 256):.1%}")
    print(f"FP8 QDQ top-1 (256 images):  {accuracy(quantized, 256):.1%}")


.. GENERATED FROM PYTHON SOURCE LINES 242-253

Step 3: float16 in PyTorch, then export
----------------------------------------

``convert_float_to_float16()`` returns a copy whose floating parameters and
buffers are float16, keeps the quantizer scales usable, and pins the few
operations that need float32 (interpolation, cumulative sums, top-k) with
casts at their boundaries. The Q/DQ nodes are emitted during export with
float16 scales, so the exported graph is float16 end to end with FP8 in
the quantized matmuls. The same half model is also saved as a ``.pt2``
through ``torch.export`` for inspection or a later re-export. The baseline
is the same conversion applied to the fused, unquantized model.

.. GENERATED FROM PYTHON SOURCE LINES 253-281

.. code-block:: Python


    from embedl_deploy.quantize import convert_float_to_float16


    def export_onnx(m: fx.GraphModule, path: Path) -> Path:
        """Export the float16 copy of `m` to ONNX, and to ``.pt2`` alongside."""
        half = convert_float_to_float16(m)
        with torch.no_grad():
            exported = torch.export.export(half, (example_input.half(),))
            torch.export.save(exported, str(path.with_suffix(".pt2")))
            torch.onnx.export(
                half,
                (example_input.half(),),
                str(path),
                dynamo=True,
                opset_version=20,
                input_names=["input"],
                output_names=["logits"],
            )
        # The exporter keeps the weights in a ``<name>.data`` file next to the ONNX.
        size = sum(p.stat().st_size for p in path.parent.glob(f"{path.name}*"))
        print(f"wrote {path} ({size / 1e6:.0f} MB)")
        return path


    fp16_onnx = export_onnx(fused, OUT_DIR / "vit_fp16.onnx")
    fp8_onnx = export_onnx(quantized, OUT_DIR / "vit_fp8.onnx")


.. GENERATED FROM PYTHON SOURCE LINES 282-291

Step 4: build, profile and score on the Thor
---------------------------------------------

``--stronglyTyped`` tells TensorRT to honour the dtypes in the ONNX exactly:
float16 tensors stay float16, the ``QuantizeLinear`` outputs stay FP8, and
nothing is re-decided by the autotuner. That is what makes the comparison
clean -- both engines run precisely the graph the export wrote. The FP16
engine is also built once weakly typed (``--fp16``), the way most people
run a float model, for reference.

.. GENERATED FROM PYTHON SOURCE LINES 291-351

.. code-block:: Python


    TRTEXEC = shutil.which("trtexec") or "/usr/src/tensorrt/bin/trtexec"
    if not Path(TRTEXEC).exists():
        print(
            f"trtexec not found -- copy {OUT_DIR} to the Thor and run this "
            "script there to build, profile and score the engines."
        )
        raise SystemExit(0)

    _COMPUTE_TIME = re.compile(
        r"GPU Compute Time: min = ([\d.]+) ms, max = ([\d.]+) ms, "
        r"mean = ([\d.]+) ms, median = ([\d.]+) ms"
    )


    def build_engine(onnx_path: Path, label: str, *flags: str) -> dict[str, float]:
        """Build `onnx_path` with `flags`, time it, return the latency summary."""
        engine = OUT_DIR / f"{label}.engine"
        log = OUT_DIR / f"trtexec_{label}.log"
        cmd = [
            TRTEXEC,
            f"--onnx={onnx_path}",
            *flags,
            f"--saveEngine={engine}",
            f"--timingCacheFile={OUT_DIR / 'timing.cache'}",
            "--useCudaGraph",
            "--noDataTransfers",
            "--useSpinWait",
            "--warmUp=1000",
            "--iterations=500",
            "--avgRuns=100",
        ]
        print(f"[{label}] {' '.join(cmd)}")
        with log.open("w") as fh:
            subprocess.run(cmd, stdout=fh, stderr=subprocess.STDOUT, check=True)
        match = _COMPUTE_TIME.search(log.read_text())
        if match is None:
            raise RuntimeError(f"no GPU compute time summary in {log}")
        minimum, maximum, mean, median = (float(v) for v in match.groups())
        print(f"[{label}] median {median:.3f} ms (mean {mean:.3f} ms)")
        return {
            "median_ms": median,
            "mean_ms": mean,
            "min_ms": minimum,
            "max_ms": maximum,
        }


    latency = {
        "fp16 (weakly typed --fp16)": build_engine(
            fp16_onnx, "fp16_weak", "--fp16"
        ),
        "fp16 (strongly typed)": build_engine(
            fp16_onnx, "fp16_strong", "--stronglyTyped"
        ),
        "fp8 + fp16 (strongly typed)": build_engine(
            fp8_onnx, "fp8_strong", "--stronglyTyped"
        ),
    }


.. GENERATED FROM PYTHON SOURCE LINES 352-355

Accuracy is measured by running the engines over the whole ImageNette
validation split (3925 images) through the TensorRT Python API. The inputs
are float16 because the strongly typed engines' input tensor is.

.. GENERATED FROM PYTHON SOURCE LINES 355-412

.. code-block:: Python


    import tensorrt as trt


    class Engine:
        """Minimal single-input, single-output TensorRT engine runner."""

        def __init__(self, path: Path) -> None:
            """Deserialize the engine at `path` and allocate its output buffer."""
            runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
            self.engine = runtime.deserialize_cuda_engine(path.read_bytes())
            self.context = self.engine.create_execution_context()
            self.input_name = self.engine.get_tensor_name(0)
            self.output_name = self.engine.get_tensor_name(1)
            self.stream = torch.cuda.Stream()
            self.output = torch.empty(
                tuple(self.engine.get_tensor_shape(self.output_name)),
                dtype=_torch_dtype(self.engine.get_tensor_dtype(self.output_name)),
                device="cuda",
            )
            self.input_dtype = _torch_dtype(
                self.engine.get_tensor_dtype(self.input_name)
            )

        def __call__(self, images: torch.Tensor) -> torch.Tensor:
            """Run one batch through the engine and return float32 logits."""
            images = images.to("cuda", self.input_dtype).contiguous()
            self.context.set_tensor_address(self.input_name, images.data_ptr())
            self.context.set_tensor_address(
                self.output_name, self.output.data_ptr()
            )
            self.context.execute_async_v3(self.stream.cuda_stream)
            self.stream.synchronize()
            return self.output.float().cpu()


    def _torch_dtype(dtype: "trt.DataType") -> torch.dtype:
        """Map a TensorRT tensor dtype to the torch dtype of its buffer."""
        return {trt.float32: torch.float32, trt.float16: torch.float16}[dtype]


    def engine_accuracy(label: str) -> float:
        """Top-1 accuracy of the engine saved under `label`."""
        engine = Engine(OUT_DIR / f"{label}.engine")
        correct = total = 0
        for images, labels in val_loader:
            correct += int((engine(images).argmax(dim=1) == labels).sum())
            total += labels.numel()
        return correct / total


    top1 = {
        "fp16 (weakly typed --fp16)": engine_accuracy("fp16_weak"),
        "fp16 (strongly typed)": engine_accuracy("fp16_strong"),
        "fp8 + fp16 (strongly typed)": engine_accuracy("fp8_strong"),
    }


.. GENERATED FROM PYTHON SOURCE LINES 413-415

Results
-------

.. GENERATED FROM PYTHON SOURCE LINES 415-435

.. code-block:: Python


    print(f"\n{'engine':<32}{'median ms':>12}{'top-1':>10}")
    for label, stats in latency.items():
        print(f"{label:<32}{stats['median_ms']:>12.3f}{top1[label]:>10.1%}")
    speedup = (
        latency["fp16 (strongly typed)"]["median_ms"]
        / latency["fp8 + fp16 (strongly typed)"]["median_ms"]
    )
    print(f"\nFP8 speed-up over the FP16 engine: {speedup:.2f}x")

    results = {
        "model": MODEL_ID,
        "tensorrt": trt.__version__,
        "latency": latency,
        "top1": top1,
        "latency_is": "GPU compute time, batch 1, CUDA graphs, H2D/D2H excluded",
        "accuracy_is": "ImageNette validation, 3925 images, ImageNet-1k labels",
    }
    (OUT_DIR / "results.json").write_text(json.dumps(results, indent=2))
    print(f"wrote {OUT_DIR / 'results.json'}")


.. _sphx_glr_download_auto_tutorials_thor_fp8_vit.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: thor_fp8_vit.ipynb <thor_fp8_vit.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: thor_fp8_vit.py <thor_fp8_vit.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: thor_fp8_vit.zip <thor_fp8_vit.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_