diff --git a/CHANGELOG.md b/CHANGELOG.md index c3c40e9..58264a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,11 @@ format follows [keep a changelog](https://keepachangelog.com/en/1.1.0/). ## unreleased ### added +- **ONNX export**: `export_to_onnx` / `earlyon export` write a portable static + multi-output graph (one output per exit plus the final classifier). Routing + stays at runtime — the graph computes every exit and the caller picks the + first confident one. Works for both conv and transformer exit heads; verified + against onnxruntime. (`onnxruntime` added to the `dev` extra for testing.) - **transformer support**: `vit_b_16_ee` wraps torchvision ViT-B/16 with two early exits (after encoder blocks 3 and 9). `EarlyExitHead` now accepts 3D token features ``(B, N, D)`` (CLS or mean pooling) and 2D vectors in addition diff --git a/CLAUDE.md b/CLAUDE.md index 37bd34c..fc7671c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -521,8 +521,9 @@ earlyon profile --model calibrated.pth earlyon analyze --model calibrated.pth --dataset cifar10 ``` -> ONNX export is not yet implemented (`earlyon/onnx.py` raises -> `NotImplementedError`); there is no `earlyon export` command. +> ONNX export writes a static multi-output graph (all exits; routing applied at +> runtime) via `earlyon export` / `earlyon.onnx.export_to_onnx`. The dynamic +> per-sample routing itself is not expressed in ONNX. --- diff --git a/README.md b/README.md index 5f66751..0cb0271 100644 --- a/README.md +++ b/README.md @@ -212,6 +212,7 @@ earlyon calibrate --model ee.pth --target-drop 0.01 --output calibrated.pth earlyon benchmark --model calibrated.pth --device cuda --runs 500 earlyon profile --model calibrated.pth --runs 200 # Jetson power + thermal earlyon analyze --model calibrated.pth # per-exit accuracy + distribution +earlyon export --model calibrated.pth --output model.onnx # static multi-exit ONNX graph ``` ## How it works @@ -233,8 +234,10 @@ earlyon analyze --model calibrated.pth # per-exit accuracy + d `forward_inference_batched(x)` for per-batch routing. - **No `torch.compile`** on the inference path — the conditional control flow is incompatible; the wrapper raises a clear error. Compile the raw backbone. -- **ONNX export not yet supported** — torch 2.x's exporter rejects the dynamic - control flow; deploy from PyTorch directly (tracked in `earlyon/onnx.py`). +- **ONNX export is static-graph only** — `earlyon export` writes a portable graph + that computes *all* exits (routing applied at runtime by the caller); the + per-sample early-exit speedup itself isn't expressed in ONNX. Use the PyTorch + wrapper when you need the actual compute saving. - **No compute-budget routing yet** — confidence and entropy ship today. ## Acknowledgements diff --git a/earlyon/cli.py b/earlyon/cli.py index 6add085..08162e7 100644 --- a/earlyon/cli.py +++ b/earlyon/cli.py @@ -288,5 +288,26 @@ def analyze(model_path: str, dataset: str, device: str) -> None: ) +@main.command() +@click.option("--model", "model_path", required=True, type=click.Path(exists=True)) +@click.option("--output", required=True, type=click.Path()) +@click.option("--input-size", default=224, type=int) +@click.option("--opset", default=17, type=int) +@click.option("--dynamic-batch/--no-dynamic-batch", default=True) +def export(model_path: str, output: str, input_size: int, opset: int, dynamic_batch: bool) -> None: + """Export all exits as a static multi-output ONNX graph (routing at runtime).""" + from earlyon.onnx import export_to_onnx + + model = load_wrapper(model_path) + names = export_to_onnx( + model, + output, + input_shape=(1, 3, input_size, input_size), + opset=opset, + dynamic_batch=dynamic_batch, + ) + click.echo(f"wrote {output} (input 1x3x{input_size}x{input_size}, outputs: {', '.join(names)})") + + if __name__ == "__main__": main() diff --git a/earlyon/onnx.py b/earlyon/onnx.py index a865668..7df6c27 100644 --- a/earlyon/onnx.py +++ b/earlyon/onnx.py @@ -1,27 +1,122 @@ -"""ONNX export for early-exit wrappers — not yet supported. +"""ONNX export for early-exit wrappers. -The honest story: ONNX has no clean answer for sample-dependent control -flow that satisfies all runtimes. We tried two approaches (per-exit static -graph shrapnel, and a TorchScript ``If`` op) and both ran into issues with -the new ``torch.export``-based exporter in torch 2.x. +ONNX has no clean, portable way to express sample-dependent control flow, so +earlyon does not try to bake the *routing* into the graph. Instead it exports a +single **static multi-output** graph that computes every exit: -A future approach is likely one of: -- ``torch.onnx.export`` with ``dynamo=False`` (legacy exporter) plus - TorchScript ``If`` to express the routing as a single ONNX graph -- TensorRT backend bypass: skip ONNX entirely and emit a TRT engine + inputs: input (N, C, H, W) + outputs: exit_0 .. exit_{k-1}, final each (N, num_classes) -Until then, deploy the wrapper from PyTorch directly. This stub exists only -so existing imports do not break. +The deploying application runs the graph and applies the same routing rule +earlyon uses at runtime — take the first exit whose softmax max ≥ threshold +(or whose entropy ≤ threshold) — picking the prediction and stopping. The graph +always computes all exits, so this trades the compute saving for a fully +portable, runtime-agnostic artifact; use the PyTorch wrapper directly when you +need the actual early-exit speedup. """ from __future__ import annotations -from typing import NoReturn +import inspect +import warnings +from pathlib import Path +from typing import Any +import torch +import torch.nn as nn -def export_to_onnx(*args: object, **kwargs: object) -> NoReturn: - raise NotImplementedError( - "ONNX export is not yet supported: the torch 2.x onnx exporter rejects " - "the dynamic control flow used by the early-exit wrapper. Deploy the " - "wrapper directly from PyTorch; track the github issue for status." - ) +from earlyon.core.wrappers import EarlyExitWrapper + + +class _AllExitsModule(nn.Module): + """Adapter exposing the wrapper's all-exits training forward as a plain + tuple-returning module, so the legacy ONNX tracer can capture a static graph. + """ + + def __init__(self, wrapper: EarlyExitWrapper) -> None: + super().__init__() + self.wrapper = wrapper + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, ...]: + outputs = self.wrapper(x, mode="training") + return tuple(outputs) + + +def export_to_onnx( + model: EarlyExitWrapper, + path: str | Path, + input_shape: tuple[int, int, int, int] = (1, 3, 224, 224), + opset: int = 17, + dynamic_batch: bool = True, +) -> list[str]: + """Export ``model`` to a static multi-output ONNX graph and return the output + names (``["exit_0", ..., "exit_{k-1}", "final"]``). + + The graph computes every exit (routing is applied by the caller at runtime — + see the module docstring). Uses the legacy TorchScript exporter + (``dynamo=False``): the all-exits forward is static, so it traces cleanly, + whereas the routing forward's dynamic control flow does not. + + Parameters + ---------- + model: + A trained :class:`EarlyExitWrapper`. + path: + Destination ``.onnx`` file. + input_shape: + Example input shape for tracing (and the static shape unless + ``dynamic_batch``). + opset: + ONNX opset version. + dynamic_batch: + If True, mark the batch dimension of the input and every output as + dynamic so the graph accepts any batch size. + + Raises ``ImportError`` if the ``onnx`` package is missing (install with + ``pip install 'earlyon[onnx]'``) — the legacy exporter loads it at call time. + """ + try: + import onnx # noqa: F401 # the legacy ONNX exporter imports it internally + except ImportError as exc: + raise ImportError( + "ONNX export requires the 'onnx' package: pip install 'earlyon[onnx]'" + ) from exc + + n_exits = len(model.config.exit_points) + output_names = [f"exit_{i}" for i in range(n_exits)] + ["final"] + + dynamic_axes: dict[str, dict[int, str]] | None = None + if dynamic_batch: + dynamic_axes = {"input": {0: "batch"}} + dynamic_axes.update({name: {0: "batch"} for name in output_names}) + + # torch.onnx.export leaves the module in train mode afterwards; save and + # restore the caller's mode so export has no surprising side effect (a + # train-mode model would re-enable dropout / update BatchNorm on next call). + was_training = model.training + model.eval() + adapter = _AllExitsModule(model) # after eval() so the adapter is eval too + dummy = torch.zeros(input_shape) + + export_kwargs: dict[str, Any] = { + "input_names": ["input"], + "output_names": output_names, + "dynamic_axes": dynamic_axes, + "opset_version": opset, + } + # ``dynamo`` only exists on torch>=2.4; older torch has the legacy + # TorchScript exporter as the only option, so omitting it is equivalent. + # On torch>=2.9 the new exporter is the default and trips the wrapper's + # _is_compiling() guard, so we must force the legacy path with dynamo=False. + if "dynamo" in inspect.signature(torch.onnx.export).parameters: + export_kwargs["dynamo"] = False + + try: + with warnings.catch_warnings(): + # the legacy exporter is deprecated on torch>=2.9 but is the only + # path that traces the wrapper; silence its expected noise. + warnings.simplefilter("ignore", DeprecationWarning) + torch.onnx.export(adapter, (dummy,), str(path), **export_kwargs) + finally: + model.train(was_training) + return output_names diff --git a/pyproject.toml b/pyproject.toml index ace862c..8cc8e47 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,7 +38,10 @@ dev = [ "isort>=5.12", "ruff>=0.1.0", "mypy>=1.0", + "onnx>=1.14", # required by export_to_onnx (legacy exporter loads it) + "onnxruntime>=1.16", # verifies exported ONNX graphs in tests ] +onnx = ["onnx>=1.14"] [project.urls] Homepage = "https://github.com/sohams25/earlyon" diff --git a/tests/test_integration.py b/tests/test_integration.py new file mode 100644 index 0000000..51d4790 --- /dev/null +++ b/tests/test_integration.py @@ -0,0 +1,127 @@ +"""End-to-end integration tests — the pieces working *together*. + +Each test drives a full pipeline (build → train → calibrate → analyze → +benchmark → save/load → ONNX) on a small model, across both routing policies and +both a factory backbone and a custom-wrapped one. +""" + +import numpy as np +import pytest +import torch +from torch.utils.data import DataLoader, TensorDataset + +from earlyon.benchmarking import benchmark_wrapper_on_loader, evaluate +from earlyon.core.thresholds import calibrate_thresholds +from earlyon.models import cifar_resnet_ee, custom_ee +from earlyon.onnx import export_to_onnx +from earlyon.training import joint_train_backbone_and_exits, stage2_train_exits +from earlyon.utils import load_wrapper, save_wrapper +from tests.fixtures.tiny_models import TinyTokenBackbone + +pytest.importorskip("onnx") # export_to_onnx needs it; skip cleanly if absent +ort = pytest.importorskip("onnxruntime") + + +def _loader(n=24, bs=8): + return DataLoader( + TensorDataset(torch.randn(n, 3, 32, 32), torch.randint(0, 10, (n,))), batch_size=bs + ) + + +def _loader1(n=24): + return DataLoader( + TensorDataset(torch.randn(n, 3, 32, 32), torch.randint(0, 10, (n,))), batch_size=1 + ) + + +def test_pipeline_cifar_resnet_confidence_train_calibrate_save_onnx(tmp_path): + """Factory backbone, confidence routing: train exits → calibrate → analyze → + benchmark → save/load (round-trips thresholds + routing) → ONNX (matches).""" + model = cifar_resnet_ee(num_classes=10, depth=8) + stage2_train_exits(model, _loader(), epochs=1, device="cpu", on_epoch_end=lambda _: None) + + result = calibrate_thresholds(model, _loader1(), target_accuracy_drop=0.5, device="cpu") + assert result.policy == "confidence" + assert model.config.confidence_thresholds == result.thresholds + + report = evaluate(model, _loader1(), device="cpu") + assert 0.0 <= report.overall_accuracy <= 1.0 + assert abs(sum(report.exit_distribution.values()) - 1.0) < 1e-6 + + bench = benchmark_wrapper_on_loader(model, _loader1(), device="cpu", num_warmup=2, num_runs=8) + assert bench.throughput_ips > 0 + + path = tmp_path / "m.pth" + save_wrapper(model, path) + loaded = load_wrapper(path) + assert loaded.config.backbone == "cifar_resnet8" + assert loaded.config.confidence_thresholds == result.thresholds + + # same weights + thresholds -> identical routing on a fixed input + x = torch.randn(1, 3, 32, 32) + model.eval() + loaded.eval() + with torch.no_grad(): + r1 = model(x, mode="inference") + r2 = loaded(x, mode="inference") + assert r1.exit_taken == r2.exit_taken + assert torch.allclose(r1.prediction, r2.prediction, atol=1e-5) + + onnx_path = tmp_path / "m.onnx" + export_to_onnx(loaded, onnx_path, input_shape=(1, 3, 32, 32)) + with torch.no_grad(): + ref = loaded(x, mode="training") + outs = ort.InferenceSession(str(onnx_path)).run(None, {"input": x.numpy()}) + for o, r in zip(outs, ref): + assert np.allclose(o, r.numpy(), atol=1e-4) + + +def test_pipeline_custom_token_entropy_joint_calibrate_onnx(tmp_path): + """Custom-wrapped transformer, entropy routing: joint train → entropy + calibrate (updates entropy_thresholds) → analyze → ONNX. Custom models are + not load_wrapper-able, which the pipeline asserts.""" + model = custom_ee( + TinyTokenBackbone(num_classes=10), + ["block0", "block1"], + num_classes=10, + input_shape=(1, 3, 32, 32), + routing_policy="entropy", + ) + joint_train_backbone_and_exits( + model, _loader(), epochs=1, device="cpu", on_epoch_end=lambda _: None + ) + + result = calibrate_thresholds(model, _loader1(), target_accuracy_drop=1.0, device="cpu") + assert result.policy == "entropy" + assert model.config.entropy_thresholds == result.thresholds + + report = evaluate(model, _loader1(), device="cpu") + assert set(report.exit_distribution).issubset({"exit_0", "exit_1", "final"}) + + # custom models save the state_dict but can't be rebuilt via load_wrapper + save_wrapper(model, tmp_path / "c.pth") + with pytest.raises(NotImplementedError, match="custom_ee"): + load_wrapper(tmp_path / "c.pth") + + # ONNX export still works (the token/3D head path) + onnx_path = tmp_path / "c.onnx" + export_to_onnx(model, onnx_path, input_shape=(1, 3, 32, 32)) + x = torch.randn(1, 3, 32, 32) + model.eval() + with torch.no_grad(): + ref = model(x, mode="training") + outs = ort.InferenceSession(str(onnx_path)).run(None, {"input": x.numpy()}) + for o, r in zip(outs, ref): + assert np.allclose(o, r.numpy(), atol=1e-4) + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device") +def test_pipeline_on_cuda_train_calibrate_analyze(tmp_path): + """The whole train → calibrate → analyze pipeline runs on CUDA end to end.""" + model = cifar_resnet_ee(num_classes=10, depth=8) + stage2_train_exits(model, _loader(), epochs=1, device="cuda", on_epoch_end=lambda _: None) + result = calibrate_thresholds(model, _loader1(), target_accuracy_drop=0.5, device="cuda") + assert len(result.thresholds) == len(model.config.exit_points) + report = evaluate(model, _loader1(), device="cuda") + assert 0.0 <= report.overall_accuracy <= 1.0 diff --git a/tests/test_onnx.py b/tests/test_onnx.py index a7872e5..f557e1c 100644 --- a/tests/test_onnx.py +++ b/tests/test_onnx.py @@ -1,10 +1,173 @@ -"""ONNX export tests — feature not yet supported (stub raises).""" +"""ONNX export tests — the exported static graph must reproduce the wrapper's +all-exits (training-mode) outputs under onnxruntime.""" +import numpy as np import pytest +import torch +from earlyon.models import custom_ee from earlyon.onnx import export_to_onnx +from tests.fixtures.tiny_models import TinyBackbone, TinyTokenBackbone +pytest.importorskip("onnx") # export_to_onnx needs it; skip cleanly if absent +ort = pytest.importorskip("onnxruntime") -def test_export_raises_not_implemented(): - with pytest.raises(NotImplementedError, match="not yet supported"): - export_to_onnx(None, "/tmp/x") + +def _cnn_model(): + return custom_ee( + TinyBackbone(num_classes=10), + ["stage1", "stage2"], + num_classes=10, + input_shape=(1, 3, 32, 32), + ).eval() + + +def _token_model(): + return custom_ee( + TinyTokenBackbone(num_classes=10), + ["block0", "block1"], + num_classes=10, + input_shape=(1, 3, 32, 32), + ).eval() + + +def _run_onnx(path, x): + session = ort.InferenceSession(str(path)) + return session.run(None, {"input": x.numpy()}) + + +@pytest.mark.parametrize("builder", [_cnn_model, _token_model], ids=["cnn", "token"]) +def test_export_outputs_match_torch(builder, tmp_path): + """onnxruntime outputs equal the wrapper's training-mode logits — for both + 4D-conv and 3D-token exit heads.""" + model = builder() + path = tmp_path / "m.onnx" + names = export_to_onnx(model, path, input_shape=(1, 3, 32, 32)) + assert names == ["exit_0", "exit_1", "final"] + assert path.exists() + + x = torch.randn(1, 3, 32, 32) + with torch.no_grad(): + ref = model(x, mode="training") + outs = _run_onnx(path, x) + assert len(outs) == len(ref) == 3 + for o, r in zip(outs, ref): + assert np.allclose(o, r.numpy(), atol=1e-4) + + +@pytest.mark.parametrize("builder", [_cnn_model, _token_model], ids=["cnn", "token"]) +def test_export_dynamic_batch_accepts_other_batch_sizes(builder, tmp_path): + """Traced at batch 1, the dynamic-batch graph must still run (and match) at + batch 4 — for both conv and transformer traces.""" + model = builder() + path = tmp_path / "dyn.onnx" + export_to_onnx(model, path, input_shape=(1, 3, 32, 32), dynamic_batch=True) + + x = torch.randn(4, 3, 32, 32) + with torch.no_grad(): + ref = model(x, mode="training") + outs = _run_onnx(path, x) + assert outs[0].shape == (4, 10) + for o, r in zip(outs, ref): + assert np.allclose(o, r.numpy(), atol=1e-4) + + +def test_export_static_batch_rejects_other_batch_sizes(tmp_path): + """Without dynamic_batch the graph is fixed to the traced batch size.""" + model = _cnn_model() + path = tmp_path / "static.onnx" + export_to_onnx(model, path, input_shape=(1, 3, 32, 32), dynamic_batch=False) + + # batch 1 works + outs = _run_onnx(path, torch.randn(1, 3, 32, 32)) + assert outs[0].shape == (1, 10) + # batch 4 is rejected by the static graph (onnxruntime dimension mismatch) + with pytest.raises(Exception, match=r"invalid dimensions|Got: 4|index"): + _run_onnx(path, torch.randn(4, 3, 32, 32)) + + +@pytest.mark.parametrize("start_training", [True, False]) +def test_export_preserves_model_mode(start_training, tmp_path): + """Export must not flip the model's train/eval mode (torch.onnx.export does; + we restore it).""" + model = _cnn_model() + model.train(start_training) + export_to_onnx(model, tmp_path / "m.onnx", input_shape=(1, 3, 32, 32)) + assert model.training is start_training + + +def _route_from_onnx(outs, thresholds, policy, temperature=1.0): + """Reproduce the wrapper's routing rule from the exported all-exits outputs: + return (chosen_exit_index, argmax_prediction).""" + + def softmax(z): + z = z - z.max(axis=-1, keepdims=True) + e = np.exp(z / temperature) + return e / e.sum(axis=-1, keepdims=True) + + for i, logits in enumerate(outs[:-1]): + p = softmax(logits) + if policy == "confidence": + if p.max() >= thresholds[i]: + return i, int(logits.argmax()) + else: # entropy + ent = float(-(p * np.log(np.clip(p, 1e-12, None))).sum()) + if ent <= thresholds[i]: + return i, int(logits.argmax()) + return -1, int(outs[-1].argmax()) + + +@pytest.mark.parametrize("force_exit", [True, False], ids=["exit0", "final"]) +def test_onnx_routing_matches_torch_confidence(force_exit, tmp_path): + """The deployment contract: applying the confidence routing rule to the ONNX + outputs reproduces the torch wrapper's inference decision and prediction.""" + model = _cnn_model() + thr = [0.0, 0.0] if force_exit else [1.01, 1.01] + model.config.confidence_thresholds = thr + path = tmp_path / "m.onnx" + export_to_onnx(model, path, input_shape=(1, 3, 32, 32)) + + x = torch.randn(1, 3, 32, 32) + idx, pred = _route_from_onnx(_run_onnx(path, x), thr, "confidence") + with torch.no_grad(): + ref = model(x, mode="inference") + assert idx == ref.exit_taken == (0 if force_exit else -1) + assert pred == int(ref.prediction.argmax()) + + +@pytest.mark.parametrize("force_exit", [True, False], ids=["exit0", "final"]) +def test_onnx_routing_matches_torch_entropy(force_exit, tmp_path): + """Same deployment contract for the entropy policy.""" + import math + + model = custom_ee( + TinyBackbone(num_classes=10), + ["stage1", "stage2"], + num_classes=10, + input_shape=(1, 3, 32, 32), + routing_policy="entropy", + ) + thr = [math.log(10) + 1.0] * 2 if force_exit else [0.0, 0.0] + model.config.entropy_thresholds = thr + model.eval() # disable head dropout so the prediction is deterministic + path = tmp_path / "e.onnx" + export_to_onnx(model, path, input_shape=(1, 3, 32, 32)) + + x = torch.randn(1, 3, 32, 32) + idx, pred = _route_from_onnx(_run_onnx(path, x), thr, "entropy") + with torch.no_grad(): + ref = model(x, mode="inference") + assert idx == ref.exit_taken == (0 if force_exit else -1) + assert pred == int(ref.prediction.argmax()) + + +def test_export_output_count_tracks_exits(tmp_path): + """A 3-exit model exports 4 outputs (exit_0..exit_2 + final).""" + model = custom_ee( + TinyBackbone(num_classes=10), + ["stage1", "stage2", "stage3"], + num_classes=10, + input_shape=(1, 3, 32, 32), + ).eval() + names = export_to_onnx(model, tmp_path / "three.onnx", input_shape=(1, 3, 32, 32)) + assert names == ["exit_0", "exit_1", "exit_2", "final"]