{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "\n# Mixed precision search\n\nPost-training quantization is not all-or-nothing: leaving a few sensitive\nlayers at full precision often recovers most of the accuracy loss for a small\ncost in latency or size. ``embedl_deploy``'s search subsystem automates that\ntrade-off -- it drives an [Optuna](https://optuna.org/) study over *which*\nactivations and weights to skip, scoring each candidate with a user-supplied\nevaluation function.\n\nTwo strategies ship today:\n\n- ``GranularQuantConfigSearchSpace`` -- one independent boolean per quantizer\n axis, the finest granularity.\n- ``PartitionedQuantConfigSearchSpace`` -- a handful of continuous parameters\n carve the model into alternating skip/keep sections, so the search\n stays constant-size regardless of model size.\n\nThis tutorial takes an ImageNet-pretrained torchvision ResNet-18, fuses it with\nthe TensorRT patterns, and searches for which quantizers to skip -- trading\nImageNette validation accuracy against how much of the model stays quantized.\n\nInstall dependencies:\n\n```bash\npip install \"embedl-deploy[tensorrt,search]\" torchvision\n```\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup: a pretrained model and ImageNette\n\nWe use an ImageNet-pretrained ResNet-18 as-is -- no fine-tuning needed.\nFor data we use [ImageNette](https://github.com/fastai/imagenette), a\n10-class subset of ImageNet, as in the other tutorials. Its labels are\nremapped to the corresponding ImageNet class indices so the pretrained\nclassifier head can be scored directly.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import itertools\nimport tarfile\nimport urllib.request\nfrom pathlib import Path\n\nimport torch\nimport torchvision\nfrom torchvision import transforms\nfrom torchvision.models import resnet18\n\nfrom embedl_deploy.backend import set_backend\n\nset_backend(\"tensorrt\")\ntorch.manual_seed(0)\n\nIMAGENETTE_URL = (\n \"https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz\"\n)\nDATA_DIR = Path(\"artifacts/data\")\nIMAGENETTE_DIR = DATA_DIR / \"imagenette2-320\"\nIMAGENETTE_TO_IMAGENET = [0, 217, 482, 491, 497, 566, 569, 571, 574, 701]\nBATCH_SIZE = 32\n\n\ndef download_imagenette() -> None:\n \"\"\"Download and extract ImageNette if not already present.\"\"\"\n if IMAGENETTE_DIR.exists():\n print(f\"ImageNette already present at {IMAGENETTE_DIR}\")\n return\n DATA_DIR.mkdir(parents=True, exist_ok=True)\n tgz_path = DATA_DIR / \"imagenette2-320.tgz\"\n print(f\"Downloading ImageNette to {tgz_path} ...\")\n urllib.request.urlretrieve(IMAGENETTE_URL, str(tgz_path))\n print(\"Extracting ...\")\n with tarfile.open(tgz_path) as tar:\n tar.extractall(DATA_DIR)\n tgz_path.unlink()\n print(\"Done.\")\n\n\ndownload_imagenette()\n\ntransform_images = transforms.Compose(\n [\n transforms.Resize(256),\n transforms.CenterCrop(224),\n transforms.ToTensor(),\n transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),\n ]\n)\n\n\ndef remap(label: int) -> int:\n \"\"\"Map an ImageNette class index to its ImageNet class index.\"\"\"\n return IMAGENETTE_TO_IMAGENET[label]\n\n\ntrain_set = torchvision.datasets.ImageFolder(\n str(IMAGENETTE_DIR / \"train\"),\n transform=transform_images,\n target_transform=remap,\n)\nval_set = torchvision.datasets.ImageFolder(\n str(IMAGENETTE_DIR / \"val\"),\n transform=transform_images,\n target_transform=remap,\n)\ntrain_loader = torch.utils.data.DataLoader(\n train_set, batch_size=BATCH_SIZE, shuffle=True\n)\n\nmodel = resnet18(weights=\"IMAGENET1K_V1\").eval()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Held-out evaluation and calibration data\n\nA fixed slice of the validation set scores every candidate on the same\nimages, and a few training batches calibrate the quant stubs. Small,\nfixed subsets keep each search trial cheap. ``ImageFolder`` yields its\nsamples in class order, so the slice is drawn through a seeded shuffle --\ntaken straight off the loader it would be one class of the ten.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "eval_loader = torch.utils.data.DataLoader(\n val_set,\n batch_size=BATCH_SIZE,\n shuffle=True,\n generator=torch.Generator().manual_seed(0),\n)\neval_batches = list(itertools.islice(iter(eval_loader), 8))\ncalib_batches = [\n images for images, _ in itertools.islice(iter(train_loader), 4)\n]\n\n\ndef accuracy(model: torch.nn.Module) -> float:\n \"\"\"Top-1 accuracy of `model` on the fixed evaluation batches.\"\"\"\n correct = total = 0\n with torch.no_grad():\n for images, labels in eval_batches:\n correct += int((model(images).argmax(dim=1) == labels).sum())\n total += labels.numel()\n return correct / total\n\n\nfp32_accuracy = accuracy(model)\nprint(f\"FP32 accuracy: {fp32_accuracy:.1%}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Fusing for quantization\n\nThe search operates on a fused, not-yet-quantized model.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from embedl_deploy import transform\nfrom embedl_deploy.tensorrt import TENSORRT_PATTERNS\n\nargs = (torch.randn(1, 3, 224, 224),)\nfused = transform(model, args, patterns=TENSORRT_PATTERNS).model\nprint(f\"Fused accuracy: {accuracy(fused):.1%} (fusion is lossless)\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Scoring a configuration\n\n``eval_fn`` receives a freshly quantized model and returns any number of\nnamed metrics. Here, ``accuracy`` is the real test-subset accuracy, and\n``active_quantizers`` counts how many quant stubs are still enabled --\nfewer skips means more of the model runs in INT8 but risks more accuracy\nloss, so the two trade off against each other as the search progresses.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from embedl_deploy.quantize import QuantStub, WeightFakeQuantize\nfrom embedl_deploy.search import Direction\n\n\ndef forward_loop(model: torch.nn.Module) -> None:\n \"\"\"Feed the calibration batches through `model` to calibrate its stubs.\"\"\"\n with torch.no_grad():\n for images in calib_batches:\n model(images)\n\n\ndef eval_fn(model: torch.nn.Module) -> dict[str, float]:\n \"\"\"Score `model` on accuracy and on how many quantizers stay active.\"\"\"\n active = sum(\n module.enabled\n for module in model.modules()\n if isinstance(module, (QuantStub, WeightFakeQuantize))\n )\n return {\"accuracy\": accuracy(model), \"active_quantizers\": active}\n\n\ndirections: dict[str, Direction] = {\n \"accuracy\": \"maximize\",\n \"active_quantizers\": \"maximize\",\n}" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## The search strategies\n\nBoth strategies discover the same underlying dimensions -- one per enabled\nweight or activation quantizer -- from the fused model alone, before any\nsearch-space decision is made. Each dimension is keyed by the module and axis it\ncontrols, and those keys double as the Optuna parameter names, so trial\nparameters stay readable in model terms everywhere downstream.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from embedl_deploy.search import get_dimensions\n\ndimensions = get_dimensions(fused)\nprint(f\"{len(dimensions)} quantizable dimensions, e.g.:\")\nfor key in list(dimensions)[:4]:\n print(f\" {key}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Running a granular search\n\n``GranularQuantConfigSearchSpace`` turns every dimension into its own boolean\nparameter.\n:meth:`~embedl_deploy.search.QuantConfigSearchSpace.search`\ndrives Optuna's ask/evaluate/tell loop, seeding two informed starting\npoints ahead of the sampled ones: skip nothing (the fully-INT8 model) and\nskip everything (effectively the FP32 baseline).\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from embedl_deploy.search import GranularQuantConfigSearchSpace\n\ngranular = GranularQuantConfigSearchSpace(\n fused, args, forward_loop=forward_loop, eval_fn=eval_fn\n)\nresults, study = granular.search(16, directions=directions)\n\nint8_result = results[0]\nif int8_result is None:\n raise RuntimeError(\"The seeded fully-INT8 trial failed.\")\nint8_accuracy = int8_result[\"accuracy\"]\nprint(f\"FP32 baseline: {fp32_accuracy:.1%}\")\nprint(f\"Fully INT8: {int8_accuracy:.1%}\")\nprint(\"Pareto-optimal trade-offs found:\")\nfor trial in sorted(study.best_trials, key=lambda t: t.values[1]):\n active = int(trial.values[1])\n print(f\" {active:2d} active quantizers -> {trial.values[0]:.1%}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Inspecting a trial\n\nBecause parameters are named after the modules they control, a trial's\n``params`` read directly in model terms. Each trial additionally records\nthe concrete manifestation of its point -- the qualified names it skips,\nper axis -- on its ``skip`` user attribute, so Optuna tooling\n(``study.trials_dataframe()``, dashboards, plot hover text) shows *what*\na trial did, not just its raw parameters.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "trial = study.trials[2]\nskipped = [key for key, keep in trial.params.items() if not keep]\nprint(f\"trial 2 skips {len(skipped)} axes, e.g. {skipped[:3]}\")\nprint(f\"manifestation: {trial.user_attrs['skip']}\")\nprint(f\"as a config: {granular.skip(trial)!r}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Constant-size search spaces\n\nFor a model with hundreds of quantizable dimensions, one boolean each is\na lot of parameters. ``PartitionedQuantConfigSearchSpace`` instead carves the\nsame dimensions into a handful of alternating skip/keep sections,\ncontrolled by only ``2 * num_enabled_sections`` continuous parameters --\nconstant-size regardless of model size. Its raw parameters are abstract\ncut positions, which is exactly where the per-trial ``skip`` user\nattribute keeps the study interpretable.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from embedl_deploy.search import PartitionedQuantConfigSearchSpace\n\npartitioned = PartitionedQuantConfigSearchSpace(\n fused, args, forward_loop=forward_loop, eval_fn=eval_fn\n)\nprint(partitioned)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Crash-resilient persistence\n\nPassing ``path`` persists every trial via\n`~optuna.storages.JournalStorage` -- a crashed run resumes exactly where\nit left off instead of re-running completed trials. The journal lives\nunder ``artifacts/`` next to the data, so re-running this script (or\nraising ``n_trials`` later) picks the study up where it stopped rather\nthan starting over.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "STUDY_PATH = Path(\"artifacts/search/partitioned_study.log\")\nSTUDY_PATH.parent.mkdir(parents=True, exist_ok=True)\n\nfirst, _ = partitioned.search(\n 4, directions=directions, path=STUDY_PATH, name=\"partitioned\"\n)\nsecond, _ = partitioned.search(\n len(first) + 4, directions=directions, path=STUDY_PATH, name=\"partitioned\"\n)\n\nif second[: len(first)] != first:\n raise RuntimeError(\"Resuming re-ran trials that `path` already held.\")\nprint(\n f\"resumed run reused {len(first)} prior trials, \"\n f\"added {len(second) - len(first)} more\"\n)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Rebuilding the best model\n\nWith two objectives there is no single best trial, only a Pareto front\n(``study.best_trials``). Pick the one you want by your own rule -- here the\nmost-quantized trial that stays within one point of the FP32 accuracy --\nand hand it back to the space. ``config`` returns the full ``QuantConfig``\nthat trial was scored with, and ``apply`` re-quantizes a fresh copy of the\nfused model with it, ready for export. Both accept a ``FrozenTrial`` from a\nstudy reloaded via ``search(0, path=..., name=...)`` in a later session,\nso the search never has to be re-run to get its model back.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "within_budget = [\n t for t in study.best_trials if t.values[0] >= fp32_accuracy - 0.01\n]\nbest = max(within_budget or study.best_trials, key=lambda t: t.values[1])\nprint(f\"best trial {best.number}: {best.user_attrs['result']}\")\nprint(f\"its config: {granular.config(best)}\")\n\nbest_model = granular.apply(best)\nprint(f\"rebuilt accuracy: {accuracy(best_model):.1%}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Visualizing the trade-off\n\nIf Plotly is installed then Optuna's own visualization utilities work\nunchanged, since ``search`` just drives a normal ``optuna.Study``.\nHovering a point shows the trial's ``skip`` user attribute -- the modules\nthat configuration leaves at full precision.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import optuna.visualization\n\nfig = optuna.visualization.plot_pareto_front(\n study,\n target_names=[\"active_quantizers\", \"accuracy\"],\n targets=lambda t: (t.values[1], t.values[0]),\n)\nfig.write_html(\"pareto_front.html\")\nprint(\"Pareto front written to pareto_front.html\")" ] } ], "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 }