Note

Go to the end to download the full example code.

Deploying vision models with Embedl Deploy#

In this tutorial, we demonstrate an end-to-end workflow for deploying vision models using embedl_deploy for TensorRT. The models selected are from the torchvision library, covering a range of architectures: ConvNeXt Tiny, ResNet-50, and ViT-B/16. We apply post-training static quantization (PTQ) using the built-in TensorRT pattern set. Embedl Deploy automatically handles special cases such as depthwise convolutions, which are memory-bound and left in FP16 to avoid quantization overhead.

The pipeline after transformation and quantization for deploying a model consists of the following steps:

  1. Export the PyTorch model to ONNX and simplify with onnxsim.

  2. Build a TensorRT engine (FP16 for baseline, --best for QDQ models).

  3. Run TensorRT inference.

  4. Measure latency with the TensorRT Python API.

Note

This tutorial requires an NVIDIA GPU with TensorRT 10.x installed.

Constants#

Pattern selection#

When no explicit pattern list is passed, transform() uses the active backend’s conversion and fusion patterns automatically. These include structural conversions (e.g. decomposing MultiheadAttention), operator fusions (Conv-BN-ReLU, Linear-ReLU, LayerNorm, etc.), and automatic handling of depthwise convolutions. Depthwise convolutions are memory-bound, so quantizing them adds TensorRT reformatting overhead that exceeds the compute gain from INT8 — Embedl Deploy detects these and keeps them in FP16 automatically.

Dataset helpers#

We use ImageNette — a 10-class subset of ImageNet — for fast evaluation. Data loaders come from tensorrt_inference.

import tarfile
import urllib.request
from pathlib import Path

import tensorrt as trt
import torch
import torchvision
from torch import nn

from embedl_deploy.backend import set_backend

try:
    from tensorrt_inference import (
        CALIBRATION_BATCHES,
        DATA_DIR,
        IMAGENETTE_DIR,
        IMAGENETTE_URL,
        ImageLoader,
        build_trt_engine,
        evaluate_trt,
        make_loaders,
        measure_latency,
    )
except ImportError:
    raise ImportError(
        "This tutorial requires the companion file 'tensorrt_inference.py' "
        "which provides utilities for downloading datasets, compiling models "
        "and measuring latencies. Please download it from the same directory "
        "as this tutorial."
    ) from None

set_backend("tensorrt")


def download_imagenette() -> None:
    """Download and extract ImageNette if not already present."""
    if IMAGENETTE_DIR.exists():
        print(f"ImageNette already present at {IMAGENETTE_DIR}")
        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))
    print("Extracting ...")
    with tarfile.open(tgz_path) as tar:
        tar.extractall(DATA_DIR)
    tgz_path.unlink()
    print("Done.")

ONNX export#

Export the model to ONNX and simplify with onnxsim.

import onnx
import onnxsim


def export_and_simplify(
    model: nn.Module, onnx_path: Path, crop_size: int = 224
) -> Path:
    """Export to ONNX + simplify with onnxsim."""
    model = model.cpu().eval()
    x = torch.randn(1, 3, crop_size, crop_size)
    torch.onnx.export(
        model,
        (x,),
        str(onnx_path),
        opset_version=20,
        input_names=["input"],
        output_names=["output"],
        dynamo=False,
    )
    onnx_model = onnx.load(str(onnx_path))
    simplified, ok = onnxsim.simplify(onnx_model)
    simp_path = onnx_path.with_name(onnx_path.stem + "_sim.onnx")
    if ok:
        onnx.save(simplified, str(simp_path))
        print(f"  Simplified ONNX: {simp_path}")
    else:
        print("  onnxsim failed; using raw export.")
        simp_path = onnx_path
    return simp_path

Quantization with Embedl Deploy#

Fuse layers, insert QDQ stubs with the selective pattern list, and calibrate on training data.

from embedl_deploy import transform
from embedl_deploy.quantize import (
    Precision,
    QuantConfig,
    QuantStub,
    TensorQuantConfig,
    WeightFakeQuantize,
    quantize,
)


def quantize_embedl(
    pretrained_model: nn.Module,
    calib_batches: list[torch.Tensor],
    crop_size: int = 224,
) -> nn.Module:
    """Fuse + quantize + calibrate using Embedl Deploy."""
    print("\n=== Embedl Deploy PTQ ===")

    example_input = torch.randn(1, 3, crop_size, crop_size)
    fused_model = transform(
        pretrained_model.cpu().eval(), (example_input,)
    ).model

    # Verify lossless fusion.
    with torch.no_grad():
        x = torch.randn(1, 3, crop_size, crop_size)
        max_diff = (pretrained_model.cpu()(x) - fused_model(x)).abs().max()
    assert max_diff < 1e-4, f"Fusion diverged: {max_diff:.2e}"
    print(f"  Fusion check passed (max diff = {max_diff:.2e})")

    print(f"  Calibrating ({len(calib_batches)} batches) ...")

    def forward_loop(model: nn.Module) -> None:
        with torch.no_grad():
            for batch in calib_batches:
                model(batch)

    quantized = quantize(
        fused_model,
        (torch.randn(1, 3, crop_size, crop_size),),
        config=QuantConfig(
            activation=TensorQuantConfig(Precision.INT8, symmetric=True),
            weight=TensorQuantConfig(
                Precision.INT8,
                symmetric=True,
                per_channel=True,
            ),
        ),
        forward_loop=forward_loop,
        freeze_weights=True,
    )

    n_act = sum(1 for m in quantized.modules() if isinstance(m, QuantStub))
    n_wt = sum(
        1 for m in quantized.modules() if isinstance(m, WeightFakeQuantize)
    )
    print(f"  QuantStubs: {n_act}, WeightFakeQuantize: {n_wt}")
    print("  Calibration complete.")
    return quantized

Benchmark runner#

Helper that runs the full export -> build -> evaluate -> latency pipeline for a single model variant.

def run_variant(
    tag: str,
    model: nn.Module,
    *,
    dest: Path,
    val_loader: ImageLoader | None,
    crop_size: int,
    best: bool = False,
    fp16: bool = True,
) -> tuple[dict[str, float], dict[str, float]]:
    """Export, build, evaluate, and measure one model variant."""
    print(f"\n{'=' * 60}\n{tag}\n{'=' * 60}")
    onnx_path = export_and_simplify(model, dest / f"{tag}.onnx", crop_size)
    engine_path = dest / f"{tag}.engine"
    build_trt_engine(onnx_path, engine_path, fp16=fp16, best=best)

    acc = (
        evaluate_trt(engine_path, val_loader)
        if val_loader is not None
        else {"top1": 0.0, "top5": 0.0}
    )
    latency = measure_latency(engine_path)

    if val_loader is not None:
        print(f"  Top-1: {acc['top1']:.2f}%  Top-5: {acc['top5']:.2f}%")
    print(
        f"  Latency: {latency['mean_latency_ms']:.3f} ms  "
        f"Throughput: {latency['throughput_qps']:.1f} qps"
    )
    return acc, latency

Configuration#

Choose the model and image sizes. Change MODEL_NAME to benchmark a different architecture (e.g. "convnext_base", "resnet50", "vit_b_16").

MODEL_NAME = "vit_b_16"
CROP_SIZE = 224
RESIZE_SIZE = 256

BENCHMARK_DIR = Path(f"artifacts/ptq/{MODEL_NAME}_benchmark")
BENCHMARK_DIR.mkdir(parents=True, exist_ok=True)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
assert device.type == "cuda", "This benchmark requires a CUDA GPU."

print(f"Model: {MODEL_NAME}  crop={CROP_SIZE}  resize={RESIZE_SIZE}")
print(f"Output: {BENCHMARK_DIR}")

Prepare data and model#

Download ImageNette and load the pretrained torchvision model.

download_imagenette()
train_dl, val_dl = make_loaders(CROP_SIZE, RESIZE_SIZE)

calib_data: list[torch.Tensor] = []
for batch_idx, (imgs, _) in enumerate(train_dl):
    if batch_idx >= CALIBRATION_BATCHES:
        break
    calib_data.append(imgs)
print(f"Collected {len(calib_data)} calibration batches.")

pretrained = torchvision.models.get_model(MODEL_NAME, weights="DEFAULT").eval()
print(
    f"Loaded {MODEL_NAME} "
    f"({sum(p.numel() for p in pretrained.parameters()):,} params)"
)

Baseline FP16#

Build and benchmark the unquantized model with FP16 precision.

baseline_acc, baseline_lat = run_variant(
    "baseline_fp16",
    pretrained,
    dest=BENCHMARK_DIR,
    val_loader=val_dl,
    crop_size=CROP_SIZE,
)

Embedl Deploy mixed-precision#

Apply the active backend’s patterns. Depthwise convolutions are automatically kept in FP16 while compute-bound operators are quantized to INT8.

embedl_model = quantize_embedl(pretrained, calib_data, CROP_SIZE)
embedl_acc, embedl_lat = run_variant(
    "embedl_mixed_precision",
    embedl_model,
    dest=BENCHMARK_DIR,
    val_loader=val_dl,
    crop_size=CROP_SIZE,
    best=True,
)

Summary#

Compare accuracy and latency across all variants.

HEADER = (
    f"{'Variant':<25s} {'Top-1':>7s} {'Top-5':>7s} "
    f"{'Latency(ms)':>12s} {'Throughput':>12s}"
)
rows = [
    ("Baseline (FP16)", baseline_acc, baseline_lat),
    ("Embedl Deploy (best)", embedl_acc, embedl_lat),
]

print(f"\n{'=' * 80}")
print(f"BENCHMARK SUMMARY - {MODEL_NAME} PTQ on ImageNette")
print("=" * 80)
print(HEADER)
print("-" * len(HEADER))
for label, row_acc, row_lat in rows:
    print(
        f"{label:<25s} {row_acc['top1']:6.2f}% {row_acc['top5']:6.2f}% "
        f"{row_lat['mean_latency_ms']:11.3f} "
        f"{row_lat['throughput_qps']:10.1f} qps"
    )

speedup_e = baseline_lat["mean_latency_ms"] / max(
    embedl_lat["mean_latency_ms"], 1e-6
)
drop_e = baseline_acc["top1"] - embedl_acc["top1"]

print()
print(
    f"  Embedl Deploy - Top-1 drop: {drop_e:+.2f}pp, speedup: {speedup_e:.2f}x"
)
print("=" * 80)

# Save results.
RESULTS_PATH = BENCHMARK_DIR / f"{MODEL_NAME}_results.txt"
lines = [
    f"TRT {trt.__version__}",
    "=" * 80,
    f"BENCHMARK SUMMARY - {MODEL_NAME} PTQ on ImageNette",
    "=" * 80,
    HEADER,
    "-" * len(HEADER),
]
for label, row_acc, row_lat in rows:
    lines.append(
        f"{label:<25s} {row_acc['top1']:6.2f}% {row_acc['top5']:6.2f}% "
        f"{row_lat['mean_latency_ms']:11.3f} "
        f"{row_lat['throughput_qps']:10.1f} qps"
    )
lines += [
    "",
    f"  Embedl Deploy - Top-1 drop: {drop_e:+.2f}pp, "
    f"speedup: {speedup_e:.2f}x",
    "=" * 80,
]
RESULTS_PATH.write_text("\n".join(lines) + "\n")
print(f"\nResults saved to {RESULTS_PATH}")

Gallery generated by Sphinx-Gallery