.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_tutorials/mixed_precision_search.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_tutorials_mixed_precision_search.py: Mixed precision search ====================== Post-training quantization is not all-or-nothing: leaving a few sensitive layers at full precision often recovers most of the accuracy loss for a small cost in latency or size. ``embedl_deploy``'s search subsystem automates that trade-off -- it drives an `Optuna `_ study over *which* activations and weights to skip, scoring each candidate with a user-supplied evaluation function. Two strategies ship today: - ``GranularQuantConfigSearchSpace`` -- one independent boolean per quantizer axis, the finest granularity. - ``PartitionedQuantConfigSearchSpace`` -- a handful of continuous parameters carve the model into alternating skip/keep sections, so the search stays constant-size regardless of model size. This tutorial takes an ImageNet-pretrained torchvision ResNet-18, fuses it with the TensorRT patterns, and searches for which quantizers to skip -- trading ImageNette validation accuracy against how much of the model stays quantized. Install dependencies: .. code-block:: bash pip install "embedl-deploy[tensorrt,search]" torchvision .. GENERATED FROM PYTHON SOURCE LINES 35-43 Setup: a pretrained model and ImageNette ----------------------------------------- We use an ImageNet-pretrained ResNet-18 as-is -- no fine-tuning needed. For data we use `ImageNette `_, a 10-class subset of ImageNet, as in the other tutorials. Its labels are remapped to the corresponding ImageNet class indices so the pretrained classifier head can be scored directly. .. GENERATED FROM PYTHON SOURCE LINES 43-117 .. code-block:: Python import itertools import tarfile import urllib.request from pathlib import Path import torch import torchvision from torchvision import transforms from torchvision.models import resnet18 from embedl_deploy.backend import set_backend set_backend("tensorrt") torch.manual_seed(0) IMAGENETTE_URL = ( "https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz" ) DATA_DIR = Path("artifacts/data") IMAGENETTE_DIR = DATA_DIR / "imagenette2-320" IMAGENETTE_TO_IMAGENET = [0, 217, 482, 491, 497, 566, 569, 571, 574, 701] BATCH_SIZE = 32 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.") download_imagenette() transform_images = transforms.Compose( [ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)), ] ) 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=transform_images, target_transform=remap, ) val_set = torchvision.datasets.ImageFolder( str(IMAGENETTE_DIR / "val"), transform=transform_images, target_transform=remap, ) train_loader = torch.utils.data.DataLoader( train_set, batch_size=BATCH_SIZE, shuffle=True ) model = resnet18(weights="IMAGENET1K_V1").eval() .. GENERATED FROM PYTHON SOURCE LINES 118-126 Held-out evaluation and calibration data ----------------------------------------- A fixed slice of the validation set scores every candidate on the same images, and a few training batches calibrate the quant stubs. Small, fixed subsets keep each search trial cheap. ``ImageFolder`` yields its samples in class order, so the slice is drawn through a seeded shuffle -- taken straight off the loader it would be one class of the ten. .. GENERATED FROM PYTHON SOURCE LINES 126-152 .. code-block:: Python eval_loader = torch.utils.data.DataLoader( val_set, batch_size=BATCH_SIZE, shuffle=True, generator=torch.Generator().manual_seed(0), ) eval_batches = list(itertools.islice(iter(eval_loader), 8)) calib_batches = [ images for images, _ in itertools.islice(iter(train_loader), 4) ] def accuracy(model: torch.nn.Module) -> float: """Top-1 accuracy of `model` on the fixed evaluation batches.""" correct = total = 0 with torch.no_grad(): for images, labels in eval_batches: correct += int((model(images).argmax(dim=1) == labels).sum()) total += labels.numel() return correct / total fp32_accuracy = accuracy(model) print(f"FP32 accuracy: {fp32_accuracy:.1%}") .. GENERATED FROM PYTHON SOURCE LINES 153-157 Fusing for quantization ------------------------ The search operates on a fused, not-yet-quantized model. .. GENERATED FROM PYTHON SOURCE LINES 157-165 .. code-block:: Python from embedl_deploy import transform from embedl_deploy.tensorrt import TENSORRT_PATTERNS args = (torch.randn(1, 3, 224, 224),) fused = transform(model, args, patterns=TENSORRT_PATTERNS).model print(f"Fused accuracy: {accuracy(fused):.1%} (fusion is lossless)") .. GENERATED FROM PYTHON SOURCE LINES 166-174 Scoring a configuration ------------------------ ``eval_fn`` receives a freshly quantized model and returns any number of named metrics. Here, ``accuracy`` is the real test-subset accuracy, and ``active_quantizers`` counts how many quant stubs are still enabled -- fewer skips means more of the model runs in INT8 but risks more accuracy loss, so the two trade off against each other as the search progresses. .. GENERATED FROM PYTHON SOURCE LINES 174-201 .. code-block:: Python from embedl_deploy.quantize import QuantStub, WeightFakeQuantize from embedl_deploy.search import Direction def forward_loop(model: torch.nn.Module) -> None: """Feed the calibration batches through `model` to calibrate its stubs.""" with torch.no_grad(): for images in calib_batches: model(images) def eval_fn(model: torch.nn.Module) -> dict[str, float]: """Score `model` on accuracy and on how many quantizers stay active.""" active = sum( module.enabled for module in model.modules() if isinstance(module, (QuantStub, WeightFakeQuantize)) ) return {"accuracy": accuracy(model), "active_quantizers": active} directions: dict[str, Direction] = { "accuracy": "maximize", "active_quantizers": "maximize", } .. GENERATED FROM PYTHON SOURCE LINES 202-210 The search strategies ---------------------- Both strategies discover the same underlying dimensions -- one per enabled weight or activation quantizer -- from the fused model alone, before any search-space decision is made. Each dimension is keyed by the module and axis it controls, and those keys double as the Optuna parameter names, so trial parameters stay readable in model terms everywhere downstream. .. GENERATED FROM PYTHON SOURCE LINES 210-218 .. code-block:: Python from embedl_deploy.search import get_dimensions dimensions = get_dimensions(fused) print(f"{len(dimensions)} quantizable dimensions, e.g.:") for key in list(dimensions)[:4]: print(f" {key}") .. GENERATED FROM PYTHON SOURCE LINES 219-228 Running a granular search -------------------------- ``GranularQuantConfigSearchSpace`` turns every dimension into its own boolean parameter. :meth:`~embedl_deploy.search.QuantConfigSearchSpace.search` drives Optuna's ask/evaluate/tell loop, seeding two informed starting points ahead of the sampled ones: skip nothing (the fully-INT8 model) and skip everything (effectively the FP32 baseline). .. GENERATED FROM PYTHON SOURCE LINES 228-247 .. code-block:: Python from embedl_deploy.search import GranularQuantConfigSearchSpace granular = GranularQuantConfigSearchSpace( fused, args, forward_loop=forward_loop, eval_fn=eval_fn ) results, study = granular.search(16, directions=directions) int8_result = results[0] if int8_result is None: raise RuntimeError("The seeded fully-INT8 trial failed.") int8_accuracy = int8_result["accuracy"] print(f"FP32 baseline: {fp32_accuracy:.1%}") print(f"Fully INT8: {int8_accuracy:.1%}") print("Pareto-optimal trade-offs found:") for trial in sorted(study.best_trials, key=lambda t: t.values[1]): active = int(trial.values[1]) print(f" {active:2d} active quantizers -> {trial.values[0]:.1%}") .. GENERATED FROM PYTHON SOURCE LINES 248-257 Inspecting a trial ------------------- Because parameters are named after the modules they control, a trial's ``params`` read directly in model terms. Each trial additionally records the concrete manifestation of its point -- the qualified names it skips, per axis -- on its ``skip`` user attribute, so Optuna tooling (``study.trials_dataframe()``, dashboards, plot hover text) shows *what* a trial did, not just its raw parameters. .. GENERATED FROM PYTHON SOURCE LINES 257-264 .. code-block:: Python trial = study.trials[2] skipped = [key for key, keep in trial.params.items() if not keep] print(f"trial 2 skips {len(skipped)} axes, e.g. {skipped[:3]}") print(f"manifestation: {trial.user_attrs['skip']}") print(f"as a config: {granular.skip(trial)!r}") .. GENERATED FROM PYTHON SOURCE LINES 265-275 Constant-size search spaces ----------------------------- For a model with hundreds of quantizable dimensions, one boolean each is a lot of parameters. ``PartitionedQuantConfigSearchSpace`` instead carves the same dimensions into a handful of alternating skip/keep sections, controlled by only ``2 * num_enabled_sections`` continuous parameters -- constant-size regardless of model size. Its raw parameters are abstract cut positions, which is exactly where the per-trial ``skip`` user attribute keeps the study interpretable. .. GENERATED FROM PYTHON SOURCE LINES 275-283 .. code-block:: Python from embedl_deploy.search import PartitionedQuantConfigSearchSpace partitioned = PartitionedQuantConfigSearchSpace( fused, args, forward_loop=forward_loop, eval_fn=eval_fn ) print(partitioned) .. GENERATED FROM PYTHON SOURCE LINES 284-293 Crash-resilient persistence ------------------------------ Passing ``path`` persists every trial via `~optuna.storages.JournalStorage` -- a crashed run resumes exactly where it left off instead of re-running completed trials. The journal lives under ``artifacts/`` next to the data, so re-running this script (or raising ``n_trials`` later) picks the study up where it stopped rather than starting over. .. GENERATED FROM PYTHON SOURCE LINES 293-311 .. code-block:: Python STUDY_PATH = Path("artifacts/search/partitioned_study.log") STUDY_PATH.parent.mkdir(parents=True, exist_ok=True) first, _ = partitioned.search( 4, directions=directions, path=STUDY_PATH, name="partitioned" ) second, _ = partitioned.search( len(first) + 4, directions=directions, path=STUDY_PATH, name="partitioned" ) if second[: len(first)] != first: raise RuntimeError("Resuming re-ran trials that `path` already held.") print( f"resumed run reused {len(first)} prior trials, " f"added {len(second) - len(first)} more" ) .. GENERATED FROM PYTHON SOURCE LINES 312-323 Rebuilding the best model -------------------------- With two objectives there is no single best trial, only a Pareto front (``study.best_trials``). Pick the one you want by your own rule -- here the most-quantized trial that stays within one point of the FP32 accuracy -- and hand it back to the space. ``config`` returns the full ``QuantConfig`` that trial was scored with, and ``apply`` re-quantizes a fresh copy of the fused model with it, ready for export. Both accept a ``FrozenTrial`` from a study reloaded via ``search(0, path=..., name=...)`` in a later session, so the search never has to be re-run to get its model back. .. GENERATED FROM PYTHON SOURCE LINES 323-334 .. code-block:: Python within_budget = [ t for t in study.best_trials if t.values[0] >= fp32_accuracy - 0.01 ] best = max(within_budget or study.best_trials, key=lambda t: t.values[1]) print(f"best trial {best.number}: {best.user_attrs['result']}") print(f"its config: {granular.config(best)}") best_model = granular.apply(best) print(f"rebuilt accuracy: {accuracy(best_model):.1%}") .. GENERATED FROM PYTHON SOURCE LINES 335-342 Visualizing the trade-off ---------------------------- If Plotly is installed then Optuna's own visualization utilities work unchanged, since ``search`` just drives a normal ``optuna.Study``. Hovering a point shows the trial's ``skip`` user attribute -- the modules that configuration leaves at full precision. .. GENERATED FROM PYTHON SOURCE LINES 342-352 .. code-block:: Python import optuna.visualization fig = optuna.visualization.plot_pareto_front( study, target_names=["active_quantizers", "accuracy"], targets=lambda t: (t.values[1], t.values[0]), ) fig.write_html("pareto_front.html") print("Pareto front written to pareto_front.html") .. _sphx_glr_download_auto_tutorials_mixed_precision_search.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: mixed_precision_search.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: mixed_precision_search.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: mixed_precision_search.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_