ONNX multi-exit export + end-to-end integration tests - #4
Merged
Conversation
ONNX export (removes the documented limitation): - export_to_onnx / `earlyon export` write a static multi-output graph (one output per exit + final). Routing stays at runtime — the graph computes every exit and the caller picks the first confident one. - legacy tracer is required (the new exporter trips the wrapper's _is_compiling guard); the `dynamo` kwarg only exists on torch>=2.4, so it's applied only when present (older torch defaults to the legacy exporter anyway). Saves/ restores the model's train/eval mode; suppresses the exporter's deprecation noise. - `onnx` is a real dependency of the exporter (loaded at call time): declared as an `[onnx]` extra with a clear ImportError, and added to `dev` so CI runs the tests instead of skipping them. - works for conv (4D) and transformer (3D-token) heads; dynamic-batch axis. Testing — in isolation AND together: - tests/test_onnx.py: export match vs torch (cnn+token), dynamic/static batch, mode preservation, output count, and the deployment contract — applying the confidence/entropy routing rule to the ONNX outputs reproduces the torch inference decision + prediction. - tests/test_integration.py: full pipelines (build → train → calibrate conf/ entropy → analyze → benchmark → save/load → ONNX) across a factory backbone and a custom transformer, plus a CUDA run. 164 CPU + 5 GPU tests, 96% coverage.
sohams25
force-pushed
the
feat-onnx-and-integration
branch
from
June 6, 2026 07:37
3549157 to
079c39e
Compare
sohams25
commented
Jun 6, 2026
sohams25
left a comment
Owner
Author
There was a problem hiding this comment.
Review — ✅ ready to merge
CI green across 3.10 / 3.11 / 3.12. This run installs onnx via the dev extra, so the ONNX tests genuinely run (not skip), validating export on CI's torch.
/review + a 4-dimension adversarial review found 11 findings (2 HIGH) — both HIGH were real ship-blockers and are fixed:
dynamo=Falsekwarg only exists on torch>=2.4 (pyproject allows >=2.0): now applied only when the param is present (older torch defaults to the legacy exporter anyway), so export works across the whole supported torch range.onnxwas an undeclared hard runtime dep (legacy exporter loads it; tests guarded ononnxruntime, so a clean runner would have failed, not skipped): now an[onnx]extra with a clear ImportError, added todev, and testsimportorskip("onnx").
Also addressed: deployment-contract test (apply the routing rule to ONNX outputs → matches torch inference, confidence + entropy), deprecation-warning suppression, dynamic-batch token coverage, narrowed the static-reject exception, eval-before-adapter, and the CLI now echoes the traced input shape.
Tested in isolation and together: 164 CPU + 5 GPU, 96% coverage. Merging.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Removes the last shipped limitation (ONNX) and adds comprehensive isolation+integration testing.
ONNX export
export_to_onnx/earlyon exportwrite a static multi-output graph (one output per exit +final). ONNX has no portable per-sample control flow, so routing stays at runtime — the graph computes every exit and the caller picks the first confident one (honest trade: portability over the early-exit compute saving).dynamo=False) traces the static all-exits forward cleanly. Works for conv (4D) and transformer (3D-token) heads, with a dynamic batch axis. Saves/restores the model's train/eval mode (torch.onnx.export otherwise silently leaves it in train).onnxruntimeadded to thedevextra to verify exported graphs.Testing — in isolation and together
tests/test_onnx.py): export cnn + token models, outputs match torch all-exits under onnxruntime, dynamic vs static batch, mode preservation, output-count tracks exits.tests/test_integration.py): full pipelines — build → train → calibrate (confidence + entropy) → analyze → benchmark → save/load → ONNX — across a factory backbone (cifar_resnet) and a custom-wrapped transformer, plus a CUDA end-to-end run.Test plan
--strict— 0 errors-m "not gpu"— 159 passed, 97% coverage-m gpu— 5 passed (RTX 4050)