ONNX multi-exit export + end-to-end integration tests

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

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.
 
 ---
 

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

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()

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

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"

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