diff --git a/CHANGELOG.md b/CHANGELOG.md index 9cb062d..5ce3442 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,9 +2,89 @@ format follows [keep a changelog](https://keepachangelog.com/en/1.1.0/). -## unreleased +## 0.3.0 - 2026-07-22 + +### changed — calibration & routing correctness (breaking-ish, migrated) +- **per-head temperatures**: `EarlyExitConfig.temperatures` maps every head + (each exit and the final classifier) to its own fitted temperature; routing + applies the right one at each exit. Previously one temperature was fit from + the final classifier's logits and reused everywhere. The legacy scalar + `temperature` remains as a constructor convenience (broadcast at + construction) and in v1 checkpoints (migrated on load). +- **explicit exit enablement**: `EarlyExitConfig.enabled_exits` replaces the + threshold sentinels (confidence `1.0` / entropy `0.0`). A disabled exit can + no longer fire on a float-saturated softmax (confidence exactly 1.0 / + entropy exactly 0.0); its head is not even evaluated at inference. v1 + checkpoints migrate sentinels to explicit flags with a warning. +- **centralized config validation**: `EarlyExitConfig.validate()` checks + every invariant (unique non-empty exit/layer names, threshold ranges, + finite positive temperatures, loss-weight shape, enablement length) and is + re-invoked at wrapper construction and checkpoint load. Invalid negative + temperatures now raise instead of being clamped to `1e-6` (which produced + an artificially razor-sharp softmax). +- **staged calibration pipeline**: collect logits once + (`collect_head_logits`), fit per-head temperatures with convergence status + (`fit_head_temperatures` / `TemperatureFitResult`), greedy grid search on + the cache, evaluate. `CalibrationResult` now reports enablement, exit + distribution, sample count, accuracy delta, objective, the honest method + name (`greedy-coordinate-grid`) and a schema version. Exits that never fire + on the calibration set stay disabled instead of shipping enabled-but-idle. + Empty calibration loaders raise immediately. +- **checkpoint format v2**: `format_version`, library version, exit-point + metadata (cross-checked against the rebuilt factory on load), enablement + and per-head temperatures. v1 files migrate deterministically with + warnings; files from a newer format are rejected; `custom_ee` models + reload via `load_wrapper(path, factory=...)`. +- **fair benchmark runner**: `benchmark_models` measures every compared + model (backbone, wrapper, static baseline, quantised variant) on the exact + same preloaded sample sequence with identical warmup, boundary + (`model-only` vs `end-to-end`, labelled) and synchronization; reports + accuracy alongside speed for labelled loaders. The legacy helpers delegate + to it. `docs/benchmarks.json` quarantines pre-v0.3 records under + `legacy_v0_2` as not methodologically comparable. +- **honest compute accounting**: `InferenceResult.computation_used` renamed + to `estimated_backbone_flops_fraction` (read alias kept). The estimator + (`FlopsEstimate`) carries method/reliability/notes, explicitly excludes + exit-head cost and routing overhead, detects reused modules (degrading to + a warned low-confidence uniform fallback), and runs lazily so ViT + construction is fast. +- package description: "production-ready" → "deployment-oriented" — the + README states the primary use case (batch-1 latency-sensitive inference) + and labels pre-v0.3 benchmark numbers as legacy. ### added +- **release-readiness verification**: independent adversarial audit tests + (`tests/test_release_audit.py`) including a global-vs-per-head temperature + counterfactual, batched-path per-head temperatures, NaN-head safety, + calibration determinism, and structural proof that the staged runtime + skips later stages; a genuine v0.2-written checkpoint fixture pinning the + v1→v2 migration; a post-install wheel smoke script + (`scripts/smoke_test.py`); a bounded seeded CUDA evidence runner + (`scripts/evidence_run.py`); `MANIFEST.in` so the sdist stops shipping a + partial test tree; `SECURITY_REVIEW.md` and `docs/MIGRATION.md`. +- **staged deployment** (`earlyon.staged`): a documented protocol (ordered + stages emitting continuation features + exit logits, routing between + stages) with a verified reference splitter for Sequential backbones — + staged execution provably matches eager routing and genuinely skips later + stages, unlike the static multi-output ONNX export (now also available + under the explicit name `export_all_exits_to_onnx`). See + `docs/STAGED_DEPLOYMENT.md`. +- **Jetson monitoring fixes + energy semantics**: `TegrastatsMonitor` is + restartable (the stop event is cleared on start; previously a restarted + monitor silently collected nothing), no-ops on double start, closes its + pipes, and degrades cleanly when tegrastats can't spawn. Missing telemetry + parses to `None` instead of invented zeros. `integrate_energy` / + `EnergySummary` distinguish instantaneous power, window-average power and + trapezoid-integrated energy; energy-per-inference is only reported when a + real integral and inference count exist (`profile_with_energy`). +- **custom_ee**: `example_args`/`example_kwargs` for the inspection forward, + device inference from the backbone's parameters (no more CPU probe of a + CUDA model), per-layer `feature_extractors` for tuple/dict/odd-rank + outputs (applied in the dry run and the routing hook), and dry-run + validation that every exit layer executes exactly once and in the listed + forward order. +- `docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md`: the split discipline, + calibration stages, estimator limitations and benchmark fairness rules. - **compute-budget calibration**: `calibrate_thresholds_for_budget` / `earlyon calibrate --target-compute`, the mirror of the accuracy-budget search. State the average FLOPs fraction the deployed model may use (e.g. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..1eeef06 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +prune tests diff --git a/OVERNIGHT_REPORT.md b/OVERNIGHT_REPORT.md new file mode 100644 index 0000000..1a4eec8 --- /dev/null +++ b/OVERNIGHT_REPORT.md @@ -0,0 +1,238 @@ +# Overnight report — earlyon hardening (branch `claude/earlyon-hardening`) + +Date: 2026-07-22. Environment: Linux, Python 3.13.9, torch 2.10.0+cu128, +CUDA available (laptop GPU), no Jetson, no TensorRT. + +## 1. Executive summary + +earlyon is now a methodologically sound, deployment-oriented toolkit for +training, calibrating, evaluating and benchmarking early-exit PyTorch CV +models. The overnight work fixed the four correctness defects called out in +the brief (shared temperature, sentinel-threshold "disabling", negative- +temperature clamping, unfair benchmark boundaries), rebuilt calibration as a +staged, testable pipeline with a strict data-split discipline, versioned the +checkpoint format with migration, made every compute number an explicitly +labelled estimate, fixed the tegrastats monitor lifecycle and energy +semantics, extended `custom_ee` safely, and added a verified staged- +deployment contract. Public claims were re-based: "production-ready" is +gone, pre-v0.3 benchmark numbers are quarantined as legacy, batch-1 +latency-sensitive inference is stated as the primary use case. + +## 2. Baseline (before any change) + +Commands and results on `main` (a5f8414): + +- `pytest -m "not gpu"` → **177 passed, 5 deselected, ~14 s** +- `mypy earlyon` (strict) → clean +- `ruff check earlyon tests` → clean +- `black --check` / `isort --check-only` on `earlyon tests` → clean + (`scripts/` was unformatted; CI does not check it — scripts were + reformatted as part of their rewrite) +- No hanging tests; ViT tests were the slowest (~1 s each) because FLOPs + analysis ran at construction. + +## 3. Major architecture & methodology changes + +1. **Per-head temperatures** — `EarlyExitConfig.temperatures: dict[head, T]` + covering every exit and the final classifier; routing applies each head's + own T. `fit_head_temperatures` fits per head from one cached-logits pass; + each fit carries convergence/fallback status (fallback = safe 1.0, never + a sharp value). +2. **Explicit enablement** — `enabled_exits: list[bool]`; disabled exits + never fire (confidence exactly 1.0 / entropy exactly 0.0 included) and + their heads aren't evaluated at inference. Sentinels retired; v1 + checkpoints migrate. +3. **Centralized validation** — `EarlyExitConfig.validate()` (all + invariants; actionable errors) re-invoked at wrapper construction and + checkpoint load. Invalid negative temperature now raises instead of the + old `max(t, 1e-6)` clamp. +4. **Staged calibration pipeline** — collect → fit → search → evaluate as + separate tested functions; the vectorized simulator is pinned exactly + against the real router. `CalibrationResult` reports enablement, exit + distribution, sample count, delta, objective, honest method name, schema + version. Never-firing exits stay disabled. Empty loaders raise. +5. **Fair benchmark runner** — `benchmark_models`: identical preloaded + samples, warmup, boundary and sync for every compared model; model-only + vs end-to-end labelled; accuracy alongside speed for labelled loaders; + batch-1 validated. CLI + scripts use it; legacy helpers delegate to it. +6. **Honest compute accounting** — `estimated_backbone_flops_fraction` + (alias `computation_used` kept) from a lazy, cached `FlopsEstimate` + carrying method/reliability/notes; reused-module detection degrades to a + warned low-confidence uniform fallback; exit-head cost explicitly + excluded. +7. **Checkpoint v2** — versioned schema with exit-point cross-check, + deterministic v1 migration (broadcast temperature; sentinel → disabled), + newer-format rejection, `factory=` loading for custom models. +8. **Jetson** — restartable monitor (stop event cleared on start — the + restart bug), `None` for missing telemetry, `integrate_energy` / + `EnergySummary` separating instantaneous power, window-average power and + trapezoid-integrated energy; per-inference energy only from a real + integral. +9. **custom_ee** — example args/kwargs, backbone-device inference, per-layer + feature extractors (dry run + routing hook), exactly-once and in-order + exit-layer validation, heads placed on the backbone's device. +10. **Staged deployment** — `earlyon.staged` protocol + reference splitter + for Sequential backbones, self-verified against the eager wrapper; + `export_all_exits_to_onnx` alias names the static ONNX export honestly. + +## 4. Bugs reproduced and corrected + +| Defect (verified against the old code) | Fix | +|---|---| +| One temperature fit from final-classifier logits, reused for every exit (docs claimed per-head) | per-head fit + routing (P0-A) | +| Threshold `1.0` "disabled" exits could fire at saturated confidence (`>=`); entropy `0.0` at entropy 0 | explicit `enabled_exits` (P0-B) | +| Negative/zero temperature clamped to `1e-6` → artificially sharp softmax | validate() raises; runtime guard falls back to 1.0 | +| `computation_used` read as measured FLOPs; module-order assumption unchecked | renamed estimate + reuse detection + provenance | +| Wrapper and backbone benchmarked on different inputs/boundaries | `benchmark_models` fairness contract | +| `TegrastatsMonitor` unrestartable (stale stop event); pipe not closed | lifecycle fix + tests | +| Missing tegrastats fields reported as 0.0 (fabricated) | `None` + energy integration semantics | +| ViT FLOPs analysis at construction (slow) | lazy + cached | +| `custom_ee` blindly probed CUDA models with CPU tensors; accepted reused/out-of-order exit layers | device inference + dry-run validation | +| Calibration could enable an exit at a threshold that never fires | never-firing exits stay disabled | +| Empty calibration loader silently returned 0.0/0.0 | raises ValueError | +| LBFGS `exp(log_t)` could overflow to inf | guarded, falls back with status | + +## 5. Files changed by subsystem + +- **core**: `types.py` (config redesign, validation, result renames), + `wrappers.py` (per-head T, enablement, adapters, lazy FLOPs, head/config + cross-checks), `thresholds.py` (staged pipeline rewrite), `temperature.py` + (fit status), `flops.py` (FlopsEstimate, reuse detection, device param) +- **benchmarking**: `throughput.py` (rewritten around `benchmark_models`), + `jetson_profiler.py` (rewritten), `__init__.py` +- **serialization**: `utils.py` (checkpoint v2 + migration + factory) +- **models**: `custom.py` (rewritten) +- **new**: `earlyon/staged.py`, `docs/STAGED_DEPLOYMENT.md`, + `docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md`, `tests/test_staged.py`, + `RELEASE_NOTES_DRAFT.md`, this report +- **cli/scripts/examples/docs**: `cli.py` (benchmark/profile commands), + `scripts/run_benchmarks.py`, `scripts/re_evaluate.py`, + `docs/benchmarks.json` (legacy quarantine), `examples/01`, `examples/02`, + `README.md`, `docs/DESIGN_DECISIONS.md`, `CHANGELOG.md`, `pyproject.toml` +- **tests**: substantial additions across `test_types`, `test_wrappers`, + `test_thresholds`, `test_temperature`, `test_utils`, `test_benchmarks`, + `test_flops`, `test_jetson_profiler`, `test_custom_ee`, `test_staged` + +## 6. Public API / config / checkpoint migration notes + +- `EarlyExitConfig`: new `temperatures` (dict) and `enabled_exits` fields; + `temperature` is constructor-only (broadcast) — mutate `temperatures`. + Constructing with both `temperature != 1.0` and `temperatures` raises. +- `InferenceResult`/`BatchedInferenceResult`/`BenchmarkResult`/`JetsonRun`: + FLOPs field renamed; `computation_used` remains a read-only alias. + Constructing these with `computation_used=` no longer works (library- + internal construction only). +- `CalibrationResult`: superset of old fields; `fitted_temperature` is now + the final head's fit (deprecated alias). +- Checkpoints: v2 written; v1 read with deterministic migration + warnings; + `load_wrapper(..., factory=...)` for custom models; newer formats rejected. +- `benchmark_wrapper` / `benchmark_backbone` / `benchmark_wrapper_on_loader` + keep signatures (now thin wrappers over `benchmark_models`). + +## 7. Tests and final gate results + +- Baseline 177 tests → **261 tests**: `pytest` (full, incl. 6 GPU tests on + the CUDA host) → **261 passed** (~16 s); `pytest -m "not gpu"` → 255 + passed. +- Coverage (CPU subset): **95%** (`--cov=earlyon`). +- `mypy earlyon` (strict): clean, 28 files. +- `ruff check earlyon tests`: clean. `black --check` / `isort --check-only`: + clean (scripts now formatted too). +- The former ViT slowness: construction no longer runs FLOPs analysis; the + ViT test module runs in ~2 s dominated by torchvision model build. + +## 8. Benchmarks actually run + +None published. All numeric records in `docs/benchmarks.json` are the +pre-existing v0.2 measurements, now explicitly quarantined under +`legacy_v0_2` with a note that they are not comparable to the v0.3 fair +runner (the old runner fed the wrapper and backbone different inputs). The +new runner is fully implemented and tested with tiny deterministic models; +regenerating real numbers is one command (`python scripts/run_benchmarks.py`) +on the target machine and was deliberately not run overnight to avoid mixing +a multi-hour training job into a correctness review. + +## 9. Hardware checks not run, and the later procedure + +No Jetson, no TensorRT, no new CUDA benchmark numbers. Procedure for later: + +1. On the Jetson (MAXN, fans steady): `pip install -e .`, verify + `tegrastats` on PATH. +2. Train/calibrate on a workstation, `save_wrapper`, copy the `.pth`. +3. `earlyon profile --model calibrated.pth --runs 500 --warmup 50` — the + energy block reports integrated mJ and mJ/inference from real samples. +4. `python scripts/re_evaluate.py` for the fair wrapper-vs-backbone table. +5. TensorRT: per-stage export + `trtexec` loop per + `docs/STAGED_DEPLOYMENT.md` (procedure documented, nothing fabricated). + +## 10. Remaining limitations / deferred work + +- Staged splitting covers Sequential backbones only; torchvision factories + need per-architecture stage maps (protocol + acceptance test exist). +- Masked per-sample batched routing is still roadmap. +- Per-stage ONNX export helper is documented but not implemented as a + one-call function (each `Stage` traces cleanly with `torch.onnx.export`). +- README benchmark table still shows legacy v0.2 numbers (labelled); + regenerate before release. +- `evaluate`/`AccuracyReport` still uses the `avg_computation_used` name + (documented as an estimate); renaming it was skipped to limit churn. + +## 11. Suggested next version & release checklist + +**0.3.0** — see `RELEASE_NOTES_DRAFT.md` (includes the checklist). + +## 12. Local commits (oldest first) + +1. `3e69bbc` feat: per-head temperatures, explicit exit enablement, staged + calibration, checkpoint v2 +2. `5f03a98` feat: fair benchmark runner with identical boundaries and inputs +3. `5724afe` feat: honest compute accounting — lazy FLOPs estimate with provenance +4. `c732544` fix: restartable tegrastats monitor, honest power/energy semantics +5. `7b84818` feat: safer, more capable custom_ee +6. `43e82b1` feat: staged deployment contract with verified reference splitter +7. `a750586` fix: cast ModuleList access in staged runtime for strict mypy +8. (docs commit) docs: repositioning, contracts, changelog, release draft, + overnight report + +## 13. Five-minute interview explanation + +**Early-exit architecture.** A CNN spends the same compute on every image, +but most images are easy. earlyon attaches small classifier heads (~10-100k +params) at intermediate layers of an unmodified backbone via forward hooks. +At inference, each enabled head computes a softmax; if its calibrated +criterion passes (max-probability above a threshold, or entropy below one), +a sentinel exception unwinds the forward and that head's logits are the +answer — later layers never run. Training is two-stage: train the backbone +normally, freeze it (parameters *and* BatchNorm stats), then train the heads +for a few epochs with a weighted multi-exit cross-entropy. + +**Calibration methodology.** Thresholding raw softmax is meaningless because +each head is differently over-confident, so we first fit one temperature per +head (Guo et al. 2017) on a held-out split — NLL minimization that rescales +confidence without changing predictions. Then a greedy per-exit grid search +on a separate calibration split picks the most aggressive thresholds that +keep routed accuracy within the user's budget (or meet a compute budget with +least accuracy loss). The search runs the network once — logits are cached +per head and every grid point is simulated vectorially, pinned bit-exact +against the real router by a test. Exits that never fire stay explicitly +disabled. Test labels are never visible to the search; reported accuracy +comes from a separate evaluation pass. + +**Why FLOP savings may not become latency savings.** The FLOPs number is a +static estimate of the backbone fraction executed — it excludes the exit +heads and, critically, the routing overhead: each enabled exit computes a +softmax and calls `.item()`, which synchronizes the GPU pipeline. On a fast +GPU with a small backbone, one host sync can cost more than the skipped +layers; kernels are also more efficient in bulk. So a model can "save" 20% +of estimated FLOPs and be *slower* in wall-clock — our own legacy ResNet50 +row shows exactly that (0.94× on real inputs). + +**How the project measures this honestly.** One benchmark core feeds every +compared model — backbone, wrapper, and ideally a smaller static model at +matched accuracy — the exact same preloaded sample sequence, same warmup, +same synchronization, with the measurement boundary labelled (model-only vs +end-to-end). Estimated FLOPs and wall-clock are reported side by side and +never conflated; noise-input runs are labelled best-case; results that +predate this methodology are quarantined as legacy; and on Jetson, energy is +only reported as the time-integral of real power samples over the timed +window — never extrapolated from a single instantaneous reading. diff --git a/PR_BODY_DRAFT.md b/PR_BODY_DRAFT.md new file mode 100644 index 0000000..f328471 --- /dev/null +++ b/PR_BODY_DRAFT.md @@ -0,0 +1,73 @@ +# Release v0.3.0: defensible calibration, fair benchmarks, and staged inference + +> Base: `main`. Head: `claude/earlyon-release-readiness` +> (includes the `claude/earlyon-hardening` work). + +## What this PR is + +A correctness and methodology release. No new backbones, no new headline +claims — the calibration math, the measurement methodology, the +serialization contract and the public claims are made defensible, and every +fix is pinned by tests (261 → 265+ tests, 95% coverage, strict mypy). + +## The four core defects fixed + +1. **One temperature for every head.** Temperature scaling was fit once from + the final classifier's logits and reused at every exit, though each exit + head is a differently miscalibrated classifier. Now: per-head fit with + convergence status, per-head application at routing time. A test proves a + single global temperature misroutes where per-head does not. +2. **"Disabled" exits could fire.** Disabling used threshold sentinels + (confidence 1.0 / entropy 0.0) while routing compares with `>=`/`<=`, so + a float32-saturated softmax fired a "disabled" exit. Now: explicit + `enabled_exits` booleans; sentinel-carrying v1 checkpoints migrate. +3. **Invalid temperatures were clamped sharp.** A negative fitted/user + temperature was clamped to `1e-6`, producing an artificially razor-sharp + softmax. Now: centralized `EarlyExitConfig.validate()` raises; runtime + guard falls back to the no-op 1.0. +4. **Benchmarks weren't comparable.** The wrapper and backbone were measured + on different inputs and boundaries. Now: `benchmark_models` feeds every + compared model the identical preloaded sample sequence with identical + warmup/boundary/sync; old records quarantined as `legacy_v0_2`. + +## Also in this PR + +- Staged calibration pipeline + rich `CalibrationResult` (enablement, exit + distribution, sample count, honest `greedy-coordinate-grid` method name). +- Checkpoint `format_version: 2` + deterministic v1 migration, tested + against a genuine v0.2-written fixture; `factory=` loading for custom + models. +- `estimated_backbone_flops_fraction` (alias kept): lazy `FlopsEstimate` + with provenance, reused-module detection, exit-head cost disclosed. +- Restartable tegrastats monitor; instantaneous power vs window-average vs + integrated energy kept distinct; missing telemetry is `None`. +- `custom_ee`: example args/kwargs, device inference, feature extractors, + exactly-once/in-order exit validation. +- Staged deployment contract + reference splitter for Sequential backbones, + proven equivalent to eager routing and proven to skip later stages. +- Packaging: wheel/sdist build + twine-clean; clean-env wheel smoke script; + sdist no longer ships a partial test tree. +- Bounded, seeded CUDA evidence run on CIFAR-10 (ResNet-18 EE vs backbone vs + static MobileNetV2) with disjoint train/temperature/calibration/test + splits — results in `docs/evidence/`. +- Docs: calibration/benchmark contract, migration guide, staged deployment + guide, security review, honest README repositioning + ("deployment-oriented", legacy labels, batch-1 scope). + +## Compatibility + +Old checkpoints load (migrated, warned). `computation_used`, +`fitted_temperature` and the single-model benchmark helpers remain as +compatible aliases/wrappers. Breaking edges are listed in +`docs/MIGRATION.md`. + +## Test plan + +- [x] `pytest` full suite (CPU + CUDA host): all green +- [x] `mypy earlyon` strict, `ruff`, `black --check`, `isort --check-only` +- [x] wheel + sdist build, `twine check`, clean-venv install + out-of-repo + smoke (`scripts/smoke_test.py`) +- [x] genuine v1 checkpoint migration fixture +- [x] bounded CUDA evidence run (`scripts/evidence_run.py`) +- [ ] optional before tagging: full-length `scripts/run_benchmarks.py` on + target hardware; Jetson procedure from `docs/STAGED_DEPLOYMENT.md` diff --git a/README.md b/README.md index 2725820..f4cddf8 100644 --- a/README.md +++ b/README.md @@ -5,9 +5,9 @@

ci python - coverage + coverage mypy - version + version license

@@ -26,11 +26,14 @@ from earlyon.models import resnet50_ee model = resnet50_ee(num_classes=10, pretrained=True).eval() result = model(torch.randn(1, 3, 224, 224), mode="inference") -result.exit_taken # which head answered; -1 means the full network ran -result.computation_used # fraction of the network's FLOPs actually spent -result.confidence # softmax confidence at the point it answered +result.exit_taken # which head answered; -1 = full network +result.estimated_backbone_flops_fraction # static estimate, not a measurement +result.confidence # softmax confidence where it answered ``` +earlyon is a research-to-deployment toolkit. Its primary target is batch-1, +latency-sensitive inference; batched routing exists but is conservative. +

earlyon demo: an easy image exits at the first head using 12% of the FLOPs, a hard image runs the whole network

@@ -92,33 +95,49 @@ model = custom_ee(backbone, exit_layers=["layer2", "layer3"], num_classes=10) Routing has two policies. `"confidence"` (default) exits when `softmax(logits).max() >= threshold`; `"entropy"` exits when the softmax entropy drops below a threshold, which reads the whole distribution instead -of just the top class. Calibration and `save_wrapper`/`load_wrapper` are -policy-aware, so a calibrated entropy model reloads as one. +of just the top class. Whether an exit may fire at all is an explicit +per-exit boolean (`enabled_exits`) — a disabled exit can never fire, not +even at softmax confidence exactly 1.0. Calibration and +`save_wrapper`/`load_wrapper` are policy-aware, so a calibrated entropy +model reloads as one; checkpoints carry a versioned schema (v2) and older +files migrate with a warning. Training is two-stage by default: train the backbone exactly as you already do, freeze it (parameters and BatchNorm stats), then train the small heads for a few epochs. `joint_train_backbone_and_exits` does end-to-end training instead when you want the last bit of accuracy and have the budget. Temperature scaling (Guo et al. 2017) can be fit before calibration to fix -the usual softmax over-confidence; pass `fit_temperature=True`. +the usual softmax over-confidence; pass `fit_temperature=True` and each head +— every exit and the final classifier — gets its own fitted temperature. +The full split discipline (what may touch which data) is written down in +[`docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md`](docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md). -## Measured results +## Measured results (legacy: earlyon v0.2 methodology) CIFAR-10, trained on an RTX 4050 Laptop GPU (6 GB) from ImageNet-pretrained weights: 3–4 epochs for the backbone, 3–4 for the exit heads, thresholds -calibrated to a 1% accuracy budget. - -| Model | Test Acc | Baseline Acc | Avg FLOPs Used | % exited early | -|-------------|---------:|-------------:|---------------:|---------------:| -| ResNet18 | 94.42% | 96.32% | 89.88% | 35.3% | -| ResNet50 | 95.88% | 97.60% | 81.42% | 58.2% | -| MobileNetV2 | 93.31% | 95.08% | 93.90% | 8.5% | - -Avg FLOPs Used is measured on the real test set, per image. It is the number -I'd want to see, so it's the one the tool reports. For ResNet50 that means -~19% of the compute gone for a 1.7% accuracy cost, with 58% of images never -reaching `layer4`. Wall-clock speedup is hardware- and input-dependent; see -the [latency appendix](#appendix-wall-clock-latency) before quoting one. +calibrated to a 1% accuracy budget. **These runs predate v0.3**: they used a +single shared temperature and the pre-fair-runner benchmark; treat them as +indicative, not as v0.3 results. Regenerate with +`python scripts/run_benchmarks.py` (records land in +`docs/benchmarks.json` under `runs`; these live under `legacy_v0_2`). + +| Model | Test Acc | Baseline Acc | Est. FLOPs fraction | % exited early | +|-------------|---------:|-------------:|--------------------:|---------------:| +| ResNet18 | 94.42% | 96.32% | 89.88% | 35.3% | +| ResNet50 | 95.88% | 97.60% | 81.42% | 58.2% | +| MobileNetV2 | 93.31% | 95.08% | 93.90% | 8.5% | + +The FLOPs fraction is a per-image *estimate* over the real test set — a +static analysis of the backbone that excludes the exit heads' own small cost +and all routing overhead. For ResNet50 it reads as ~19% of estimated compute +gone for a 1.7% accuracy cost, with 58% of images never reaching `layer4`. +An estimated FLOPs saving is not a latency claim: wall-clock speedup is +hardware- and input-dependent; see the +[latency appendix](#appendix-wall-clock-latency-legacy) before quoting one. +When judging a deployment, also benchmark a smaller static model at matched +accuracy (pass it as a third entry to `benchmark_models`) — if a plain +ResNet18 matches your routed ResNet50, ship the ResNet18. The MobileNetV2 row is a loss, and it's in the table anyway. An already-compressed backbone leaves little for early exit to skim: only 8.5% @@ -130,6 +149,25 @@ Reproduce with `python scripts/run_benchmarks.py` (trains and benchmarks) or `python scripts/re_evaluate.py` (from checkpoints). Raw numbers live in [`docs/benchmarks.json`](docs/benchmarks.json). +## Measured results (v0.3 fair runner — bounded evidence run) + +One seeded, deliberately small validation run (2 epochs per stage, RTX 4050 +Laptop GPU, CIFAR-10 at 224px, disjoint train/temperature/calibration/test +splits, identical samples and boundaries for all three models). Full data +and honest interpretation: [`docs/evidence/CUDA_EVIDENCE.md`](docs/evidence/CUDA_EVIDENCE.md). + +| Model | Test acc | p50 | p95 | Throughput | Est. FLOPs fraction | +|---|---:|---:|---:|---:|---:| +| ResNet-18 backbone | 94.07% | 1.33 ms | 1.51 ms | 712 ips | 1.00 | +| ResNet-18 early-exit | 92.92% | 1.43 ms | 1.46 ms | 782 ips (**1.10×**) | 0.88 (estimate) | +| MobileNetV2 static baseline | 92.67% | 1.40 ms | 1.48 ms | 702 ips (0.99×) | 1.00 | + +Read the negative parts too: the early-exit **median** latency is worse than +the backbone's (routing overhead on the 70% of images that don't exit); the +gain is in throughput and the tail (p95/p99), and the static baseline is +competitive at this training budget. That's the honest shape of early exit +on a fast GPU — reproduce with `python scripts/evidence_run.py`. + ## How it compares Compression makes the model cheaper for every input; early exit spends @@ -219,6 +257,35 @@ Every command takes `--help`. The rest of the toolkit: `train joint`, `calibrate --target-compute`, `analyze` (per-exit accuracy), `profile` (Jetson power), `export` (ONNX). +## Architecture + +``` + input (batch 1) + │ + ┌──────────▼──────────┐ hook ┌────────────┐ softmax(logits / T_e0) + │ backbone stage 0..k │──────────▶│ exit head 0│──┐ enabled? confidence ≥ thr? + └──────────┬──────────┘ └────────────┘ │ yes → return logits, STOP + │ (only if exit 0 didn't fire) │ no ▼ continue + ┌──────────▼──────────┐ hook ┌────────────┐ │ + │ backbone stage k..m │──────────▶│ exit head 1│──┤ each head has its OWN + └──────────┬──────────┘ └────────────┘ │ temperature T and its + │ │ own enabled flag + ┌──────────▼──────────┐ │ + │ rest of backbone │───────▶ final classifier ─┘ (T_final) + └─────────────────────┘ + + training mode: every head produces logits (no routing) → weighted CE loss + calibration: collect logits once → fit T per head → greedy threshold grid + deployment: eager wrapper (above), or staged stages (docs/STAGED_DEPLOYMENT.md) +``` + +Three concerns, kept separate: **training** (two-stage by default: your +backbone recipe untouched, then small heads on a frozen backbone), +**calibration** (per-head temperatures + threshold/enablement search on held- +out splits), and **routing** (pure inference-time policy driven by the +config). The full contract is in +[`docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md`](docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md). + ## How it works 1. Forward hooks attach the exit heads at the named layers, so the backbone @@ -238,19 +305,54 @@ The reasoning for each choice, including the ugly parts, is in ## Limitations -- `forward(mode="inference")` is batch-size 1. `forward_inference_batched` - handles batches, but conservatively (see above); masked per-sample routing - inside a batch is on the roadmap. +- `forward(mode="inference")` is batch-size 1 — the latency-sensitive edge + case earlyon targets. `forward_inference_batched` handles batches, but + conservatively (see above); masked per-sample routing inside a batch is on + the roadmap. +- Eager routing has real overhead: every enabled exit evaluates its head and + synchronises the host (`.item()`) to decide. On a GPU that synchronisation + can cost more than the skipped layers save, particularly on small + backbones — theoretical FLOP savings do not guarantee latency savings. + The fair runner reports both so you can see it. - `torch.compile` cannot trace the routing control flow. The wrapper raises a clear error instead of silently falling back; compile the raw backbone if you need it. -- ONNX export writes a static graph that computes every exit and leaves - routing to the caller. The per-sample compute saving only exists in the - PyTorch wrapper. +- ONNX export (`export_all_exits_to_onnx`) writes a static graph that + computes **every** exit and leaves routing to the caller — it does not + short-circuit compute. For a deployable split that genuinely skips later + stages, see [`docs/STAGED_DEPLOYMENT.md`](docs/STAGED_DEPLOYMENT.md) + (reference implementation for Sequential backbones). +- `computation_used` / `estimated_backbone_flops_fraction` is a static + estimate: exit-head cost and routing overhead excluded; backbones that + reuse modules degrade to a warned low-confidence uniform estimate. - Compute budgets from `calibrate_thresholds_for_budget` hold as an average over the calibration distribution. Nothing caps FLOPs per sample. -- Benchmarks are CIFAR-10 on a laptop GPU. ImageNet-scale numbers and a real - Jetson table don't exist yet; treat any wall-clock claim accordingly. +- Benchmarks are CIFAR-10 on a laptop GPU, and the published table predates + the v0.3 fair runner (see the legacy labels). ImageNet-scale numbers and a + real Jetson table don't exist yet; treat any wall-clock claim accordingly. + +## Checkpoint migration + +Checkpoints are versioned (`format_version: 2`). Files written by earlyon +≤ 0.2 load with an automatic, deterministic migration (scalar temperature +broadcast per head; legacy "disabled" threshold sentinels become explicit +`enabled_exits=False`) and a warning describing what changed — pinned +against a genuine v0.2-written fixture in the test suite. Details and code +changes: [`docs/MIGRATION.md`](docs/MIGRATION.md). + +## Reproducibility + +- `pytest -m "not gpu"` — full CPU suite; `pytest -m gpu` on a CUDA host. +- `python scripts/smoke_test.py` — post-install smoke against an installed + wheel (run it from outside the repo). +- `python scripts/evidence_run.py` — the bounded, seeded CIFAR-10 evidence + experiment (CUDA; explicit train/temperature/calibration/test splits, + epoch caps, wall-clock budget). Output lands in `docs/evidence/`. +- `python scripts/run_benchmarks.py` — the longer train+calibrate+benchmark + pipeline that regenerates `docs/benchmarks.json` (`runs` key). +- Jetson: the exact procedure (tegrastats profiling, per-stage TensorRT) is + in [`docs/STAGED_DEPLOYMENT.md`](docs/STAGED_DEPLOYMENT.md); no Jetson + numbers are published because none have been measured. ## Contributing @@ -281,7 +383,7 @@ If earlyon saves your model some FLOPs, a citation is welcome: year = {2026}, publisher = {GitHub}, howpublished = {\url{https://github.com/sohams25/earlyon}}, - note = {Version 0.2.0} + note = {Version 0.3.0} } ``` @@ -291,7 +393,12 @@ BranchyNet (Teerapittayanon et al. 2016) and the ACM 2024 early-exit survey (10.1145/3698767) for the ideas; torchvision (BSD) for the backbones; fvcore for FLOPs accounting. earlyon itself is MIT. -## Appendix: wall-clock latency +## Appendix: wall-clock latency (legacy) + +**Legacy numbers**: measured with the pre-v0.3 runner, which benchmarked the +wrapper and backbone on *different* random inputs — not methodologically +comparable to `benchmark_models` results, which feed every compared model +the identical sample sequence. Kept for transparency until re-measured. Throughput below uses a random-noise input, which can trigger spurious early exits, so read it as a best-case bound rather than a claim. RTX 4050 Laptop diff --git a/RELEASE_NOTES_DRAFT.md b/RELEASE_NOTES_DRAFT.md new file mode 100644 index 0000000..33e5b65 --- /dev/null +++ b/RELEASE_NOTES_DRAFT.md @@ -0,0 +1,133 @@ +# earlyon v0.3.0 — release notes + +(This file is the canonical text for the v0.3.0 GitHub Release.) + +earlyon 0.3.0 is a correctness release for the early-exit toolkit. It makes +the calibration methodology defensible (per-head temperatures, explicit exit +enablement, strict data-split discipline), the measurements fair (one +benchmark core, identical samples and boundaries for every compared model), +and the claims honest (estimated FLOPs labelled as estimates, legacy numbers +quarantined, "production-ready" retired in favor of "deployment-oriented"). + +## Install + +```bash +pip install earlyon==0.3.0 +``` + +Python 3.10+, torch 2.0+. Existing checkpoints load with automatic +migration — see **Migration** below. + +## Methodological corrections + +- **One temperature per head.** Temperature scaling used to be fit once from + the final classifier's logits and reused at every exit, though each exit + head is a differently miscalibrated classifier. 0.3 fits and applies a + temperature per head (`config.temperatures`), with per-fit convergence + status; a regression test proves a single global temperature misroutes + where per-head fitting does not. +- **Disabled exits cannot fire.** Disabling via threshold sentinels + (confidence `1.0` / entropy `0.0`) was numerically unsound: a + float32-saturated softmax met the criterion. Enablement is now an explicit + boolean per exit (`config.enabled_exits`), honored by the confidence and + entropy policies on the single-sample, batched and staged paths. +- **Staged calibration pipeline.** Collect logits once → fit per-head + temperatures → greedy threshold grid search on the cache → evaluate; each + stage separately testable, the simulator pinned bit-exact against the real + router, threshold search never sees test labels, never-firing exits stay + disabled, empty loaders fail fast. The method is named what it is + (`greedy-coordinate-grid`) — no global-optimality claim. +- **Centralized validation.** `EarlyExitConfig.validate()` raises actionable + errors for invalid config (a negative temperature used to be silently + clamped to `1e-6`, producing an artificially sharp softmax). + +## Fair benchmarking + +`benchmark_models` measures every compared model — full backbone, early-exit +wrapper, a static smaller baseline, a quantised variant — on the **exact +same preloaded sample sequence** with identical warmup, boundary +(model-only vs end-to-end, labelled) and synchronization, reporting +p50/p95/p99, throughput, exit distribution and accuracy for labelled +loaders. Pre-0.3 benchmark records are quarantined as `legacy_v0_2` in +`docs/benchmarks.json`: they fed compared models different inputs and are +not comparable. + +## Honest compute terminology + +`computation_used` is renamed `estimated_backbone_flops_fraction` (read +alias kept). It is a static estimate — it excludes exit-head cost and all +routing overhead, and it is neither a latency nor an energy measurement. +The estimator (`FlopsEstimate`) reports its method and reliability, detects +reused modules (degrading to a warned low-confidence fallback), and runs +lazily so large-model construction is fast. + +## Staged execution + +`earlyon.staged` defines a staged-deployment contract (ordered stages +emitting continuation features + exit logits, routing applied between +stages) with a verified reference splitter for `nn.Sequential` backbones. +Staged execution provably matches eager routing and *structurally* skips +later stages — unlike the all-exits ONNX export +(`export_all_exits_to_onnx`), which always computes every exit and is a +portability artifact, not conditional execution. The supported contract is +deliberately narrow; unsupported graphs are rejected. See +`docs/STAGED_DEPLOYMENT.md`. + +## Checkpoints and migration + +Checkpoints now carry `format_version: 2` (exit-point identity, routing +policy, thresholds, enablement, per-head temperatures, library version) and +are cross-checked against the rebuilt factory on load. Files written by +earlyon ≤ 0.2 migrate automatically and deterministically — scalar +temperature broadcast per head, legacy sentinels → explicit disables — with +a warning; the migration is pinned against a genuine v0.2-written fixture. +`custom_ee` models reload via `load_wrapper(path, factory=...)`. +Details and code-level changes: `docs/MIGRATION.md`. + +## Also fixed + +- `custom_ee`: example args/kwargs for inspection, device inference (no more + CPU probes of CUDA models), per-layer feature extractors for + tuple/dict/odd-rank outputs, and dry-run validation that every exit layer + runs exactly once in the listed order. +- Jetson monitoring: `TegrastatsMonitor` is restartable (a restarted monitor + previously collected nothing), missing telemetry parses to `None` instead + of invented zeros, and `profile_with_energy` distinguishes instantaneous + power, window-average power and trapezoid-integrated energy. + +## Real CUDA evidence (one seeded, bounded run — not a universal claim) + +A deliberately small validation run (CIFAR-10 at 224px, ResNet-18 EE +fine-tuned 2+2 epochs from ImageNet weights, seed 42, disjoint +train/temperature/calibration/test splits, RTX 4050 Laptop GPU; full data in +`docs/evidence/`): + +- early-exit ResNet-18: 92.92% test accuracy vs 94.07% full backbone; + **1.10× throughput** with **worse median latency** (1.43 vs 1.33 ms) and + better tail latency (p95/p99) — routing overhead is real on the ~70% of + images that ran the full network; +- a static MobileNetV2 baseline trained the same way was competitive + (92.67%, 0.99×) — smaller static models remain an essential comparison; +- a noise-input 2.05× figure exists only as a labelled best-case bound, not + a practical speed claim. + +Early exit is **not always faster**; the fair runner exists so you can see +when it isn't. + +## Limitations + +- Dynamic routing is primarily eager PyTorch execution (hooks + host sync); + `torch.compile` is refused with a clear error. +- All-exits ONNX export computes every exit — no conditional execution. +- Staged reference execution supports Sequential backbones only. +- The Jetson/tegrastats procedure is documented but **unmeasured** — no + Jetson numbers exist; TensorRT deployment is not claimed. +- Estimated FLOPs fractions are estimates, not measured latency or energy. +- `forward(mode="inference")` is batch-1; batched routing is conservative. + +## Verification behind this release + +270 tests (incl. CUDA suite on a CUDA host), 95% coverage, mypy strict, +ruff/black/isort clean; wheel + sdist twine-clean and smoked from a clean +environment (`scripts/smoke_test.py`); security review in +`SECURITY_REVIEW.md`; full audit trail in `RELEASE_READINESS_REPORT.md`. diff --git a/RELEASE_READINESS_REPORT.md b/RELEASE_READINESS_REPORT.md new file mode 100644 index 0000000..9d555b4 --- /dev/null +++ b/RELEASE_READINESS_REPORT.md @@ -0,0 +1,231 @@ +# Release readiness report — earlyon 0.3 release candidate + +Date: 2026-07-22. + +## 1. Executive status + +**COMPLETE.** Every hardening claim from the handoff was independently +re-verified (and, where under-tested, adversarially tested); one packaging +defect and one evidence-runner defect were found and fixed; the package +builds, twine-checks, installs into a clean environment and passes an +out-of-repo smoke; v1 checkpoint migration is proven against a genuine +prior-version file; a bounded, seeded CUDA evidence run produced real +CIFAR-10 numbers with the negatives preserved. No remote state was touched; +no Jetson/TensorRT claims exist because no such hardware exists here. + +## 2. Branches / commits + +- Starting point: `claude/earlyon-hardening` @ `db3832e` (8 hardening + commits on `main` @ `a5f8414`), working tree clean, untouched. +- Work branch: `claude/earlyon-release-readiness` (created from hardening + HEAD; hardening commits preserved, not rebased/squashed). +- Final HEAD: the last commit listed in §16. + +## 3. Environment + +Python 3.13.9 · torch 2.10.0+cu128 · torchvision 0.25.0+cu128 · CUDA 12.8 · +NVIDIA RTX 4050 Laptop GPU (6 GB) · onnx 1.21.0 / onnxruntime 1.25.1 · +fvcore 0.1.5 · Linux 6.8. **Absent:** tegrastats, TensorRT, Jetson (all +verified absent; nothing was claimed from them). + +## 4. Claims independently verified + +All of the handoff's headline claims reproduced from scratch: + +- 261 tests passing at baseline (now 270), 95% coverage (badge said 96% — + corrected), strict mypy clean, ruff/black/isort clean for the CI scope. +- **Per-head temperatures (A):** independent per-head fit with convergence + status; deterministic head mapping; routing applies each head's own T on + the single-sample, batched (newly pinned) and staged paths; persisted in + v2 checkpoints; scalar v1 migrates deterministically; invalid values + raise / fall back safely. New decisive test: a single global temperature + demonstrably misroutes (accuracy 1.0 → <0.5) where per-head does not. +- **Explicit enablement (B):** disabled exits cannot fire at confidence + exactly 1.0 or entropy exactly 0.0 (both policies, both paths); all- + disabled routing reaches the final head; no sentinel thresholds remain; + migration preserves state. +- **Staged calibration (C):** collect/fit/search/evaluate are separable and + contracted; simulator pinned bit-exact against the real router; no test + leakage (split discipline documented and followed by the evidence run); + search is deterministic (newly pinned); edge cases covered: empty loader + raises, single-class data works, no-feasible-threshold → disabled, + NaN-emitting head safely disabled (new), NaN logits rejected in fitting, + non-converged fits fall back to 1.0 with status. +- **Benchmark fairness (D):** identical preloaded samples/order/warmup/ + sync/boundary proven by spy-module tests; batch-1 validated; model-only + vs end-to-end labelled; negative speedups preserved (ratio, no clamp); + full backbone + early-exit + static baseline compared in the evidence run. +- **Compute estimates (E):** `estimated_backbone_flops_fraction` naming, + `FlopsEstimate` provenance, exit-head exclusion disclosed, reused-module + detection → warned low-confidence fallback, lazy + cached analysis + (construction fast). Cache invalidation limitation documented (§17). +- **Checkpoint v2 (F):** format/library version, exit-point identity + cross-check, routing/enablement/temperatures/loss metadata, deterministic + v1 migration, future-version rejection, malformed-file rejection, + `weights_only=True` trust model documented in SECURITY_REVIEW.md. Newly + verified against a **genuine** prior-version checkpoint written by the + actual historical `save_wrapper` (committed fixture, ~1.2 MB, provenance + documented). +- **custom_ee (G):** CPU + CUDA (GPU test passes on this host), kwargs, + tuple and dict outputs via feature extractors (dict newly pinned), + missing/reused/out-of-order exits fail loudly; claims match tested scope + (single-tensor runtime contract stated in the docstring/README). +- **Jetson (H):** restartability, double-start no-op, thread/subprocess/pipe + cleanup, missing-tool degradation, conservative parsing (None ≠ 0), + trapezoidal integration with correct units, instantaneous vs average vs + integrated energy kept distinct — all covered by hardware-free tests. No + non-Jetson output is presented as Jetson evidence anywhere. +- **Staged deployment (I):** newly proven that later stages are *never + executed* when an exit fires (counting-module test); staged == eager + predictions/exit choices across policies, temperatures, enablement; + unsupported graphs rejected; all-exits ONNX vs conditional execution + distinction explicit in API name and docs. + +## 5. Defects discovered (beyond the handoff) + +1. **sdist shipped a partial test tree** (test modules without the fixtures + package) — broken if anyone ran tests from the sdist. +2. **The v1-migration fixture wasn't actually tracked** (`*.pth` gitignore), + so the committed migration test would fail on a fresh clone. +3. **Evidence runner OOM'd on 6 GB GPUs** (MobileNetV2 at batch 128 with the + wrapper still resident). +4. Lint/format debt in `scripts/` and `examples/` (outside CI's checked + scope): unused import, E402s, ambiguous lambda names, unformatted files. +5. README coverage badge overstated coverage (96% vs measured 95%). + +## 6. Fixes made + +(1) `MANIFEST.in` prunes tests from the sdist. (2) fixture force-added with +a provenance README. (3) baseline batch 32 + off-GPU staging between phases. +(4) fixed and brought into the format gate. (5) badge corrected. All +conservative; no behavior changes to library code were needed — the audit +found no functional defect in the hardening work itself. + +## 7. Tests, coverage, static checks (final) + +- `pytest` (full, incl. 6 CUDA tests): **270 passed**, ~13 s. +- Coverage: **95%** (`--cov=earlyon`). +- `mypy earlyon` (strict): clean, 28 files. +- `ruff check earlyon tests scripts examples`: clean. +- `black --check` / `isort --check-only` over the same scope: clean. + +## 8. Packaging + +`python -m build` → `earlyon-0.2.0-py3-none-any.whl` (73 KB) + +`earlyon-0.2.0.tar.gz`; both `twine check` PASSED; `pip check` clean. +Wheel installed into a fresh virtualenv and smoked **from outside the +repo** (`scripts/smoke_test.py`): import/version, model build, training + +routed inference, deterministic calibration, v2 round-trip, v1 migration, +staged==eager, benchmark smoke — all PASS; CLI entry point and all +subcommand `--help`s work from the installed console script. Caveat: the +clean venv used `--system-site-packages` for the torch stack (a from-scratch +torch download was out of scope); earlyon itself and its metadata were +installed with `--no-deps` from the wheel. + +## 9. v1 migration result + +PASS — against both hand-built payloads and the genuine historical fixture: +scalar temperature 1.6 broadcast to all four heads, sentinel threshold 1.0 +→ `enabled_exits=[True, False, True]`, warning emitted, model routes, +re-save produces a silent v2 round-trip. + +## 10. CUDA evidence + +Run completed (seed 42, 2+2+2 epochs, 10.4 min compute; full data in +`docs/evidence/cuda_evidence.json`, interpretation in +`docs/evidence/CUDA_EVIDENCE.md`): + +- ResNet-18 backbone 94.07% test acc, 1.33 ms p50, 712 ips. +- Early-exit ResNet-18 92.92% (−1.15% vs the 1% budget calibrated on the + calibration split — generalization gap reported), est. FLOPs fraction + 0.88, **1.10× throughput but worse median latency** (1.43 ms) with better + tails; 30.5% of images exited early. +- Static MobileNetV2 92.67%, 0.99× — competitive, as the docs predict at + tiny training budgets. +- Noise-input 2.05× recorded separately as a best-case bound. Peak CUDA + memory 350 MB during the timed comparison. No result was tuned on test + data; no negative was suppressed. + +## 11. Jetson / TensorRT status + +No hardware, no runtime, therefore **no claims**: monitor lifecycle, parser +and energy integration are verified hardware-free; the exact device +procedure is documented (`docs/STAGED_DEPLOYMENT.md` §Jetson, README +Reproducibility). Restartability on a real Jetson and any TensorRT staged +run remain external work. + +## 12. Security / privacy + +`SECURITY_REVIEW.md` created. Findings: no secrets/keys/tokens; no private +paths or employer identifiers; datasets and build artifacts gitignored; only +committed binary is the documented 1.2 MB untrained migration fixture; +`torch.load(weights_only=True)` + structural validation on every load path; +single subprocess (tegrastats) with fixed argv and cleanup; no eval/exec/ +shell=True; wheel/sdist contain no tests, data or weights. + +## 13. Public API changes (this branch) + +None. All library-code changes on this branch are tests, scripts, packaging +and docs. (API changes from the hardening branch are summarized in +`docs/MIGRATION.md` and `RELEASE_NOTES_DRAFT.md`.) + +## 14. Documentation changes + +README (architecture diagram, evidence-run results table with negative- +result interpretation, checkpoint-migration + reproducibility sections, +corrected badge), `docs/MIGRATION.md`, `SECURITY_REVIEW.md`, +`PR_BODY_DRAFT.md`, `docs/evidence/CUDA_EVIDENCE.md`, updated +`RELEASE_NOTES_DRAFT.md` and `CHANGELOG.md`, `tests/fixtures/README.md`. + +## 15. Files changed (this branch, 17 files) + +- tests: `tests/test_release_audit.py` (new, 10 tests), + `tests/fixtures/v1_cifar_resnet20.pth` + `README.md` (new) +- scripts: `smoke_test.py` (new), `evidence_run.py` (new), + `run_benchmarks.py` / `re_evaluate.py` / examples (lint/format only) +- packaging: `MANIFEST.in` (new) +- docs: README, CHANGELOG, RELEASE_NOTES_DRAFT, MIGRATION, SECURITY_REVIEW, + PR_BODY_DRAFT, docs/evidence/* (this report) + +## 16. Local commits (`claude/earlyon-hardening..HEAD`) + +1. `5d821aa` chore: bring scripts and examples up to the lint/format gate +2. `de8ce1c` test: independent release audit — adversarial verification +3. `1a0d90f` feat: packaging validation — sdist pruning, wheel smoke, v1 fixture, evidence runner +4. `e381db2` docs: security review, migration guide, PR draft, README architecture +5. `5ce0427` docs: coverage badge fix, changelog, release-notes audit section +6. `e38b073` feat: bounded CUDA evidence run — real seeded CIFAR-10 results +7. (final) docs: release readiness report + +## 17. Remaining external work + +- Jetson: run `earlyon profile` + the staged TensorRT procedure on a real + device before making any Jetson claim. +- Optional: full-length `scripts/run_benchmarks.py` (hours) to refresh + `docs/benchmarks.json` `runs` and retire the legacy README table. +- Torchvision-factory stage maps for `earlyon.staged` (protocol + acceptance + test exist). +- Known limitation (documented): `wrapper.flops_estimate` caches per + instance and does not invalidate if the backbone is structurally mutated + after first inference. +- Version bump to 0.3.0 + changelog dating at actual release time. + +## 18. Commands for the user (NOT executed) + +```bash +# inspect +git log --oneline main..claude/earlyon-release-readiness +git diff main...claude/earlyon-release-readiness + +# push + PR +git push -u origin claude/earlyon-release-readiness +gh pr create --base main --head claude/earlyon-release-readiness \ + --title "earlyon 0.3: correctness hardening + release readiness" \ + --body-file PR_BODY_DRAFT.md + +# after merge: version bump, tag, publish +# edit pyproject.toml + earlyon/__init__.py -> 0.3.0; date the CHANGELOG +git tag -a v0.3.0 -m "earlyon 0.3.0" +git push origin v0.3.0 +python -m build && python -m twine check dist/* && python -m twine upload dist/* +``` diff --git a/SECURITY_REVIEW.md b/SECURITY_REVIEW.md new file mode 100644 index 0000000..a11ab9f --- /dev/null +++ b/SECURITY_REVIEW.md @@ -0,0 +1,72 @@ +# Security review — earlyon 0.3 release candidate + +Date: 2026-07-22. Scope: the `earlyon` package, scripts, tests, docs and +packaged artifacts on branch `claude/earlyon-release-readiness`. + +## Checkpoint trust assumptions + +- `load_wrapper` calls `torch.load(..., map_location="cpu", weights_only=True)`. + With `weights_only=True`, PyTorch's restricted unpickler refuses arbitrary + object construction, so a malicious `.pth` cannot execute code through the + earlyon loading path. The payload is then structurally validated + (`format_version`, config shapes, exit-point cross-check) before any state + is applied. +- This protects the *loading* step only. A checkpoint still controls model + *behavior* (weights, thresholds, enablement, temperatures); loading an + untrusted checkpoint yields an untrusted model. Treat checkpoints like + data, not like code — but do not deploy models from untrusted sources. +- `save_wrapper` writes only tensors, plain Python scalars/containers and + strings; no custom classes are pickled. + +## Model-code execution assumptions + +- `custom_ee`, `EarlyExitWrapper` and `staged_model` execute the + **user-supplied backbone's `forward`** (dry runs, FLOPs probes, routing). + Wrapping a model is running its code. earlyon adds no sandboxing and makes + no claim to; only wrap modules you would run anyway. +- `load_wrapper(..., factory=...)` calls the user-supplied factory — same + trust level as the caller's own code. +- No use of `eval`, `exec`, dynamic imports of user strings, or `pickle` + outside `torch.save/load` as described above. + +## Dataset and weight provenance + +- Datasets: CIFAR-10 via `torchvision.datasets` (checksummed download by + torchvision) into a gitignored `data/` directory. No datasets are + committed or redistributed. +- Weights: ImageNet-pretrained weights come from torchvision's official hub + (HTTPS, hash-pinned filenames) when `pretrained=True`. earlyon + redistributes no third-party weights. The only committed binary is + `tests/fixtures/v1_cifar_resnet20.pth` (~1.2 MB): a *seeded, untrained* + model saved with the historical v1 format purely to test checkpoint + migration; it contains no trained knowledge and no external data. + +## Subprocess boundaries + +- The only subprocess is `tegrastats` (Jetson telemetry): fixed argv list, + no `shell=True`, no user-controlled arguments (the interval is a validated + integer), stdout parsed with anchored regexes, stderr discarded, process + terminated/killed with timeouts and pipes closed on stop. Spawn failure + degrades to "unavailable" rather than raising mid-benchmark. + +## Filesystem / temp / network + +- Library code performs no network access. Downloads happen only through + torchvision (datasets/weights) when the user opts in via CLI/scripts. +- Tests use `tmp_path`/`TemporaryDirectory`; no world-writable fixed paths. +- No archive extraction is performed by earlyon itself. + +## Scan results (this review) + +- No credentials, API keys, tokens or private keys in the tree. +- No private/internal paths, employer identifiers, or non-public material. + The author contact in `pyproject.toml` is deliberate public metadata. +- Coverage output, datasets, build artifacts (`htmlcov/`, `data/`, `dist/`, + `.coverage`) are gitignored and untracked. +- The wheel contains only the `earlyon` package; the sdist prunes `tests/` + (`MANIFEST.in`), so the binary fixture ships in neither artifact. + +## Reporting + +Vulnerabilities: follow `SECURITY.md` (private disclosure; do not open +public issues for security reports). diff --git a/docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md b/docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md new file mode 100644 index 0000000..8a05a2c --- /dev/null +++ b/docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md @@ -0,0 +1,93 @@ +# Calibration & benchmark contract + +The rules earlyon holds itself to when it fits a routing policy and when it +publishes a number. Everything here is enforced by code and pinned by tests. + +## Data splits + +| Split | Used for | Never used for | +|---|---|---| +| train | backbone + exit-head training | calibration, evaluation | +| temperature (held-out; optional) | per-head temperature fitting | threshold search | +| calibration / validation | threshold search, enablement, calibration-split metrics | reported accuracy claims | +| test | final reported accuracy + exit distribution (`earlyon.benchmarking.evaluate`) | anything the search touches | + +`calibrate_thresholds(fit_temperature=True)` without a separate +`temperature_loader` reuses the calibration split and **warns** — the +temperature fit leaks into the threshold accuracy estimate. All accuracies in +a `CalibrationResult` are calibration-split numbers, labelled as such. + +## Calibration pipeline (four separable stages) + +1. **Collect** — `collect_head_logits` runs the loader once through the + all-exits forward, caching raw logits for every head (exits + final). + Empty loaders raise immediately. +2. **Fit temperatures** — `fit_head_temperatures`: one temperature per head + (each head is a differently miscalibrated classifier). Fits carry a + convergence/fallback status; a diverged fit falls back to 1.0 (no + calibration) with a warning — never to an artificially sharp value. +3. **Search** — greedy coordinate descent over a fixed grid, simulated + against the cache. It is named `greedy-coordinate-grid` in results because + that is what it is; no global-optimality claim is made. A threshold that + fires no calibration sample leaves its exit **disabled** + (`enabled_exits[i] = False`) — an explicit boolean, not a sentinel, so a + float-saturated softmax (confidence exactly 1.0, entropy exactly 0.0) can + never fire a disabled exit. +4. **Evaluate** — the selected policy is scored on the calibration split + (accuracy, estimated compute, exit distribution) into `CalibrationResult`. + The simulator is pinned bit-for-bit against the real routing path by + `test_eval_cache_matches_real_router_exactly`. + +## Compute numbers are estimates + +`estimated_backbone_flops_fraction` (formerly `computation_used`, kept as a +read alias) comes from a static fvcore analysis of the backbone: + +- exit-head FLOPs and all routing overhead are **excluded**; +- the per-layer attribution assumes each leaf module runs exactly once, in + registration order (true for the shipped torchvision backbones). A reused + module is detected and degrades the estimate to a warned, low-confidence + uniform fallback (`FlopsEstimate.reliable = False`); +- the analysis is lazy (first inference), cached, and its provenance is on + `wrapper.flops_estimate`. + +A FLOPs saving is **not** a latency claim. The eager router evaluates an +exit head and synchronises the host (`.item()`) at every enabled exit; on a +GPU that synchronisation can cost more than the skipped layers save, +especially on small backbones. That is why benchmark results and the README +report wall-clock and estimated-FLOPs numbers separately. + +## Benchmark fairness rules (`benchmark_models`) + +1. Every compared model sees the **exact same preloaded sample sequence** — + same order, same tensors — with the same warmup count. +2. Boundaries are identical and labelled: `model-only` (input staged on the + device before timing) or `end-to-end` (H2D copy inside the timed region). + Loader/preprocessing time is never silently included. +3. Same `eval()` + `torch.inference_mode()` + per-iteration CUDA + synchronisation for every model. +4. Batch 1 is the primary supported mode (the latency-sensitive edge case + earlyon targets) and is validated. +5. Reported: p50/p95/p99 latency, throughput, run count, device, dtype, + input shape, warmup, sync policy, exit distribution, input source + (`loader` vs `random-noise`), and accuracy when the loader is labelled. +6. Random-noise input is labelled a best-case bound: trained heads can fire + spuriously on noise. +7. A speedup is only quoted between results produced by the same + `benchmark_models` call. Pre-v0.3 records in `docs/benchmarks.json` live + under `legacy_v0_2` and are marked not comparable (the old runner fed the + wrapper and backbone different inputs). +8. An early-exit model should also be compared against a **smaller static + model** at matched accuracy — if a ResNet18 matches your routed ResNet50's + accuracy and latency, the honest deployment is the ResNet18. Pass it as a + third entry to `benchmark_models`. + +## Serialization + +Checkpoints are `format_version: 2`: routing policy, both threshold lists, +`enabled_exits`, per-head `temperatures`, loss weights, exit-point metadata +(cross-checked against the rebuilt factory), library version. Unversioned v1 +files migrate deterministically (scalar temperature broadcast; sentinel +thresholds → explicit disabled) with warnings. `custom_ee` models reload only +via a user-supplied `factory=` — an arbitrary backbone is never silently +"reconstructed". diff --git a/docs/DESIGN_DECISIONS.md b/docs/DESIGN_DECISIONS.md index 2ad8667..8ea0580 100644 --- a/docs/DESIGN_DECISIONS.md +++ b/docs/DESIGN_DECISIONS.md @@ -78,19 +78,61 @@ returns `budget_met=False` with the least-compute configuration found, and because it is measured on the calibration set, the guarantee is average-case over that distribution. The README's limitations section says both out loud. -## Temperature scaling before calibration - -**Decision:** optionally fit a single scalar temperature (Guo et al. 2017) on -held-out data before threshold calibration, per exit head. - -**Why:** modern CNN heads are systematically over-confident; softmax maxima -cluster near 1.0 whether the prediction is right or wrong. Thresholding an -uncalibrated confidence wastes most of the [0, 1] range and makes the greedy -sweep's grid resolution meaningless. One scalar per head fixes the ranking -cheaply without changing argmax predictions, so accuracy is untouched by -construction. We chose temperature scaling over Platt/isotonic because it's -the simplest method with a strong published track record on exactly this -failure mode. +## Temperature scaling before calibration — one temperature per head + +**Decision:** optionally fit a temperature (Guo et al. 2017) on held-out data +before threshold calibration — one **per head**: every exit head and the +final classifier get their own fitted scalar +(`EarlyExitConfig.temperatures`). + +**Why per-head:** each exit head is a different classifier reading different +features, with its own miscalibration; a shallow head is typically far more +over-confident than the final classifier. Earlier versions fitted a single +temperature from the final classifier's logits and reused it for every exit +— defensible for none of them. The fit (`fit_head_temperatures`) runs on +cached logits from one pass, and each `TemperatureFitResult` carries a +convergence/fallback status: a diverged LBFGS falls back to 1.0 (no +calibration) with a warning, never to an artificially sharp value. + +**Why temperature scaling at all:** modern CNN heads are systematically +over-confident; softmax maxima cluster near 1.0 whether the prediction is +right or wrong. Thresholding an uncalibrated confidence wastes most of the +[0, 1] range and makes the greedy sweep's grid resolution meaningless. One +scalar per head fixes the ranking cheaply without changing argmax +predictions, so accuracy is untouched by construction. We chose temperature +scaling over Platt/isotonic because it's the simplest method with a strong +published track record on exactly this failure mode. + +## Explicit exit enablement instead of threshold sentinels + +**Decision:** whether an exit may fire is a boolean +(`EarlyExitConfig.enabled_exits`), not a magic threshold value. + +**Why:** the old convention — confidence threshold `1.0` (or entropy `0.0`) +means "disabled" — was numerically unsound: the router fires on +`confidence >= threshold`, and a float32-saturated softmax produces +confidence exactly 1.0, so a "disabled" exit could still fire. Explicit +booleans make disablement absolute (the head isn't even evaluated at +inference, saving its overhead too), keep thresholds meaning only +thresholds, and let calibration state its decision (`enabled_exits` on the +result) instead of encoding it. v1 checkpoints migrate their sentinels to +explicit flags on load, with a warning. Disabled heads still produce logits +in training mode — enablement is a routing concept, not a training one. + +## `estimated_backbone_flops_fraction`, not "computation used" + +**Decision:** the compute number every result reports is named as what it is +— a static *estimate* of the backbone FLOPs fraction — with the estimator's +provenance attached (`wrapper.flops_estimate`: method, reliability, notes). + +**Why:** the old name `computation_used` read like a measurement of the +inference that just ran. It never was: it is a one-time fvcore analysis of +the backbone, excludes the exit heads and all routing overhead, and assumes +each leaf module runs once in registration order. Backbones that reuse +modules break that assumption — the estimator now detects reuse and degrades +to a warned, low-confidence uniform fallback instead of a precise-looking +wrong number. The analysis is lazy (first inference) so constructing large +models (ViT) stays fast. `computation_used` survives as a read alias. ## Entropy routing alongside confidence diff --git a/docs/MIGRATION.md b/docs/MIGRATION.md new file mode 100644 index 0000000..16bda77 --- /dev/null +++ b/docs/MIGRATION.md @@ -0,0 +1,55 @@ +# Migrating to earlyon 0.3 + +0.3 changes calibration/routing semantics. Old checkpoints load and migrate +automatically; code that touched the changed fields needs the edits below. + +## Checkpoints (v1 → v2): automatic + +`load_wrapper` detects unversioned (≤0.2) files and migrates them +deterministically, emitting a `UserWarning` describing what changed: + +| v1 field | v2 result | +|---|---| +| scalar `temperature: T` | `temperatures = {every exit: T, "final": T}` | +| confidence threshold `1.0` (active policy) | that exit's `enabled_exits[i] = False` | +| entropy threshold `0.0` (active policy) | that exit's `enabled_exits[i] = False` | +| non-finite/non-positive `temperature` | `1.0` (uncalibrated) + warning | + +Behavioral note: under v1, a "disabled" (threshold-1.0) exit could still fire +when float32 softmax saturated at exactly 1.0. After migration it is strictly +disabled. If you relied on that edge case, re-enable the exit and set a real +threshold. + +Re-saving a migrated model writes `format_version: 2`; v2 files round-trip +with no warnings. Files with a *newer* format version are rejected with an +upgrade message. This is pinned against a genuine v0.2-written fixture in +`tests/test_release_audit.py`. + +`custom_ee` checkpoints cannot be rebuilt from a string — pass a factory: + +```python +model = load_wrapper("ckpt.pth", factory=lambda: custom_ee(build_backbone(), ...)) +``` + +## Code changes + +- **Temperature:** read/write `config.temperatures` (a dict keyed by exit + name plus `"final"`), not `config.temperature`. The scalar is still + accepted by the constructor (broadcast once); mutating it after + construction has no effect on routing. +- **Disabling an exit:** set `config.enabled_exits[i] = False`. Do not use + threshold sentinels; `validate()` requires thresholds in-range. +- **Compute number:** `result.computation_used` still reads, but the field + is `estimated_backbone_flops_fraction` — and it is an estimate (see + `docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md`). +- **CalibrationResult:** `fitted_temperature` is now the final head's fit + (deprecated); use `result.temperatures` / `result.temperature_fits`. + New fields: `enabled_exits`, `exit_distribution`, `num_samples`, + `objective`, `method`, `accuracy_delta`, `schema_version`. +- **Benchmarking:** prefer `benchmark_models({...}, loader=...)`; the old + single-model helpers still work and now share the fair measurement core. + `BenchmarkResult.avg_computation_used` → `avg_estimated_flops_fraction` + (alias kept). +- **Invalid config now raises** (e.g. negative temperature, out-of-range + threshold) from `EarlyExitConfig.validate()` instead of being silently + clamped at inference time. diff --git a/docs/STAGED_DEPLOYMENT.md b/docs/STAGED_DEPLOYMENT.md new file mode 100644 index 0000000..81ac28d --- /dev/null +++ b/docs/STAGED_DEPLOYMENT.md @@ -0,0 +1,80 @@ +# Staged deployment + +How to get a *real* early-exit compute saving outside of eager PyTorch. + +## The three execution modes, honestly labelled + +| Mode | Skips later layers? | Exportable? | Where | +|---|---|---|---| +| Eager wrapper (`model(x, mode="inference")`) | **yes** — hooks + a sentinel exception short-circuit the backbone | no (`torch.compile`/export refuse with a clear error) | `earlyon.core.wrappers` | +| Static multi-output ONNX (`earlyon export`, `export_to_onnx`) | **no** — the graph always computes every exit; the caller applies routing to the outputs | yes, single portable graph | `earlyon.onnx` | +| Staged deployment | **yes** — stages are separate modules; routing runs between them, later stages are never invoked | yes, one artifact per stage | `earlyon.staged` | + +## The protocol + +An early-exit model with `k` exits becomes `k + 1` ordered **stages**. + +- Stage `i` (`0 <= i < k`) consumes the previous stage's **continuation + features** (stage 0 consumes the input) and produces + `(continuation_features, exit_logits_i)` where `exit_logits_i` comes from + exit head `i`. +- Stage `k` (final) produces the final classifier's logits. +- Between stages the runtime applies exactly earlyon's routing rule, in this + order: if `enabled_exits[i]` is false, skip the head and continue; else + compute `softmax(logits / temperatures[exit_name])` and fire on + `confidence >= confidence_thresholds[i]` (or entropy `<= threshold` for the + entropy policy). The first firing exit's logits are the prediction; later + stages are not executed. + +The data classes live in `earlyon.staged`: `StageSpec` (index, exit name, +covered modules), `Stage` (trunk + head, `forward -> (features, logits)`), +and `StagedModel` (the reference runtime, `infer(x) -> InferenceResult`). + +## Reference implementation and its scope + +`earlyon.staged.staged_model(wrapper)` splits a wrapper whose backbone is +literally an `nn.Sequential` (exit layers at top level, `final_classifier` +identity — what `custom_ee` produces for a Sequential backbone). It refuses +anything else: arbitrary graph partitioning is out of scope, and a wrong +split is worse than no split. As a belt, the builder runs a probe input +through both the staged and eager paths and raises on any disagreement; +`tests/test_staged.py` pins equivalence across policies, per-head +temperatures and enablement. + +```python +from earlyon.models import custom_ee +from earlyon.staged import staged_model + +wrapper = custom_ee(sequential_backbone, exit_layers=["2", "5"], num_classes=10) +staged = staged_model(wrapper) # raises if the split can't be proven +result = staged.infer(x) # same InferenceResult as the wrapper +``` + +For the torchvision factories (ResNet, MobileNet, ViT), the backbone forward +is not a plain Sequential; implementing their stage splits is future work — +the protocol above is what such an implementation must satisfy, and the +equivalence test is the acceptance criterion. + +## Exporting stages + +Each `Stage` is a plain `nn.Module` with tensor-in/tensors-out, so it traces +cleanly with `torch.onnx.export` (outputs: `features`, `logits`). Your +deployment runtime (ONNX Runtime sessions, or TensorRT engines built from the +per-stage ONNX files) then implements the routing loop above on the host — +one engine invocation per stage until an exit fires. + +## Jetson / TensorRT procedure (not run here — no fabricated numbers) + +On the target device: + +1. Export per-stage ONNX files from a calibrated checkpoint. +2. `trtexec --onnx=stage_k.onnx --saveEngine=stage_k.plan --fp16` per stage. +3. Runtime loop: run stage 0's engine, apply the threshold on the host + (logits are tiny — the copy is negligible), continue only if no fire. +4. Benchmark with the same discipline as `earlyon.benchmarking`: fixed + sample set, ≥50-iteration warmup, batch 1, report p50/p95/p99 and the + exit distribution; compare against a single-engine full-backbone build + *and* a smaller static model at matched accuracy. + +Nothing in this repository claims TensorRT numbers; the section above is the +procedure, to be filled in from a real device run. diff --git a/docs/benchmarks.json b/docs/benchmarks.json index d8786fd..eb0db8c 100644 --- a/docs/benchmarks.json +++ b/docs/benchmarks.json @@ -1,94 +1,99 @@ { - "resnet18": { - "backbone": "resnet18", - "dataset": "cifar10", - "device": "NVIDIA GeForce RTX 4050 Laptop GPU", - "stage1_seconds": 449.7, - "stage2_seconds": 147.2, - "calibrate_seconds": 183.9, - "thresholds": [ - 0.7, - 0.85 - ], - "baseline_accuracy_val": 0.9632, - "test_accuracy": 0.9442, - "test_avg_computation_used": 0.8988, - "test_exit_distribution": { - "final": 0.6471, - "exit_1": 0.2584, - "exit_0": 0.0945 + "_schema": 2, + "_note": "Records under 'legacy_v0_2' were measured with the pre-v0.3 benchmark runner, which fed the wrapper and the raw backbone DIFFERENT inputs and used non-identical measurement boundaries; they are not methodologically comparable to 'runs' records produced by the v0.3 fair runner (benchmark_models: identical preloaded samples, identical warmup/sync/boundary for every compared model). Regenerate with scripts/run_benchmarks.py or scripts/re_evaluate.py.", + "legacy_v0_2": { + "resnet18": { + "backbone": "resnet18", + "dataset": "cifar10", + "device": "NVIDIA GeForce RTX 4050 Laptop GPU", + "stage1_seconds": 449.7, + "stage2_seconds": 147.2, + "calibrate_seconds": 183.9, + "thresholds": [ + 0.7, + 0.85 + ], + "baseline_accuracy_val": 0.9632, + "test_accuracy": 0.9442, + "test_avg_computation_used": 0.8988, + "test_exit_distribution": { + "final": 0.6471, + "exit_1": 0.2584, + "exit_0": 0.0945 + }, + "throughput_backbone_ips": 758.4, + "throughput_wrapper_ips": 1906.6, + "speedup": 2.575, + "latency_backbone_p50_ms": 1.316, + "latency_wrapper_p50_ms": 0.51, + "latency_wrapper_p95_ms": 0.547, + "throughput_wrapper_ips_noise_input": 1853.0, + "throughput_speedup_noise_input": 2.443, + "latency_wrapper_p50_ms_noise_input": 0.508, + "note": "throughput uses random noise input which may trigger spurious exits in trained heads; test_avg_computation_used is the honest signal" }, - "throughput_backbone_ips": 758.4, - "throughput_wrapper_ips": 1906.6, - "speedup": 2.575, - "latency_backbone_p50_ms": 1.316, - "latency_wrapper_p50_ms": 0.51, - "latency_wrapper_p95_ms": 0.547, - "throughput_wrapper_ips_noise_input": 1853.0, - "throughput_speedup_noise_input": 2.443, - "latency_wrapper_p50_ms_noise_input": 0.508, - "note": "throughput uses random noise input which may trigger spurious exits in trained heads; test_avg_computation_used is the honest signal" - }, - "resnet50": { - "backbone": "resnet50", - "dataset": "cifar10", - "device": "NVIDIA GeForce RTX 4050 Laptop GPU", - "stage1_seconds": 1009.3, - "stage2_seconds": 345.7, - "calibrate_seconds": 413.8, - "thresholds": [ - 0.7, - 0.85, - 0.99 - ], - "baseline_accuracy_val": 0.976, - "test_accuracy": 0.9588, - "test_avg_computation_used": 0.8142, - "test_exit_distribution": { - "final": 0.4177, - "exit_2": 0.4247, - "exit_0": 0.0579, - "exit_1": 0.0997 + "resnet50": { + "backbone": "resnet50", + "dataset": "cifar10", + "device": "NVIDIA GeForce RTX 4050 Laptop GPU", + "stage1_seconds": 1009.3, + "stage2_seconds": 345.7, + "calibrate_seconds": 413.8, + "thresholds": [ + 0.7, + 0.85, + 0.99 + ], + "baseline_accuracy_val": 0.976, + "test_accuracy": 0.9588, + "test_avg_computation_used": 0.8142, + "test_exit_distribution": { + "final": 0.4177, + "exit_2": 0.4247, + "exit_0": 0.0579, + "exit_1": 0.0997 + }, + "throughput_backbone_ips": 355.1, + "throughput_wrapper_ips": 329.4, + "speedup": 0.938, + "latency_backbone_p50_ms": 2.814, + "latency_wrapper_p50_ms": 3.025, + "latency_wrapper_p95_ms": 3.082, + "throughput_wrapper_ips_noise_input": 1942.9, + "throughput_speedup_noise_input": 5.472, + "latency_wrapper_p50_ms_noise_input": 0.509, + "note": "throughput uses random noise input which may trigger spurious exits in trained heads; test_avg_computation_used is the honest signal" }, - "throughput_backbone_ips": 355.1, - "throughput_wrapper_ips": 329.4, - "speedup": 0.938, - "latency_backbone_p50_ms": 2.814, - "latency_wrapper_p50_ms": 3.025, - "latency_wrapper_p95_ms": 3.082, - "throughput_wrapper_ips_noise_input": 1942.9, - "throughput_speedup_noise_input": 5.472, - "latency_wrapper_p50_ms_noise_input": 0.509, - "note": "throughput uses random noise input which may trigger spurious exits in trained heads; test_avg_computation_used is the honest signal" + "mobilenetv2": { + "backbone": "mobilenetv2", + "dataset": "cifar10", + "device": "NVIDIA GeForce RTX 4050 Laptop GPU", + "stage1_seconds": 577.9, + "stage2_seconds": 165.8, + "calibrate_seconds": 225.5, + "thresholds": [ + 0.75, + 0.99 + ], + "baseline_accuracy_val": 0.9508, + "test_accuracy": 0.9331, + "test_avg_computation_used": 0.939, + "test_exit_distribution": { + "final": 0.9155, + "exit_0": 0.0801, + "exit_1": 0.0044 + }, + "throughput_backbone_ips": 751.4, + "throughput_wrapper_ips": 2866.5, + "speedup": 3.809, + "latency_backbone_p50_ms": 1.323, + "latency_wrapper_p50_ms": 0.337, + "latency_wrapper_p95_ms": 0.414, + "throughput_wrapper_ips_noise_input": 2685.9, + "throughput_speedup_noise_input": 3.574, + "latency_wrapper_p50_ms_noise_input": 0.365, + "note": "throughput uses random noise input which may trigger spurious exits in trained heads; test_avg_computation_used is the honest signal" + } }, - "mobilenetv2": { - "backbone": "mobilenetv2", - "dataset": "cifar10", - "device": "NVIDIA GeForce RTX 4050 Laptop GPU", - "stage1_seconds": 577.9, - "stage2_seconds": 165.8, - "calibrate_seconds": 225.5, - "thresholds": [ - 0.75, - 0.99 - ], - "baseline_accuracy_val": 0.9508, - "test_accuracy": 0.9331, - "test_avg_computation_used": 0.939, - "test_exit_distribution": { - "final": 0.9155, - "exit_0": 0.0801, - "exit_1": 0.0044 - }, - "throughput_backbone_ips": 751.4, - "throughput_wrapper_ips": 2866.5, - "speedup": 3.809, - "latency_backbone_p50_ms": 1.323, - "latency_wrapper_p50_ms": 0.337, - "latency_wrapper_p95_ms": 0.414, - "throughput_wrapper_ips_noise_input": 2685.9, - "throughput_speedup_noise_input": 3.574, - "latency_wrapper_p50_ms_noise_input": 0.365, - "note": "throughput uses random noise input which may trigger spurious exits in trained heads; test_avg_computation_used is the honest signal" - } + "runs": {} } \ No newline at end of file diff --git a/docs/evidence/CUDA_EVIDENCE.md b/docs/evidence/CUDA_EVIDENCE.md new file mode 100644 index 0000000..4dfa315 --- /dev/null +++ b/docs/evidence/CUDA_EVIDENCE.md @@ -0,0 +1,81 @@ +# Bounded CUDA evidence run — 2026-07-22 + +Machine-readable source: [`cuda_evidence.json`](cuda_evidence.json). +Reproduce: `python scripts/evidence_run.py` (seed 42; ~10 min compute on the +hardware below after the CIFAR-10 download). + +**Purpose.** A bounded, seeded, end-to-end *system validation* — training → +per-head temperature fit → threshold calibration → held-out evaluation → +fair 3-model benchmark — under the v0.3 methodology. It is deliberately +under-trained (2 epochs per stage) and is **not** a tuned headline +benchmark. Negative and mixed results below are reported as-is. + +## Setup + +| | | +|---|---| +| Hardware | NVIDIA RTX 4050 Laptop GPU (6 GB), CUDA 12.8 | +| Software | Python 3.13.9, torch 2.10.0+cu128, torchvision 0.25.0, earlyon 0.2.0 (0.3 RC branch) | +| Dataset | CIFAR-10 (torchvision download), images upsampled to 224×224 | +| Splits | train 45,000 · temperature 2,500 · calibration 2,500 · test 10,000 (official) — disjoint; the threshold search never saw test labels | +| Model | `resnet18_ee` (exits after `layer2`, `layer3`; heads: hidden 128, dropout 0.2), ImageNet-pretrained, fine-tuned 2+2 epochs (seed 42) | +| Static baseline | `mobilenet_v2` backbone, ImageNet-pretrained, fine-tuned 2 epochs (batch 32 for memory) | +| Determinism | seed fixes init/data order; bitwise GPU determinism not enforced (cuDNN) | + +## Calibration (never touches test data) + +Per-head fitted temperatures: e0 = 0.730, e1 = 0.641, final = 0.888 (all +fits converged) — three genuinely different values, which is the point of +per-head fitting. Greedy grid (accuracy budget 1%) selected thresholds +[0.7, 0.9] with both exits enabled. Calibration-split accuracy: 94.44% +(final head only) → 93.60% routed (−0.84%, inside the 1% budget on that +split). Calibration took 4.6 s (one cached-logits pass). + +## Held-out test results (10,000 images) + +| Model | Test accuracy | Est. backbone FLOPs fraction | +|---|---:|---:| +| ResNet-18 backbone (full) | **94.07%** | 1.000 | +| ResNet-18 early-exit (routed) | 92.92% | 0.880 (estimate; excludes heads/routing) | +| MobileNetV2 static baseline | 92.67% | 1.000 (of its own, smaller network) | + +Exit distribution: 5.6% at exit 0, 24.9% at exit 1, 69.5% ran the full +network. The routed accuracy drop on *test* is 1.15% — slightly beyond the +1% budget calibrated on the calibration split. That gap is expected +generalization error of the calibration estimate and is reported, not +hidden. + +## Fair benchmark (identical real test samples, model-only boundary, batch 1, 300 runs, 50 warmup, per-iteration CUDA sync) + +| Model | p50 | p95 | p99 | Throughput | Accuracy on timed samples | +|---|---:|---:|---:|---:|---:| +| ResNet-18 backbone | 1.33 ms | 1.51 ms | 1.65 ms | 711.9 ips | 93.7% | +| ResNet-18 early-exit | 1.43 ms | 1.46 ms | 1.48 ms | 781.8 ips | 91.7% | +| MobileNetV2 static | 1.40 ms | 1.48 ms | 1.50 ms | 702.1 ips | 93.0% | + +Speedups (same samples, same boundary): early-exit vs backbone **1.10×** +throughput; static baseline vs backbone 0.99×. Peak CUDA memory during the +timed comparison: 350 MB. Noise-input variant (best-case bound, spurious +exits possible): 2.05× — reported separately in the JSON and *not* a +practical-speed claim. + +## Honest interpretation + +- The early-exit model's **median latency is worse** than the raw backbone + (1.43 vs 1.33 ms) even though throughput is 10% better: when no exit + fires (70% of images) the wrapper pays the full backbone *plus* two head + evaluations and host syncs. The win comes from the 30% of images that + exit early, which also compresses the tail (p95/p99 are better). This is + exactly the routing-overhead behavior the documentation warns about. +- A ~12% estimated FLOPs saving translated to a ~10% throughput gain on + this GPU — estimates and wall-clock are close here, but that is not a + general law. +- At this (deliberately tiny) training budget the static MobileNetV2 is a + serious competitor: 92.67% accuracy at 0.99× speed vs the early-exit's + 92.92% at 1.10×. The early-exit model wins on both axes, but narrowly — + with longer training or a different backbone the comparison could flip, + which is why the fair runner makes the static baseline a first-class + citizen. +- These numbers characterize *this* bounded run on *this* laptop GPU. They + are not Jetson numbers, not ImageNet numbers, and not a general speedup + claim. diff --git a/docs/evidence/cuda_evidence.json b/docs/evidence/cuda_evidence.json new file mode 100644 index 0000000..1f8c69e --- /dev/null +++ b/docs/evidence/cuda_evidence.json @@ -0,0 +1,263 @@ +{ + "purpose": "bounded system-validation run; not a tuned headline benchmark", + "seed": 42, + "determinism_note": "seed fixes init/data order; bitwise GPU determinism not enforced (cuDNN)", + "dataset": { + "name": "CIFAR-10 (torchvision, images upsampled to 224)", + "train": 45000, + "temperature": 2500, + "calibration": 2500, + "test": 10000 + }, + "weights": "ImageNet-pretrained torchvision weights, fine-tuned as recorded below", + "epochs": { + "stage1_backbone": 2, + "stage2_exits": 2, + "static_baseline": 2 + }, + "model": { + "backbone": "resnet18", + "exit_points": [ + { + "name": "e0", + "layer": "layer2", + "in_channels": 128 + }, + { + "name": "e1", + "layer": "layer3", + "in_channels": 256 + } + ], + "head_config": "EarlyExitHead(hidden_dim=128, dropout=0.2)" + }, + "calibration": { + "objective": "accuracy_budget", + "target_accuracy_drop": 0.01, + "method": "greedy-coordinate-grid", + "thresholds": [ + 0.7, + 0.9 + ], + "enabled_exits": [ + true, + true + ], + "temperatures": { + "e0": 0.7298606743211031, + "e1": 0.6408520879825857, + "final": 0.8881403542803983 + }, + "temperature_fit_converged": { + "e0": true, + "e1": true, + "final": true + }, + "calibration_split_baseline_acc": 0.9444, + "calibration_split_routed_acc": 0.936 + }, + "test": { + "early_exit_routed_accuracy": 0.9292, + "backbone_accuracy": 0.9407, + "static_mobilenetv2_accuracy": 0.9267, + "early_exit_estimated_flops_fraction": 0.8799, + "exit_distribution": { + "final": 0.6953, + "exit_1": 0.2491, + "exit_0": 0.0556 + } + }, + "benchmark_real_input": { + "backbone_resnet18": { + "throughput_ips": 711.8584093111814, + "latency_median_ms": 1.328067499343888, + "latency_p50_ms": 1.3283369989949279, + "latency_p95_ms": 1.5121870001166826, + "latency_p99_ms": 1.6473269988637185, + "avg_estimated_flops_fraction": 1.0, + "exit_distribution": { + "final": 1.0 + }, + "num_runs": 300, + "device": "cuda", + "num_warmup": 50, + "batch_size": 1, + "input_shape": [ + 1, + 3, + 224, + 224 + ], + "dtype": "torch.float32", + "boundary": "model-only", + "sync": "cuda-synchronize-per-iteration", + "input_source": "loader", + "accuracy": 0.9366666666666666, + "num_labelled_samples": 300 + }, + "early_exit_resnet18": { + "throughput_ips": 781.8321391791311, + "latency_median_ms": 1.4273374999902444, + "latency_p50_ms": 1.4274259992816951, + "latency_p95_ms": 1.4558230013790308, + "latency_p99_ms": 1.4794830003665993, + "avg_estimated_flops_fraction": 0.8855555555555557, + "exit_distribution": { + "final": 0.73, + "exit_1": 0.19666666666666666, + "exit_0": 0.07333333333333333 + }, + "num_runs": 300, + "device": "cuda", + "num_warmup": 50, + "batch_size": 1, + "input_shape": [ + 1, + 3, + 224, + 224 + ], + "dtype": "torch.float32", + "boundary": "model-only", + "sync": "cuda-synchronize-per-iteration", + "input_source": "loader", + "accuracy": 0.9166666666666666, + "num_labelled_samples": 300 + }, + "static_mobilenetv2": { + "throughput_ips": 702.0757905132444, + "latency_median_ms": 1.398074999997334, + "latency_p50_ms": 1.3982290001877118, + "latency_p95_ms": 1.4844550005364, + "latency_p99_ms": 1.5049649991851766, + "avg_estimated_flops_fraction": 1.0, + "exit_distribution": { + "final": 1.0 + }, + "num_runs": 300, + "device": "cuda", + "num_warmup": 50, + "batch_size": 1, + "input_shape": [ + 1, + 3, + 224, + 224 + ], + "dtype": "torch.float32", + "boundary": "model-only", + "sync": "cuda-synchronize-per-iteration", + "input_source": "loader", + "accuracy": 0.93, + "num_labelled_samples": 300 + } + }, + "benchmark_real_speedups": { + "early_exit_vs_backbone": 1.098, + "static_vs_backbone": 0.986 + }, + "benchmark_noise_input": { + "note": "best-case bound: trained heads may fire spuriously on noise", + "backbone_resnet18": { + "throughput_ips": 755.3621345047061, + "latency_median_ms": 1.3031150001552305, + "latency_p50_ms": 1.3031210000917781, + "latency_p95_ms": 1.3612429993372643, + "latency_p99_ms": 1.9350050006323727, + "avg_estimated_flops_fraction": 1.0, + "exit_distribution": { + "final": 1.0 + }, + "num_runs": 300, + "device": "cuda", + "num_warmup": 50, + "batch_size": 1, + "input_shape": [ + 1, + 3, + 224, + 224 + ], + "dtype": "torch.float32", + "boundary": "model-only", + "sync": "cuda-synchronize-per-iteration", + "input_source": "random-noise", + "accuracy": null, + "num_labelled_samples": 0 + }, + "early_exit_resnet18": { + "throughput_ips": 1551.8133886081237, + "latency_median_ms": 0.6742770001437748, + "latency_p50_ms": 0.6742940004187403, + "latency_p95_ms": 0.7781620006426238, + "latency_p99_ms": 0.7984900003066286, + "avg_estimated_flops_fraction": 0.3333333333333333, + "exit_distribution": { + "exit_0": 1.0 + }, + "num_runs": 300, + "device": "cuda", + "num_warmup": 50, + "batch_size": 1, + "input_shape": [ + 1, + 3, + 224, + 224 + ], + "dtype": "torch.float32", + "boundary": "model-only", + "sync": "cuda-synchronize-per-iteration", + "input_source": "random-noise", + "accuracy": null, + "num_labelled_samples": 0 + }, + "static_mobilenetv2": { + "throughput_ips": 681.161147706457, + "latency_median_ms": 1.4052340002308483, + "latency_p50_ms": 1.4053710001462605, + "latency_p95_ms": 1.5782589998707408, + "latency_p99_ms": 3.479250999589567, + "avg_estimated_flops_fraction": 1.0, + "exit_distribution": { + "final": 1.0 + }, + "num_runs": 300, + "device": "cuda", + "num_warmup": 50, + "batch_size": 1, + "input_shape": [ + 1, + 3, + 224, + 224 + ], + "dtype": "torch.float32", + "boundary": "model-only", + "sync": "cuda-synchronize-per-iteration", + "input_source": "random-noise", + "accuracy": null, + "num_labelled_samples": 0 + }, + "early_exit_vs_backbone_speedup": 2.054 + }, + "peak_cuda_memory_mb_during_real_benchmark": 350.4, + "stage_seconds": { + "stage1_seconds": 220.7, + "stage2_seconds": 71.9, + "baseline_seconds": 284.4, + "baseline_batch_size": 32, + "calibrate_seconds": 4.6, + "test_eval_seconds": 16.8 + }, + "total_wall_seconds": 622.7, + "environment": { + "python": "3.13.9", + "torch": "2.10.0+cu128", + "torchvision": "0.25.0+cu128", + "cuda": "12.8", + "gpu": "NVIDIA GeForce RTX 4050 Laptop GPU", + "earlyon": "0.2.0", + "os": "Linux-6.8.0-106-generic-x86_64-with-glibc2.35" + } +} \ No newline at end of file diff --git a/earlyon/__init__.py b/earlyon/__init__.py index 737074b..7ced7a8 100644 --- a/earlyon/__init__.py +++ b/earlyon/__init__.py @@ -4,7 +4,7 @@ from earlyon.core.types import BatchedInferenceResult, EarlyExitConfig, ExitPoint, InferenceResult from earlyon.core.wrappers import EarlyExitWrapper -__version__ = "0.2.0" +__version__ = "0.3.0" __all__ = [ "EarlyExitWrapper", diff --git a/earlyon/benchmarking/__init__.py b/earlyon/benchmarking/__init__.py index 3bc1de4..f0c751c 100644 --- a/earlyon/benchmarking/__init__.py +++ b/earlyon/benchmarking/__init__.py @@ -2,13 +2,17 @@ from earlyon.benchmarking.jetson_profiler import JetsonProfiler, JetsonRun, JetsonSample from earlyon.benchmarking.throughput import ( BenchmarkResult, + ComparisonResult, benchmark_backbone, + benchmark_models, benchmark_wrapper, benchmark_wrapper_on_loader, ) __all__ = [ "BenchmarkResult", + "ComparisonResult", + "benchmark_models", "benchmark_wrapper", "benchmark_wrapper_on_loader", "benchmark_backbone", diff --git a/earlyon/benchmarking/jetson_profiler.py b/earlyon/benchmarking/jetson_profiler.py index 998170b..eaedbcb 100644 --- a/earlyon/benchmarking/jetson_profiler.py +++ b/earlyon/benchmarking/jetson_profiler.py @@ -1,9 +1,20 @@ """NVIDIA Jetson power/temperature profiling via tegrastats. tegrastats blocks forever without an end condition, so we spawn it as a -background subprocess, read lines via threads, and tear it down explicitly. -On non-Jetson systems all profile calls return zeros so the rest of the -benchmark suite still runs. +background subprocess, read lines via a thread, and tear it down explicitly. +On non-Jetson systems all profile calls return empty telemetry (``None`` +fields) so the rest of the benchmark suite still runs — missing data is +reported as missing, never invented as zeros. + +Power/energy semantics (three distinct quantities, never conflated): + +* **instantaneous power** — one tegrastats sample (``JetsonSample.power_mw``, + ``JetsonRun.power_mw``): the module power at that sampling instant. +* **average power over a window** — the mean of the instantaneous samples + that fall inside a benchmark window (``EnergySummary.avg_power_mw``). +* **integrated energy** — the trapezoidal integral of instantaneous power + over the window (``EnergySummary.energy_mj``). Energy-per-inference is only + reported when this integral exists and the inference count is known. """ from __future__ import annotations @@ -24,21 +35,53 @@ @dataclass class JetsonSample: + """One parsed tegrastats line. ``None`` means the field was absent from + the line (formats differ across JetPack releases) — not zero.""" + timestamp: float - gpu_util_pct: float - temp_c: float - power_mw: float + gpu_util_pct: Optional[float] + temp_c: Optional[float] + power_mw: Optional[float] @dataclass class JetsonRun: + """Per-inference telemetry. Power/temp/util are the *latest instantaneous* + tegrastats sample at the time the inference finished — not an average and + not an energy; see :class:`EnergySummary` for integrated quantities.""" + latency_ms: float exit_taken: int confidence: float - computation_used: float - gpu_util_pct: float - temp_c: float - power_mw: float + estimated_backbone_flops_fraction: float + gpu_util_pct: Optional[float] + temp_c: Optional[float] + power_mw: Optional[float] + + @property + def computation_used(self) -> float: + """Deprecated alias for ``estimated_backbone_flops_fraction``.""" + return self.estimated_backbone_flops_fraction + + +@dataclass +class EnergySummary: + """Power/energy over one benchmark window. + + ``energy_mj`` is the trapezoidal integral of the instantaneous power + samples inside the window; it (and ``energy_per_inference_mj``) is + ``None`` when fewer than two power samples landed in the window — a + single sample cannot be integrated, and inventing a rectangle would be a + fabricated measurement. + """ + + num_inferences: int + window_seconds: float + sampling_interval_ms: int + num_power_samples: int + avg_power_mw: Optional[float] + energy_mj: Optional[float] + energy_per_inference_mj: Optional[float] _GR3D_RE = re.compile(r"GR3D_FREQ\s+(\d+)%") @@ -52,48 +95,121 @@ def _parse_line(line: str) -> JetsonSample: power = _POWER_RE.search(line) return JetsonSample( timestamp=time.monotonic(), - gpu_util_pct=float(gpu.group(1)) if gpu else 0.0, - temp_c=float(temp.group(1)) if temp else 0.0, - power_mw=float(power.group(1)) if power else 0.0, + gpu_util_pct=float(gpu.group(1)) if gpu else None, + temp_c=float(temp.group(1)) if temp else None, + power_mw=float(power.group(1)) if power else None, + ) + + +def integrate_energy( + samples: list[JetsonSample], + window_start: float, + window_end: float, + num_inferences: int, + sampling_interval_ms: int, +) -> EnergySummary: + """Integrate instantaneous power samples over a benchmark window. + + Pure function (unit-testable without hardware). Only samples with a + parsed power value that fall inside ``[window_start, window_end]`` are + used; trapezoidal integration over their timestamps yields millijoules. + """ + window_seconds = max(window_end - window_start, 0.0) + in_window = [ + s for s in samples if s.power_mw is not None and window_start <= s.timestamp <= window_end + ] + avg_power: Optional[float] = None + energy_mj: Optional[float] = None + per_inference: Optional[float] = None + if in_window: + powers = [s.power_mw for s in in_window if s.power_mw is not None] + avg_power = sum(powers) / len(powers) + if len(in_window) >= 2: + energy = 0.0 # mW * s == mJ + for a, b in zip(in_window, in_window[1:]): + assert a.power_mw is not None and b.power_mw is not None + energy += 0.5 * (a.power_mw + b.power_mw) * (b.timestamp - a.timestamp) + energy_mj = energy + if num_inferences > 0: + per_inference = energy / num_inferences + return EnergySummary( + num_inferences=num_inferences, + window_seconds=window_seconds, + sampling_interval_ms=sampling_interval_ms, + num_power_samples=len(in_window), + avg_power_mw=avg_power, + energy_mj=energy_mj, + energy_per_inference_mj=per_inference, ) class TegrastatsMonitor: - """Spawn tegrastats, parse the latest line on demand, terminate cleanly.""" + """Spawn tegrastats, parse lines in a reader thread, terminate cleanly. + + Reusable: ``start()`` after ``stop()`` begins a fresh session (the stop + event is cleared, the sample buffer is reset). A second ``start()`` while + running is a no-op rather than a duplicate process/thread. + """ - def __init__(self, interval_ms: int = 100, max_samples: int = 256) -> None: + def __init__(self, interval_ms: int = 100, max_samples: int = 4096) -> None: + if interval_ms <= 0: + raise ValueError(f"interval_ms must be > 0 (got {interval_ms})") + if max_samples <= 0: + raise ValueError(f"max_samples must be > 0 (got {max_samples})") self.interval_ms = interval_ms + self._max_samples = max_samples self._samples: deque[JetsonSample] = deque(maxlen=max_samples) self._proc: Optional[subprocess.Popen[str]] = None self._thread: Optional[threading.Thread] = None self._stop = threading.Event() self.available = shutil.which("tegrastats") is not None + @property + def running(self) -> bool: + return self._proc is not None + def start(self) -> None: - if not self.available: + if not self.available or self._proc is not None: + return + # a fresh session: clear the previous session's stop flag and samples. + # Without this, a restarted monitor's reader thread exited immediately + # (the stop event stayed set) and silently collected nothing. + self._stop.clear() + self._samples = deque(maxlen=self._max_samples) + try: + self._proc = subprocess.Popen( + ["tegrastats", "--interval", str(self.interval_ms)], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + bufsize=1, + ) + except OSError: + # binary disappeared or failed to spawn: degrade to unavailable + self._proc = None + self.available = False return - self._proc = subprocess.Popen( - ["tegrastats", "--interval", str(self.interval_ms)], - stdout=subprocess.PIPE, - stderr=subprocess.DEVNULL, - text=True, - bufsize=1, - ) self._thread = threading.Thread(target=self._reader, daemon=True) self._thread.start() def _reader(self) -> None: - assert self._proc is not None and self._proc.stdout is not None - for line in self._proc.stdout: + proc = self._proc + if proc is None or proc.stdout is None: + return + for line in proc.stdout: if self._stop.is_set(): break self._samples.append(_parse_line(line)) def latest(self) -> JetsonSample: if not self._samples: - return JetsonSample(time.monotonic(), 0.0, 0.0, 0.0) + return JetsonSample(time.monotonic(), None, None, None) return self._samples[-1] + def samples(self) -> list[JetsonSample]: + """Snapshot of the current session's samples.""" + return list(self._samples) + def stop(self) -> None: self._stop.set() if self._proc is not None: @@ -102,6 +218,9 @@ def stop(self) -> None: self._proc.wait(timeout=2) except subprocess.TimeoutExpired: self._proc.kill() + self._proc.wait(timeout=2) + if self._proc.stdout is not None: + self._proc.stdout.close() # join the reader so we don't drop a live thread holding the pipe if self._thread is not None: self._thread.join(timeout=1) @@ -113,12 +232,12 @@ class JetsonProfiler: """Profile a wrapper on Jetson and produce per-run telemetry. If tegrastats is not available (non-Jetson host), profiling still runs - but power/temp/util fields are zero. This keeps tests and CI usable on - a dev laptop. + but power/temp/util fields are ``None``. This keeps tests and CI usable + on a dev laptop without fabricating hardware numbers. """ - def __init__(self) -> None: - self.monitor = TegrastatsMonitor() + def __init__(self, interval_ms: int = 100) -> None: + self.monitor = TegrastatsMonitor(interval_ms=interval_ms) def profile( self, @@ -128,6 +247,29 @@ def profile( num_runs: int = 100, device: str = "cuda" if torch.cuda.is_available() else "cpu", ) -> list[JetsonRun]: + runs, _ = self.profile_with_energy( + model, input_shape=input_shape, num_warmup=num_warmup, num_runs=num_runs, device=device + ) + return runs + + def profile_with_energy( + self, + model: EarlyExitWrapper, + input_shape: tuple[int, int, int, int] = (1, 3, 224, 224), + num_warmup: int = 50, + num_runs: int = 100, + device: str = "cuda" if torch.cuda.is_available() else "cpu", + ) -> tuple[list[JetsonRun], EnergySummary]: + """Per-run telemetry plus power integrated over the timed window. + + The energy summary covers only the timed loop (warmup excluded) and + divides by the true inference count; on hosts without tegrastats it + reports zero samples and ``None`` energy. + """ + if num_warmup < 0: + raise ValueError(f"num_warmup must be >= 0 (got {num_warmup})") + if num_runs <= 0: + raise ValueError(f"num_runs must be > 0 (got {num_runs})") model = model.to(device).eval() dummy = torch.randn(input_shape, device=device) @@ -140,6 +282,7 @@ def profile( torch.cuda.synchronize() runs: list[JetsonRun] = [] + window_start = time.monotonic() with torch.no_grad(): for _ in range(num_runs): t0 = time.perf_counter() @@ -153,12 +296,22 @@ def profile( latency_ms=latency * 1000, exit_taken=result.exit_taken, confidence=result.confidence, - computation_used=result.computation_used, + estimated_backbone_flops_fraction=( + result.estimated_backbone_flops_fraction + ), gpu_util_pct=sample.gpu_util_pct, temp_c=sample.temp_c, power_mw=sample.power_mw, ) ) - return runs + window_end = time.monotonic() + summary = integrate_energy( + self.monitor.samples(), + window_start, + window_end, + num_inferences=num_runs, + sampling_interval_ms=self.monitor.interval_ms, + ) + return runs, summary finally: self.monitor.stop() diff --git a/earlyon/benchmarking/throughput.py b/earlyon/benchmarking/throughput.py index ceb5666..fec37c7 100644 --- a/earlyon/benchmarking/throughput.py +++ b/earlyon/benchmarking/throughput.py @@ -1,13 +1,28 @@ -"""Throughput / latency benchmarking with proper CUDA sync. - -Single-sample inference benchmark (batch=1) per v0.1 routing constraint. Use -``num_warmup>=50`` on CUDA; GPU clocks need ~30 iterations to stabilize. - -``benchmark_wrapper`` runs a fixed random-noise input and is fast to -reproduce; trained heads may fire spuriously on noise, so it's a best-case -upper bound. ``benchmark_wrapper_on_loader`` drives real samples through the -wrapper and is the honest input-distribution signal — see the v0.2 README -appendix for the dual numbers. +"""Throughput / latency benchmarking with identical boundaries per model. + +Fairness contract +----------------- +:func:`benchmark_models` is the one measurement core. Every model in a +comparison sees **the same sample sequence** (preloaded once, before any +timing), the same warmup count, the same ``eval()`` + ``torch.inference_mode`` +setup, the same device placement, and the same per-iteration synchronization. +Two boundaries are supported and always labelled on the result: + +* ``"model-only"`` (default) — samples are moved to the device *before* the + timed region; the measurement covers the forward pass (and, for a wrapper, + its routing logic + host synchronization) only. +* ``"end-to-end"`` — samples stay on CPU and the timed region includes the + host-to-device copy. Loader/preprocessing time is never measured by either + boundary. + +The legacy helpers (:func:`benchmark_wrapper`, :func:`benchmark_backbone`, +:func:`benchmark_wrapper_on_loader`) are thin wrappers over the core and keep +their signatures. Note that a wrapper benchmarked on random noise is a +best-case bound — trained heads may fire spuriously on noise — so prefer the +loader variants for honest input-distribution numbers. + +Batch size 1 is the primary supported mode: it is the latency-sensitive edge +case earlyon targets, and the single-sample router requires it. """ from __future__ import annotations @@ -15,30 +30,75 @@ import statistics import time from collections import Counter -from dataclasses import asdict, dataclass -from typing import Any, Iterator +from dataclasses import asdict, dataclass, field +from typing import Any, Mapping import torch +import torch.nn as nn from torch.utils.data import DataLoader from earlyon.core.types import Batch, exit_label from earlyon.core.wrappers import EarlyExitWrapper +BOUNDARY_MODEL_ONLY = "model-only" +BOUNDARY_END_TO_END = "end-to-end" + @dataclass class BenchmarkResult: + """One model's measurement under the fairness contract. + + ``avg_estimated_flops_fraction`` is the average of the wrapper's static + per-exit FLOPs estimate (1.0 for plain modules) — an estimate, not a + measurement; see ``InferenceResult.estimated_backbone_flops_fraction``. + ``accuracy`` is populated only when the samples came from a labelled + loader. + """ + throughput_ips: float latency_median_ms: float latency_p50_ms: float latency_p95_ms: float latency_p99_ms: float - avg_computation_used: float + avg_estimated_flops_fraction: float exit_distribution: dict[str, float] num_runs: int device: str + num_warmup: int = 0 + batch_size: int = 1 + input_shape: tuple[int, ...] = () + dtype: str = "torch.float32" + boundary: str = BOUNDARY_MODEL_ONLY + sync: str = "none" # "cuda-synchronize-per-iteration" on CUDA devices + input_source: str = "random-noise" # or "loader" + accuracy: float | None = None + num_labelled_samples: int = 0 + + @property + def avg_computation_used(self) -> float: + """Deprecated alias for ``avg_estimated_flops_fraction``.""" + return self.avg_estimated_flops_fraction def to_dict(self) -> dict[str, Any]: - return asdict(self) + d = asdict(self) + d["input_shape"] = list(self.input_shape) + return d + + +@dataclass +class ComparisonResult: + """Results of benchmarking several models on identical samples. + + ``speedup_vs`` divides two models' throughputs; it is only meaningful + because the fairness contract guarantees identical inputs and boundaries. + """ + + results: dict[str, BenchmarkResult] = field(default_factory=dict) + + def speedup_vs(self, model: str, baseline: str) -> float: + return self.results[model].throughput_ips / max( + self.results[baseline].throughput_ips, 1e-12 + ) def _sync(device: str) -> None: @@ -51,126 +111,219 @@ def _percentile(values: list[float], p: float) -> float: return sorted(values)[idx] -def benchmark_wrapper( - model: EarlyExitWrapper, - input_shape: tuple[int, int, int, int] = (1, 3, 224, 224), - device: str = "cpu", - num_warmup: int = 50, - num_runs: int = 500, -) -> BenchmarkResult: - """Benchmark a wrapper in inference mode.""" +def _draw_samples( + loader: DataLoader[Batch] | None, + input_shape: tuple[int, int, int, int] | None, + count: int, + seed: int, +) -> tuple[list[torch.Tensor], list[torch.Tensor] | None]: + """Materialize the exact sample sequence every model will see. + + From a loader: cycle it until ``count`` samples are drawn (keeping + labels). Without one: fixed-seed random noise, generated once. + """ + if loader is not None: + if loader.batch_size is None or loader.batch_size != 1: + raise ValueError( + f"benchmarking requires loader.batch_size=1 (got {loader.batch_size}); " + "batch-1 latency is the primary supported mode" + ) + try: + if len(loader) == 0: + raise ValueError("benchmarking requires a non-empty loader") + except TypeError: + pass # unsized (IterableDataset) — cannot pre-check + images: list[torch.Tensor] = [] + labels: list[torch.Tensor] = [] + while len(images) < count: + drew_any = False + for x, y in loader: + drew_any = True + images.append(x) + labels.append(y) + if len(images) >= count: + break + if not drew_any: + raise ValueError("benchmarking requires a non-empty loader") + return images, labels + if input_shape is None: + raise ValueError("provide either a loader or an input_shape") if input_shape[0] != 1: - raise ValueError("v0.1 inference benchmark requires batch_size=1") + raise ValueError( + f"benchmarking requires batch_size=1 (got input_shape={input_shape}); " + "batch-1 latency is the primary supported mode" + ) + gen = torch.Generator().manual_seed(seed) + return [torch.randn(input_shape, generator=gen) for _ in range(count)], None + + +def _run_one( + model: nn.Module, + samples: list[torch.Tensor], + labels: list[torch.Tensor] | None, + device: str, + num_warmup: int, + num_runs: int, + boundary: str, + input_source: str, +) -> BenchmarkResult: + is_wrapper = isinstance(model, EarlyExitWrapper) model = model.to(device).eval() - dummy = torch.randn(input_shape, device=device) - # warmup - with torch.no_grad(): - for _ in range(num_warmup): - _ = model(dummy, mode="inference") + if boundary == BOUNDARY_MODEL_ONLY: + staged = [s.to(device) for s in samples] + elif boundary == BOUNDARY_END_TO_END: + staged = [s.cpu() for s in samples] # timed region includes .to(device) + else: + raise ValueError(f"unknown boundary {boundary!r}") + + def forward(x: torch.Tensor) -> Any: + if boundary == BOUNDARY_END_TO_END: + x = x.to(device) + if is_wrapper: + return model(x, mode="inference") + return model(x) + + with torch.inference_mode(): + for i in range(num_warmup): + forward(staged[i]) _sync(device) latencies: list[float] = [] exits: Counter[int] = Counter() - comp_used: list[float] = [] + flops_fraction: list[float] = [] + correct = 0 + labelled = 0 - with torch.no_grad(): + with torch.inference_mode(): wall_start = time.perf_counter() - for _ in range(num_runs): + for i in range(num_runs): + x = staged[num_warmup + i] t0 = time.perf_counter() - result = model(dummy, mode="inference") + out = forward(x) _sync(device) latencies.append(time.perf_counter() - t0) - exits[result.exit_taken] += 1 - comp_used.append(result.computation_used) + if is_wrapper: + exits[out.exit_taken] += 1 + flops_fraction.append(out.estimated_backbone_flops_fraction) + prediction = out.prediction + else: + exits[-1] += 1 + flops_fraction.append(1.0) + prediction = out + if labels is not None: + target = int(labels[num_warmup + i].item()) + correct += int(int(prediction.argmax(dim=-1).item()) == target) + labelled += 1 wall_total = time.perf_counter() - wall_start - # exit_distribution keys: "exit_0", "exit_1", ..., "final" dist = {exit_label(k): v / num_runs for k, v in exits.items()} - return BenchmarkResult( throughput_ips=num_runs / wall_total, latency_median_ms=statistics.median(latencies) * 1000, latency_p50_ms=_percentile(latencies, 0.50) * 1000, latency_p95_ms=_percentile(latencies, 0.95) * 1000, latency_p99_ms=_percentile(latencies, 0.99) * 1000, - avg_computation_used=sum(comp_used) / len(comp_used), + avg_estimated_flops_fraction=sum(flops_fraction) / len(flops_fraction), exit_distribution=dist, num_runs=num_runs, device=device, + num_warmup=num_warmup, + batch_size=1, + input_shape=tuple(samples[0].shape), + dtype=str(samples[0].dtype), + boundary=boundary, + sync="cuda-synchronize-per-iteration" if device.startswith("cuda") else "none", + input_source=input_source, + accuracy=(correct / labelled) if labelled else None, + num_labelled_samples=labelled, ) -def benchmark_wrapper_on_loader( - model: EarlyExitWrapper, - loader: DataLoader[Batch], +def benchmark_models( + models: Mapping[str, nn.Module], + loader: DataLoader[Batch] | None = None, + input_shape: tuple[int, int, int, int] | None = None, device: str = "cpu", num_warmup: int = 50, num_runs: int = 500, -) -> BenchmarkResult: - """Benchmark a wrapper in inference mode using real samples from a loader. + boundary: str = BOUNDARY_MODEL_ONLY, + seed: int = 0, +) -> ComparisonResult: + """Benchmark several models under identical boundaries and inputs. + + ``models`` maps a label to a module: the raw backbone, an + :class:`EarlyExitWrapper` (routed via its inference path), a smaller + static baseline, a quantised variant — anything callable on the samples. + Typical usage:: + + cmp = benchmark_models( + {"backbone": model.backbone, "early_exit": model, "small": resnet18}, + loader=test_loader, device="cuda", + ) + cmp.speedup_vs("early_exit", "backbone") - The loader must yield (image, label) batches with ``batch_size == 1`` to - match v0.2's single-sample routing. The first ``num_warmup`` samples are - discarded; the next ``num_runs`` are timed. The loader is cycled if it - runs out before ``num_warmup + num_runs`` samples have been drawn. + The exact same preloaded sample sequence is fed to every model, in the + same order, with the same warmup, boundary and synchronization. Samples + come from ``loader`` (cycled; labels kept, accuracy reported) or, when no + loader is given, from fixed-seed random noise (best-case bound: trained + exit heads may fire spuriously on noise). """ - if loader.batch_size is None or loader.batch_size != 1: - raise ValueError( - f"benchmark_wrapper_on_loader requires loader.batch_size=1 " - f"(got {loader.batch_size}); v0.2 inference is single-sample" + if num_warmup < 0: + raise ValueError(f"num_warmup must be >= 0 (got {num_warmup})") + if num_runs <= 0: + raise ValueError(f"num_runs must be > 0 (got {num_runs})") + if not models: + raise ValueError("models must be non-empty") + samples, labels = _draw_samples(loader, input_shape, num_warmup + num_runs, seed) + input_source = "loader" if loader is not None else "random-noise" + comparison = ComparisonResult() + for name, model in models.items(): + comparison.results[name] = _run_one( + model, samples, labels, device, num_warmup, num_runs, boundary, input_source ) - # an empty loader would make the cycling sample_stream spin forever; fail fast. - try: - is_empty = len(loader) == 0 - except TypeError: - is_empty = False # unsized (IterableDataset) — cannot pre-check - if is_empty: - raise ValueError("benchmark_wrapper_on_loader requires a non-empty loader") - model = model.to(device).eval() + return comparison - def sample_stream() -> Iterator[torch.Tensor]: - while True: - for images, _targets in loader: - yield images.to(device) - stream = sample_stream() +# ---------------- legacy single-model helpers ---------------- - with torch.no_grad(): - for _ in range(num_warmup): - x = next(stream) - _ = model(x, mode="inference") - _sync(device) - latencies: list[float] = [] - exits: Counter[int] = Counter() - comp_used: list[float] = [] +def benchmark_wrapper( + model: EarlyExitWrapper, + input_shape: tuple[int, int, int, int] = (1, 3, 224, 224), + device: str = "cpu", + num_warmup: int = 50, + num_runs: int = 500, +) -> BenchmarkResult: + """Benchmark a wrapper on fixed-seed random noise (best-case bound).""" + return benchmark_models( + {"model": model}, + input_shape=input_shape, + device=device, + num_warmup=num_warmup, + num_runs=num_runs, + ).results["model"] - with torch.no_grad(): - wall_start = time.perf_counter() - for _ in range(num_runs): - x = next(stream) - t0 = time.perf_counter() - result = model(x, mode="inference") - _sync(device) - latencies.append(time.perf_counter() - t0) - exits[result.exit_taken] += 1 - comp_used.append(result.computation_used) - wall_total = time.perf_counter() - wall_start - dist = {exit_label(k): v / num_runs for k, v in exits.items()} +def benchmark_wrapper_on_loader( + model: EarlyExitWrapper, + loader: DataLoader[Batch], + device: str = "cpu", + num_warmup: int = 50, + num_runs: int = 500, +) -> BenchmarkResult: + """Benchmark a wrapper on real samples from a batch-1 loader. - return BenchmarkResult( - throughput_ips=num_runs / wall_total, - latency_median_ms=statistics.median(latencies) * 1000, - latency_p50_ms=_percentile(latencies, 0.50) * 1000, - latency_p95_ms=_percentile(latencies, 0.95) * 1000, - latency_p99_ms=_percentile(latencies, 0.99) * 1000, - avg_computation_used=sum(comp_used) / len(comp_used), - exit_distribution=dist, - num_runs=num_runs, + This is the honest input-distribution signal; accuracy over the timed + samples is reported alongside speed. + """ + return benchmark_models( + {"model": model}, + loader=loader, device=device, - ) + num_warmup=num_warmup, + num_runs=num_runs, + ).results["model"] def benchmark_backbone( @@ -181,32 +334,10 @@ def benchmark_backbone( num_runs: int = 500, ) -> BenchmarkResult: """Benchmark the raw backbone (no early-exit logic). Used as baseline.""" - backbone = backbone.to(device).eval() - dummy = torch.randn(input_shape, device=device) - - with torch.no_grad(): - for _ in range(num_warmup): - _ = backbone(dummy) - _sync(device) - - latencies: list[float] = [] - with torch.no_grad(): - wall_start = time.perf_counter() - for _ in range(num_runs): - t0 = time.perf_counter() - _ = backbone(dummy) - _sync(device) - latencies.append(time.perf_counter() - t0) - wall_total = time.perf_counter() - wall_start - - return BenchmarkResult( - throughput_ips=num_runs / wall_total, - latency_median_ms=statistics.median(latencies) * 1000, - latency_p50_ms=_percentile(latencies, 0.50) * 1000, - latency_p95_ms=_percentile(latencies, 0.95) * 1000, - latency_p99_ms=_percentile(latencies, 0.99) * 1000, - avg_computation_used=1.0, - exit_distribution={"final": 1.0}, - num_runs=num_runs, + return benchmark_models( + {"model": backbone}, + input_shape=input_shape, device=device, - ) + num_warmup=num_warmup, + num_runs=num_runs, + ).results["model"] diff --git a/earlyon/cli.py b/earlyon/cli.py index 0b70577..5729b24 100644 --- a/earlyon/cli.py +++ b/earlyon/cli.py @@ -11,8 +11,7 @@ from earlyon import __version__ from earlyon.benchmarking import ( JetsonProfiler, - benchmark_backbone, - benchmark_wrapper, + benchmark_models, evaluate, ) from earlyon.core.thresholds import calibrate_thresholds, calibrate_thresholds_for_budget @@ -253,17 +252,26 @@ def calibrate( def benchmark( model_path: str, device: str, runs: int, warmup: int, input_size: int, json_out: str | None ) -> None: - """Throughput + latency benchmark, single-sample (batch=1).""" + """Throughput + latency benchmark, single-sample (batch=1). + + The wrapper and its raw backbone are measured on the exact same + fixed-seed noise samples with identical boundaries. Noise input is a + best-case bound (trained heads may fire spuriously); use the Python API + with a real loader for the honest input-distribution signal. + """ dev = _device(device) model = load_wrapper(model_path) shape = (1, 3, input_size, input_size) - wrap_r = benchmark_wrapper( - model, input_shape=shape, device=dev, num_warmup=warmup, num_runs=runs - ) - bb_r = benchmark_backbone( - model.backbone, input_shape=shape, device=dev, num_warmup=warmup, num_runs=runs + cmp_r = benchmark_models( + {"early_exit": model, "backbone": model.backbone}, + input_shape=shape, + device=dev, + num_warmup=warmup, + num_runs=runs, ) - speedup = wrap_r.throughput_ips / max(bb_r.throughput_ips, 1e-9) + wrap_r = cmp_r.results["early_exit"] + bb_r = cmp_r.results["backbone"] + speedup = cmp_r.speedup_vs("early_exit", "backbone") click.echo( json.dumps( { @@ -295,11 +303,16 @@ def benchmark( @click.option("--input-size", default=224, type=int) @click.option("--device", default="auto") def profile(model_path: str, runs: int, warmup: int, input_size: int, device: str) -> None: - """Jetson profile — reads tegrastats when available.""" + """Jetson profile — reads tegrastats when available. + + Reports instantaneous-power medians and, separately, integrated energy + over the timed window. Missing telemetry (non-Jetson host) is reported as + null, never as zero. + """ dev = _device(device) model = load_wrapper(model_path) profiler = JetsonProfiler() - runs_out = profiler.profile( + runs_out, energy = profiler.profile_with_energy( model, input_shape=(1, 3, input_size, input_size), num_warmup=warmup, @@ -309,15 +322,22 @@ def profile(model_path: str, runs: int, warmup: int, input_size: int, device: st import statistics lat = [r.latency_ms for r in runs_out] - power = [r.power_mw for r in runs_out] - temp = [r.temp_c for r in runs_out] + power = [r.power_mw for r in runs_out if r.power_mw is not None] + temp = [r.temp_c for r in runs_out if r.temp_c is not None] click.echo( json.dumps( { "runs": len(runs_out), "latency_median_ms": statistics.median(lat), - "power_median_mw": statistics.median(power), - "temp_median_c": statistics.median(temp), + "instantaneous_power_median_mw": statistics.median(power) if power else None, + "temp_median_c": statistics.median(temp) if temp else None, + "energy": { + "window_seconds": energy.window_seconds, + "num_power_samples": energy.num_power_samples, + "avg_power_mw": energy.avg_power_mw, + "energy_mj": energy.energy_mj, + "energy_per_inference_mj": energy.energy_per_inference_mj, + }, "tegrastats_available": profiler.monitor.available, }, indent=2, diff --git a/earlyon/core/flops.py b/earlyon/core/flops.py index 3858dc5..c8d00f7 100644 --- a/earlyon/core/flops.py +++ b/earlyon/core/flops.py @@ -1,44 +1,135 @@ -"""FLOPs accounting helpers. +"""Static FLOPs *estimation* helpers. -Uses fvcore for per-module flops. The non-trivial part is: when an exit point -is a nested module like ``features.3`` (MobileNetV2), we need to sum the flops -of every LEAF op executed before and including ``features.3`` in forward -order, without double-counting the container module. +What this module produces is an estimate, not a measurement. The estimator +(fvcore's ``FlopCountAnalysis``) statically attributes FLOPs per module; the +cumulative fraction at an exit layer is derived from module *registration +order* filtered to leaves, which matches execution order for straight-line +(sequential-style) backbones — the torchvision CNNs and ViT earlyon ships. + +Known limitations (all surfaced, none silent): + +* **Exit-head FLOPs are excluded.** Only the backbone is analysed; the small + exit heads (~10-100k params) and all routing overhead are not counted. +* **Reused modules break the ordering assumption.** If any leaf module runs + more than once in a forward pass (weight sharing, recursive blocks), a + cumulative "up to layer L" sum over registration order is not meaningful. + The probe detects this and returns a *low-confidence uniform estimate* + (``method="uniform-fallback"``, ``reliable=False``) with a warning, instead + of a precise-looking but wrong fraction. +* fvcore does not count some ops (e.g. scaled-dot-product attention matmuls), + so ViT fractions are a slight under-count; they remain monotone. + +The analysis costs one or two forward passes of the backbone, so callers +(``EarlyExitWrapper``) run it lazily on first use, not at construction. """ from __future__ import annotations import warnings +from dataclasses import dataclass, field from typing import Sequence import torch import torch.nn as nn +METHOD_FVCORE = "fvcore-leaf-walk" +METHOD_UNIFORM = "uniform-fallback" -def per_layer_flops( + +@dataclass(frozen=True) +class FlopsEstimate: + """Cumulative backbone-FLOPs fractions per exit layer, with provenance. + + ``fractions[layer]`` is (estimated FLOPs of everything executed up to and + including ``layer``) / (total estimated backbone FLOPs), in [0, 1]. + ``reliable`` is False when the estimator had to fall back to uniform + spacing (fvcore missing, zero total, or a reused/multi-call module that + invalidates the ordering assumption); treat those fractions as layout + placeholders, not compute estimates. + """ + + fractions: dict[str, float] + method: str + reliable: bool + notes: tuple[str, ...] = field(default_factory=tuple) + excludes_exit_heads: bool = True + + +def _uniform(layer_names: Sequence[str], note: str) -> FlopsEstimate: + n = len(layer_names) + return FlopsEstimate( + fractions={name: (i + 1) / (n + 1) for i, name in enumerate(layer_names)}, + method=METHOD_UNIFORM, + reliable=False, + notes=(note,), + ) + + +def _count_leaf_calls( + backbone: nn.Module, input_shape: tuple[int, int, int, int], device: str +) -> dict[str, int]: + """One dummy forward counting how many times each leaf module executes.""" + counts: dict[str, int] = {} + handles = [] + for name, mod in backbone.named_modules(): + if name and len(list(mod.children())) == 0: + counts[name] = 0 + + def _hook(_m: nn.Module, _i: object, _o: object, _name: str = name) -> None: + counts[_name] += 1 + + handles.append(mod.register_forward_hook(_hook)) + try: + with torch.no_grad(): + backbone(torch.zeros(input_shape, device=device)) + finally: + for h in handles: + h.remove() + return counts + + +def estimate_layer_flops( backbone: nn.Module, layer_names: Sequence[str], input_shape: tuple[int, int, int, int] = (1, 3, 224, 224), -) -> dict[str, float]: - """Return the cumulative FLOPs fraction at each named layer. + device: str = "cpu", +) -> FlopsEstimate: + """Estimate the cumulative FLOPs fraction at each named layer. - At layer L the value is (FLOPs of every leaf op executed up to and - including L) / (total backbone FLOPs). The result is monotonic in - forward order and bounded in [0, 1]. + ``device`` is where the probe input is created — it must match the + backbone's parameters (the wrapper passes its backbone's current device, + since the estimate is computed lazily and the model may have moved). + + See the module docstring for the estimator's assumptions and limitations. """ try: from fvcore.nn import FlopCountAnalysis except ImportError: - n = len(layer_names) - return {name: (i + 1) / (n + 1) for i, name in enumerate(layer_names)} + return _uniform(layer_names, "fvcore not installed; uniform spacing used") # eval() for a deterministic FLOPs probe, but restore the caller's mode: - # this runs inside EarlyExitWrapper.__init__, and silently leaving a - # training backbone in eval would freeze BatchNorm in a custom training loop. + # silently leaving a training backbone in eval would freeze BatchNorm in a + # custom training loop. was_training = backbone.training backbone.eval() try: - dummy = torch.zeros(input_shape) + # Reused-module detection: a leaf that executes more than once makes + # "cumulative FLOPs up to layer L in registration order" meaningless. + call_counts = _count_leaf_calls(backbone, input_shape, device) + reused = sorted(name for name, c in call_counts.items() if c > 1) + if reused: + shown = ", ".join(reused[:5]) + (", ..." if len(reused) > 5 else "") + warnings.warn( + f"backbone reuses module(s) [{shown}] within one forward pass; " + "per-layer cumulative FLOPs cannot be attributed reliably. " + "Falling back to a low-confidence uniform estimate " + "(FlopsEstimate.reliable=False).", + RuntimeWarning, + stacklevel=2, + ) + return _uniform(layer_names, f"reused modules invalidate ordering: {shown}") + + dummy = torch.zeros(input_shape, device=device) fca = FlopCountAnalysis(backbone, dummy) fca.unsupported_ops_warnings(False) fca.uncalled_modules_warnings(False) @@ -46,8 +137,7 @@ def per_layer_flops( by_module = fca.by_module() if total == 0: - n = len(layer_names) - return {name: (i + 1) / (n + 1) for i, name in enumerate(layer_names)} + return _uniform(layer_names, "estimator reported zero total FLOPs") # Forward order = order returned by named_modules. Filter to leaves only # (modules with no children) to avoid double-counting containers. @@ -82,6 +172,7 @@ def per_layer_flops( # The "end" of an exit at layer L is the last leaf whose name belongs to # L (equals L or is nested under it). Anything past that point hasn't run # yet when the exit fires. + notes: list[str] = [] cumulative: dict[str, float] = {} for layer in layer_names: last_idx = -1 @@ -93,6 +184,7 @@ def per_layer_flops( # uniform spacing for this entry n = len(layer_names) cumulative[layer] = (list(layer_names).index(layer) + 1) / (n + 1) + notes.append(f"layer {layer!r} has no leaves; uniform value used") continue running = sum(int(by_module.get(ordered_leaves[i], 0)) for i in range(last_idx + 1)) running += sum(flops for pos, flops in nonleaf_self if pos <= last_idx) @@ -107,8 +199,27 @@ def per_layer_flops( RuntimeWarning, stacklevel=2, ) + notes.append(f"overcount clamped at {layer!r}") cumulative[layer] = min(ratio, 1.0) - return cumulative + return FlopsEstimate( + fractions=cumulative, + method=METHOD_FVCORE, + reliable=not notes, + notes=tuple(notes), + ) finally: backbone.train(was_training) + + +def per_layer_flops( + backbone: nn.Module, + layer_names: Sequence[str], + input_shape: tuple[int, int, int, int] = (1, 3, 224, 224), +) -> dict[str, float]: + """Backward-compatible wrapper returning just the fraction mapping. + + Prefer :func:`estimate_layer_flops`, which also reports the method used + and whether the estimate is reliable. + """ + return estimate_layer_flops(backbone, layer_names, input_shape).fractions diff --git a/earlyon/core/temperature.py b/earlyon/core/temperature.py index 96889e8..d6adabc 100644 --- a/earlyon/core/temperature.py +++ b/earlyon/core/temperature.py @@ -5,18 +5,34 @@ in log-space to keep T positive and use LBFGS, which converges in a handful of iterations. -The fitted temperature is stored on ``EarlyExitConfig.temperature`` and used -by the wrapper at inference time. +Fitted temperatures are stored per head on ``EarlyExitConfig.temperatures`` +and used by the wrapper at inference time. """ from __future__ import annotations import math +from dataclasses import dataclass import torch import torch.nn.functional as F +@dataclass(frozen=True) +class TemperatureFitResult: + """Outcome of one per-head temperature fit. + + ``fallback`` is True when LBFGS diverged (non-finite iterate) and the fit + kept the last finite value — or 1.0 if even the first step blew up. A + fallback is safe (never an artificially sharp temperature) but means the + head is effectively uncalibrated; the value is reported rather than hidden. + """ + + temperature: float + converged: bool # iterates stabilised within tolerance + fallback: bool # divergence forced the last-finite / 1.0 fallback + + def fit_temperature( logits: torch.Tensor, targets: torch.Tensor, @@ -25,6 +41,20 @@ def fit_temperature( max_outer_steps: int = 10, tol: float = 1e-6, ) -> float: + """Backward-compatible scalar wrapper around :func:`fit_temperature_full`.""" + return fit_temperature_full( + logits, targets, max_iter=max_iter, lr=lr, max_outer_steps=max_outer_steps, tol=tol + ).temperature + + +def fit_temperature_full( + logits: torch.Tensor, + targets: torch.Tensor, + max_iter: int = 50, + lr: float = 1.0, + max_outer_steps: int = 10, + tol: float = 1e-6, +) -> TemperatureFitResult: """Find T > 0 minimizing NLL(logits / T, targets). ``lr`` defaults to 1.0 (the PyTorch LBFGS default); the previous 0.01 @@ -60,14 +90,28 @@ def closure() -> torch.Tensor: # blew up — i.e. no calibration rather than a broken one). best_log_t = 0.0 prev = float("inf") + converged = False + fallback = False for _ in range(max_outer_steps): optimizer.step(closure) # type: ignore[no-untyped-call] # torch stub gap current = log_t.item() if not math.isfinite(current): + fallback = True break best_log_t = current if abs(current - prev) < tol: + converged = True break prev = current - return float(math.exp(best_log_t)) + try: + temperature = float(math.exp(best_log_t)) + except OverflowError: + temperature = float("inf") + if not math.isfinite(temperature) or temperature <= 0.0: + # a finite log_t can still overflow exp(); never return a value the + # config validator would reject — fall back to the no-op temperature. + temperature = 1.0 + converged = False + fallback = True + return TemperatureFitResult(temperature=temperature, converged=converged, fallback=fallback) diff --git a/earlyon/core/thresholds.py b/earlyon/core/thresholds.py index 35c56d7..3a2aacc 100644 --- a/earlyon/core/thresholds.py +++ b/earlyon/core/thresholds.py @@ -1,59 +1,190 @@ -"""Greedy threshold calibration. - -Two calibration modes share one greedy coordinate-descent search: +"""Threshold calibration: staged, testable, greedy-by-construction. + +The pipeline is four separate stages, each independently callable and tested: + +1. **Collect** — :func:`collect_head_logits` runs the calibration loader once + through the all-exits (training-mode) forward and caches every head's raw + logits plus targets. The network runs exactly once. +2. **Fit temperatures** — :func:`fit_head_temperatures` fits one temperature + per head (every exit head *and* the final classifier) on cached logits via + post-hoc temperature scaling (Guo et al. 2017). Each head is a different + classifier with its own miscalibration; a single shared temperature is not + defensible. +3. **Search thresholds** — a greedy coordinate-descent sweep over a fixed + grid, simulated vectorially against the cache. It is *not* a joint or + globally optimal search and is not claimed to be one; see + ``docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md``. +4. **Evaluate** — the selected policy is scored on the calibration split + (accuracy, estimated compute, exit distribution). Held-out *test* + evaluation is a separate concern: use + :func:`earlyon.benchmarking.evaluate` on a test loader; the threshold + search never sees test labels. + +Two public calibrators share the search: * :func:`calibrate_thresholds` — accuracy budget. For each exit (low-index to - high-index), lower the threshold through a fixed grid until the resulting - accuracy drop on the validation set exceeds ``target_accuracy_drop``. + high-index), lower the threshold through the grid; keep the most aggressive + value whose accuracy drop stays within ``target_accuracy_drop``. * :func:`calibrate_thresholds_for_budget` — compute budget. For each exit, - keep the threshold that brings average compute within ``target_computation`` - while losing the least accuracy; exits that cannot help stay disabled. - -The search is coordinate-descent, not joint-optimal; this is fast and gives -consistently good results in practice. Trials are evaluated against -:class:`_EvalCache` — one batched forward pass collects every head's logits, -and each threshold setting is then simulated vectorially, so the network runs -once instead of once per grid point. - -When ``fit_temperature=True``, post-hoc temperature scaling (Guo et al. 2017) -is fit on a held-out set BEFORE the threshold search begins. The fitted -scalar lands on ``model.config.temperature`` so every subsequent softmax -(inside the hook and at the final classifier) is calibrated. + keep the threshold that brings average estimated compute within + ``target_computation`` while losing the least accuracy. + +Exits for which no grid value passes are left **disabled** via +``config.enabled_exits`` — an explicit boolean, not a threshold sentinel, so +a saturated softmax (confidence exactly 1.0) can never fire a disabled exit. """ from __future__ import annotations import math import warnings -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Callable import torch from torch.utils.data import DataLoader -from earlyon.core.temperature import fit_temperature as _fit_temperature -from earlyon.core.types import Batch +from earlyon.core.temperature import TemperatureFitResult, fit_temperature_full +from earlyon.core.types import FINAL_HEAD, Batch, exit_label from earlyon.core.wrappers import EarlyExitWrapper, _safe_temperature +CALIBRATION_SCHEMA_VERSION = 2 + +# confidence grid: ordered conservative -> aggressive (lower threshold fires more) +DEFAULT_GRID: tuple[float, ...] = (0.99, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.6, 0.5) + +# entropy grid: fractions of H_max = ln(num_classes), ordered conservative -> +# aggressive (higher entropy threshold fires more, since exit fires on H <= thr) +DEFAULT_ENTROPY_FRACTIONS: tuple[float, ...] = (0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7) + @dataclass class CalibrationResult: + """Everything a calibration run decided and measured. + + ``avg_computation_used`` is the average *estimated backbone FLOPs + fraction* over the calibration set (see + ``InferenceResult.estimated_backbone_flops_fraction`` for the estimator's + limitations); it is not a wall-clock or measured-FLOPs quantity. + + ``baseline_accuracy`` is the routed accuracy with every exit disabled, + i.e. the final classifier's accuracy on the calibration split. + ``final_accuracy`` is the routed accuracy under the selected policy on the + same split. Both are calibration-split numbers — report test-set accuracy + from a separate held-out loader. + """ + thresholds: list[float] baseline_accuracy: float final_accuracy: float avg_computation_used: float - fitted_temperature: float | None = None + enabled_exits: list[bool] = field(default_factory=list) + accuracy_delta: float = 0.0 + exit_distribution: dict[str, float] = field(default_factory=dict) + num_samples: int = 0 + fitted_temperature: float | None = None # deprecated: final head's fitted T + temperatures: dict[str, float] | None = None # per-head fitted temperatures + temperature_fits: dict[str, TemperatureFitResult] | None = None policy: str = "confidence" # which routing policy the thresholds calibrate + objective: str = "accuracy_budget" # "accuracy_budget" | "compute_budget" + method: str = "greedy-coordinate-grid" # honestly named: not a joint optimum + target_accuracy_drop: float | None = None target_computation: float | None = None # set by budget calibration budget_met: bool = True # False when target_computation was unattainable + schema_version: int = CALIBRATION_SCHEMA_VERSION -# confidence grid: ordered conservative -> aggressive (lower threshold fires more) -DEFAULT_GRID: tuple[float, ...] = (0.99, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.6, 0.5) +# ---------------- stage 1: collect ---------------- + + +@dataclass(frozen=True) +class HeadLogits: + """Cached raw (temperature-1) logits for every head on one dataset pass. + + ``head_names`` is the exit names in forward order plus ``"final"`` last; + ``logits[i]`` is the ``(N, num_classes)`` CPU tensor for ``head_names[i]``; + ``flops`` is each head's cumulative estimated backbone-FLOPs fraction + (1.0 for the final head). + """ + + head_names: list[str] + logits: list[torch.Tensor] + targets: torch.Tensor + flops: torch.Tensor + + @property + def num_samples(self) -> int: + return int(self.targets.shape[0]) + + +def collect_head_logits( + model: EarlyExitWrapper, loader: DataLoader[Batch], device: str +) -> HeadLogits: + """One batched training-mode pass caching every head's logits + targets. + + Training mode never fires an exit (the hook only appends logits), so + every head — including disabled ones — is captured, and temperatures are + irrelevant here (they apply at softmax time, not logit time). + + Raises ``ValueError`` on an empty loader: silently calibrating on zero + samples would produce a policy chosen from nothing. + """ + model = model.to(device) + per_head: list[list[torch.Tensor]] = [] + targets: list[torch.Tensor] = [] + was_training = model.training + model.eval() + try: + with torch.no_grad(): + for images, tgt in loader: + outputs = model(images.to(device), mode="training") + if not per_head: + per_head = [[] for _ in outputs] + for idx, logits in enumerate(outputs): + per_head[idx].append(logits.detach().cpu()) + targets.append(tgt.cpu()) + finally: + model.train(was_training) + if not targets: + raise ValueError( + "calibration loader yielded no batches; calibration needs a " + "non-empty validation/calibration split" + ) + exit_flops = [model._flops_at[ep.layer_name] for ep in model.config.exit_points] + return HeadLogits( + head_names=model.config.head_names, + logits=[torch.cat(chunks) for chunks in per_head], + targets=torch.cat(targets), + flops=torch.tensor(exit_flops + [1.0], dtype=torch.float64), + ) + + +# ---------------- stage 2: fit per-head temperatures ---------------- -# entropy grid: fractions of H_max = ln(num_classes), ordered conservative -> -# aggressive (higher entropy threshold fires more, since exit fires on H <= thr) -DEFAULT_ENTROPY_FRACTIONS: tuple[float, ...] = (0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7) + +def fit_head_temperatures(cache: HeadLogits) -> dict[str, TemperatureFitResult]: + """Fit one temperature per head (exits and final) from cached logits. + + Fitting is per-head NLL minimisation; each result carries its + convergence/fallback status. Callers decide what to do with fallbacks — + :func:`calibrate_thresholds` warns and keeps the safe 1.0. + """ + return { + name: fit_temperature_full(logits, cache.targets) + for name, logits in zip(cache.head_names, cache.logits) + } + + +# ---------------- stage 3/4: simulate + evaluate ---------------- + + +@dataclass(frozen=True) +class PolicyEvaluation: + """Routed metrics for one (thresholds, enabled, temperatures) setting.""" + + accuracy: float + avg_flops_fraction: float + exit_distribution: dict[str, float] def _evaluate_with_router( @@ -61,8 +192,8 @@ def _evaluate_with_router( ) -> tuple[float, float]: """Ground truth: route every sample through the real inference path. - Kept as the oracle that pins :class:`_EvalCache` correctness in the test - suite; calibration itself uses the cache.""" + Kept as the oracle that pins :func:`simulate_policy` correctness in the + test suite; calibration itself uses the cache.""" correct = 0 total = 0 comp = 0.0 @@ -75,110 +206,78 @@ def _evaluate_with_router( result = model(img.unsqueeze(0), mode="inference") pred = int(result.prediction.argmax(dim=-1).item()) correct += int(pred == int(tgt.item())) - comp += result.computation_used + comp += result.estimated_backbone_flops_fraction total += 1 return correct / max(total, 1), comp / max(total, 1) -class _EvalCache: - """One batched training-mode pass caches every head's logits; threshold - trials are then simulated against the cache instead of re-running the - network. +def simulate_policy(cache: HeadLogits, model: EarlyExitWrapper) -> PolicyEvaluation: + """Vectorially replay the router's decision rule against cached logits. - The simulation reproduces the router's math exactly: softmax in float32 - (as the hook computes it), criterion comparison in float64 (as the hook's - ``.item()`` produces), first firing exit claims the sample. A test pins - equivalence against :func:`_evaluate_with_router`. + Reproduces the routing math exactly: per-head temperature, softmax in + float32 (as the hook computes it), criterion comparison in float64 (as the + hook's ``.item()`` produces), first firing *enabled* exit claims the + sample. A test pins equivalence against :func:`_evaluate_with_router`. + """ + cfg = model.config + n_exits = len(cfg.exit_points) + n = cache.num_samples + chosen = torch.full((n,), n_exits, dtype=torch.long) + for idx in range(n_exits - 1, -1, -1): # earlier exits override later ones + if not cfg.enabled_exits[idx]: + continue + temp = _safe_temperature(cfg.temperatures[cfg.exit_points[idx].name]) + probs = torch.softmax(cache.logits[idx] / temp, dim=-1) + if cfg.routing_policy == "confidence": + confidence = probs.max(dim=-1).values.double() + fires = confidence >= cfg.confidence_thresholds[idx] + else: + entropy = -(probs * probs.clamp_min(1e-12).log()).sum(dim=-1) + fires = entropy.double() <= cfg.entropy_thresholds[idx] + chosen[fires] = idx + preds = torch.stack([h.argmax(dim=-1) for h in cache.logits]) + pred = preds[chosen, torch.arange(n)] + acc = (pred == cache.targets).double().mean().item() + comp = cache.flops[chosen].mean().item() + counts = torch.bincount(chosen, minlength=n_exits + 1) + dist = { + exit_label(i if i < n_exits else -1): counts[i].item() / n + for i in range(n_exits + 1) + if counts[i] > 0 + } + return PolicyEvaluation(accuracy=acc, avg_flops_fraction=comp, exit_distribution=dist) - Memory: stores ``(n_exits + 1) x N x num_classes`` logits on CPU, which is - fine for calibration-sized validation sets. - ponytail: turns grid search from O(exits * grid) full validation passes at - batch size 1 into one batched pass plus cheap tensor math. - """ +class _EvalCache: + """Adapter binding a :class:`HeadLogits` cache to a model for the greedy + search loops. ``evaluate()`` returns ``(accuracy, avg flops fraction)`` + under the model's *current* config.""" def __init__(self, model: EarlyExitWrapper, loader: DataLoader[Batch], device: str) -> None: self._model = model - per_head: list[list[torch.Tensor]] = [] - targets: list[torch.Tensor] = [] - model.eval() - with torch.no_grad(): - for images, tgt in loader: - outputs = model(images.to(device), mode="training") - if not per_head: - per_head = [[] for _ in outputs] - for idx, logits in enumerate(outputs): - per_head[idx].append(logits.detach().cpu()) - targets.append(tgt.cpu()) - if targets: - self._heads = [torch.cat(chunks) for chunks in per_head] - self._targets = torch.cat(targets) - else: # empty loader: preserve the old 0.0/0.0 contract - self._heads = [] - self._targets = torch.empty(0, dtype=torch.long) - exit_flops = [model._flops_at[ep.layer_name] for ep in model.config.exit_points] - self._flops = torch.tensor(exit_flops + [1.0], dtype=torch.float64) + self.cache = collect_head_logits(model, loader, device) def evaluate(self) -> tuple[float, float]: - """(accuracy, avg computation used) under the model's current config.""" - if self._targets.numel() == 0: - return 0.0, 0.0 - cfg = self._model.config - temp = _safe_temperature(cfg.temperature) - n_exits = len(cfg.exit_points) - n = self._targets.shape[0] - chosen = torch.full((n,), n_exits, dtype=torch.long) - for idx in range(n_exits - 1, -1, -1): # earlier exits override later ones - probs = torch.softmax(self._heads[idx] / temp, dim=-1) - if cfg.routing_policy == "confidence": - confidence = probs.max(dim=-1).values.double() - fires = confidence >= cfg.confidence_thresholds[idx] - else: - entropy = -(probs * probs.clamp_min(1e-12).log()).sum(dim=-1) - fires = entropy.double() <= cfg.entropy_thresholds[idx] - chosen[fires] = idx - preds = torch.stack([h.argmax(dim=-1) for h in self._heads]) - pred = preds[chosen, torch.arange(n)] - acc = (pred == self._targets).double().mean().item() - comp = self._flops[chosen].mean().item() - return acc, comp - - -def _collect_final_logits( - model: EarlyExitWrapper, loader: DataLoader[Batch], device: str -) -> tuple[torch.Tensor, torch.Tensor]: - """Run ``loader`` through the wrapper in training mode and harvest the - final classifier's logits + targets. Training mode never fires an exit (the - hook only appends logits and never raises), so ``outputs[-1]`` is always the - final classifier and no threshold reset is needed.""" - model = model.to(device) - all_logits: list[torch.Tensor] = [] - all_targets: list[torch.Tensor] = [] - model.eval() - with torch.no_grad(): - for images, targets in loader: - images = images.to(device) - targets = targets.to(device) - outputs = model(images, mode="training") - # outputs is [exit_0_logits, ..., exit_n_logits, final_logits] - all_logits.append(outputs[-1].detach()) - all_targets.append(targets.detach()) - return torch.cat(all_logits, dim=0), torch.cat(all_targets, dim=0) + ev = simulate_policy(self.cache, self._model) + return ev.accuracy, ev.avg_flops_fraction + + def evaluate_full(self) -> PolicyEvaluation: + return simulate_policy(self.cache, self._model) + + +# ---------------- search setup ---------------- @dataclass(frozen=True) class _PolicySearch: - """Per-policy search setup. ``seed`` is a strictly non-firing value used - while measuring the baseline and untried exits — it must NEVER fire even on - a float32-saturated softmax (max prob == 1.0 / entropy == -0.0), or the - no-exit baseline would be corrupted. ``no_exit`` is the in-range "disabled" - value the seed is clamped back to in the result. ``more_aggressive`` decides - which passing threshold to keep (smaller for confidence, larger for - entropy).""" + """Per-policy search setup. ``disabled_placeholder`` is the value recorded + in the threshold list for exits that end up disabled — routing never reads + it (``enabled_exits`` is authoritative), it only keeps the list fully + populated and in-range. ``more_aggressive`` decides which passing + threshold to keep (smaller for confidence, larger for entropy).""" field: str - seed: float - no_exit: float + disabled_placeholder: float grid: tuple[float, ...] more_aggressive: Callable[[float, float], bool] @@ -190,10 +289,11 @@ def _policy_search(model: EarlyExitWrapper, grid: tuple[float, ...] | None) -> _ raise ValueError("custom grid must be non-empty (or None for the default)") policy = model.config.routing_policy if policy == "confidence": + if grid is not None and any(not math.isfinite(t) or not 0.0 <= t <= 1.0 for t in grid): + raise ValueError(f"confidence grid values must lie in [0, 1]; got {grid}") return _PolicySearch( field="confidence_thresholds", - seed=2.0, # softmax.max() <= 1.0, so confidence >= 2.0 can never fire - no_exit=1.0, # documented "no early exit at this point" value + disabled_placeholder=1.0, grid=DEFAULT_GRID if grid is None else grid, more_aggressive=lambda new, cur: new < cur, ) @@ -202,7 +302,7 @@ def _policy_search(model: EarlyExitWrapper, grid: tuple[float, ...] | None) -> _ if grid is None: search_grid = tuple(round(f * h_max, 6) for f in DEFAULT_ENTROPY_FRACTIONS) else: - if any(t < 0 or t > h_max + 1e-9 for t in grid): + if any(not math.isfinite(t) or t < 0 or t > h_max + 1e-9 for t in grid): raise ValueError( f"entropy grid values must lie in [0, ln(num_classes)] = " f"[0, {h_max:.4f}]; got {grid}" @@ -210,8 +310,7 @@ def _policy_search(model: EarlyExitWrapper, grid: tuple[float, ...] | None) -> _ search_grid = grid return _PolicySearch( field="entropy_thresholds", - seed=-1.0, # entropy >= 0, so H <= -1.0 can never fire (even at -0.0) - no_exit=0.0, + disabled_placeholder=0.0, grid=search_grid, more_aggressive=lambda new, cur: new > cur, ) @@ -220,13 +319,13 @@ def _policy_search(model: EarlyExitWrapper, grid: tuple[float, ...] | None) -> _ ) -def _maybe_fit_temperature( +def _maybe_fit_temperatures( model: EarlyExitWrapper, fit_temperature: bool, temperature_loader: DataLoader[Batch] | None, val_loader: DataLoader[Batch], device: str, -) -> float | None: +) -> dict[str, TemperatureFitResult] | None: if not fit_temperature: return None if temperature_loader is None: @@ -240,18 +339,23 @@ def _maybe_fit_temperature( ) temperature_loader = val_loader - # Training-mode collection never fires exits (the hook only appends in - # training mode), so thresholds are irrelevant here; only the - # temperature must be neutral while harvesting uncalibrated logits. - original_temp = model.config.temperature - model.config.temperature = 1.0 - try: - logits, targets = _collect_final_logits(model, temperature_loader, device) - finally: - model.config.temperature = original_temp - fitted = _fit_temperature(logits.cpu(), targets.cpu()) - model.config.temperature = fitted - return fitted + cache = collect_head_logits(model, temperature_loader, device) + fits = fit_head_temperatures(cache) + bad = [name for name, fit in fits.items() if fit.fallback] + if bad: + warnings.warn( + f"temperature fitting diverged for head(s) {bad}; those heads keep " + "the last finite iterate (1.0 if none) and are effectively " + "uncalibrated. See TemperatureFitResult.fallback in the result.", + UserWarning, + stacklevel=3, + ) + model.config.temperatures = {name: fit.temperature for name, fit in fits.items()} + model.config.validate() + return fits + + +# ---------------- public calibrators ---------------- def calibrate_thresholds( @@ -275,63 +379,99 @@ def calibrate_thresholds( Default grid is :data:`DEFAULT_ENTROPY_FRACTIONS` scaled by ``ln(num_classes)``. - The result is written to whichever threshold list the active policy reads — - calibrating an entropy-routed model under the old confidence-only code was a - silent no-op. + Exits for which no grid value stays inside the accuracy budget are left + disabled in ``model.config.enabled_exits``. Calibration decides enablement + from scratch: any pre-existing enabled/disabled state is overwritten, + because the search's job is to select the deployment policy. + + ``val_loader`` is the *calibration* split. Do not pass your test set: the + threshold search consumes its labels. Evaluate the calibrated model on a + separate test loader with :func:`earlyon.benchmarking.evaluate`. Parameters ---------- + target_accuracy_drop: + Maximum acceptable drop (on the calibration split) from the + final-classifier baseline, in [0, 1]. grid: - Optional custom search grid. If ``None`` a policy-appropriate default is - used. For entropy a custom grid must lie in ``[0, ln(num_classes)]``. + Optional custom search grid. If ``None`` a policy-appropriate default + is used. Confidence grids must lie in [0, 1]; entropy grids in + ``[0, ln(num_classes)]``. fit_temperature: - If True, fit post-hoc temperature scaling before the threshold search. - The fitted scalar is written to ``model.config.temperature``. + If True, fit post-hoc per-head temperature scaling (every exit head + and the final classifier) before the threshold search. Fitted values + land on ``model.config.temperatures``. temperature_loader: - DataLoader used to fit the temperature. If ``None`` and + DataLoader used to fit the temperatures. If ``None`` and ``fit_temperature`` is True, ``val_loader`` is reused — this leaks the temperature fit into the threshold accuracy estimate, so a warning is emitted. Prefer a separate held-out split when available. """ + if not math.isfinite(target_accuracy_drop) or not 0.0 <= target_accuracy_drop <= 1.0: + raise ValueError(f"target_accuracy_drop must be in [0, 1], got {target_accuracy_drop}") model = model.to(device) - n = len(model.config.exit_points) + cfg = model.config + n = len(cfg.exit_points) search = _policy_search(model, grid) - fitted_t = _maybe_fit_temperature( - model, fit_temperature, temperature_loader, val_loader, device - ) + fits = _maybe_fit_temperatures(model, fit_temperature, temperature_loader, val_loader, device) cache = _EvalCache(model, val_loader, device) - # start fully conservative -- no exits will fire - best = [search.seed] * n - setattr(model.config, search.field, list(best)) + # baseline: every exit disabled — the final classifier's routed accuracy + best = [search.disabled_placeholder] * n + enabled = [False] * n + cfg.enabled_exits = list(enabled) + setattr(cfg, search.field, list(best)) baseline_acc, _ = cache.evaluate() # iterate the full grid per exit. an early break would miss thresholds # that pass after a transient miss (val accuracy isn't monotone in threshold - # on small sets). keep the most aggressive passing threshold seen. + # on small sets). keep the most aggressive passing threshold seen. A + # threshold that fires no calibration sample (compute unchanged vs. the + # exit disabled) is useless — the exit stays disabled instead of shipping + # an enabled-but-idle exit that might fire on out-of-distribution inputs. for exit_idx in range(n): + setattr(cfg, search.field, list(best)) + cfg.enabled_exits = list(enabled) + _, comp_without = cache.evaluate() + chosen: float | None = None for thr in search.grid: trial = list(best) trial[exit_idx] = thr - setattr(model.config, search.field, list(trial)) - acc, _ = cache.evaluate() - if baseline_acc - acc <= target_accuracy_drop and search.more_aggressive( - thr, best[exit_idx] + trial_enabled = list(enabled) + trial_enabled[exit_idx] = True + setattr(cfg, search.field, list(trial)) + cfg.enabled_exits = trial_enabled + acc, comp = cache.evaluate() + useful = comp < comp_without + if ( + useful + and baseline_acc - acc <= target_accuracy_drop + and (chosen is None or search.more_aggressive(thr, chosen)) ): - best[exit_idx] = thr - - # clamp any exit that never found a passing threshold from the non-firing - # search seed back to the in-range disabled value. - best = [search.no_exit if t == search.seed else t for t in best] - setattr(model.config, search.field, list(best)) - final_acc, avg_comp = cache.evaluate() + chosen = thr + if chosen is not None: + best[exit_idx] = chosen + enabled[exit_idx] = True + + setattr(cfg, search.field, list(best)) + cfg.enabled_exits = list(enabled) + cfg.validate() + final = cache.evaluate_full() return CalibrationResult( thresholds=best, + enabled_exits=list(enabled), baseline_accuracy=baseline_acc, - final_accuracy=final_acc, - avg_computation_used=avg_comp, - fitted_temperature=fitted_t, - policy=model.config.routing_policy, + final_accuracy=final.accuracy, + avg_computation_used=final.avg_flops_fraction, + accuracy_delta=baseline_acc - final.accuracy, + exit_distribution=final.exit_distribution, + num_samples=cache.cache.num_samples, + fitted_temperature=fits[FINAL_HEAD].temperature if fits else None, + temperatures=dict(cfg.temperatures) if fits else None, + temperature_fits=fits, + policy=cfg.routing_policy, + objective="accuracy_budget", + target_accuracy_drop=target_accuracy_drop, ) @@ -348,38 +488,39 @@ def calibrate_thresholds_for_budget( The mirror image of :func:`calibrate_thresholds`: there you state an accuracy budget and get the biggest compute saving inside it; here you - state a compute budget — ``target_computation``, the average fraction of - the backbone's FLOPs the deployed model may run per sample, in ``(0, 1]`` - — and the search keeps validation accuracy as high as it can while - meeting it. + state a compute budget — ``target_computation``, the average *estimated* + fraction of the backbone's FLOPs the deployed model may run per sample, in + ``(0, 1]`` — and the search keeps validation accuracy as high as it can + while meeting it. Greedy per exit, earliest first (early exits save the most compute). At - each exit, among grid values whose average compute is within budget, the - one with the highest validation accuracy is kept and the search stops; - later exits stay disabled. If no value at this exit reaches the budget, - the value with the largest strict compute reduction is kept and the search - moves to the next exit; an exit that reduces nothing stays disabled. If - the budget is unattainable after all exits, a ``UserWarning`` is emitted - and the result carries ``budget_met=False`` with the least-compute - configuration found (which may be the plain backbone). + each exit, among grid values whose average estimated compute is within + budget, the one with the highest validation accuracy is kept and the + search stops; later exits stay disabled. If no value at this exit reaches + the budget, the value with the largest strict compute reduction is kept + and the search moves to the next exit; an exit that reduces nothing stays + disabled. If the budget is unattainable after all exits, a ``UserWarning`` + is emitted and the result carries ``budget_met=False`` with the + least-compute configuration found (which may be the plain backbone). Routing-policy aware exactly like :func:`calibrate_thresholds`, and the ``grid`` / ``fit_temperature`` / ``temperature_loader`` parameters behave identically. """ - if not 0.0 < target_computation <= 1.0: + if not math.isfinite(target_computation) or not 0.0 < target_computation <= 1.0: raise ValueError(f"target_computation must be in (0, 1], got {target_computation}") model = model.to(device) - n = len(model.config.exit_points) + cfg = model.config + n = len(cfg.exit_points) search = _policy_search(model, grid) - fitted_t = _maybe_fit_temperature( - model, fit_temperature, temperature_loader, val_loader, device - ) + fits = _maybe_fit_temperatures(model, fit_temperature, temperature_loader, val_loader, device) cache = _EvalCache(model, val_loader, device) # start fully conservative -- no exits fire, compute is the full backbone - best = [search.seed] * n - setattr(model.config, search.field, list(best)) + best = [search.disabled_placeholder] * n + enabled = [False] * n + cfg.enabled_exits = list(enabled) + setattr(cfg, search.field, list(best)) baseline_acc, current_comp = cache.evaluate() budget_met = current_comp <= target_computation @@ -394,7 +535,10 @@ def calibrate_thresholds_for_budget( for thr in search.grid: trial = list(best) trial[exit_idx] = thr - setattr(model.config, search.field, list(trial)) + trial_enabled = list(enabled) + trial_enabled[exit_idx] = True + setattr(cfg, search.field, list(trial)) + cfg.enabled_exits = trial_enabled acc, comp = cache.evaluate() if comp <= target_computation and (within is None or acc > within[0]): within = (acc, thr, comp) @@ -402,36 +546,44 @@ def calibrate_thresholds_for_budget( fallback = (comp, thr) if within is not None: best[exit_idx] = within[1] + enabled[exit_idx] = True current_comp = within[2] budget_met = True elif fallback is not None: best[exit_idx] = fallback[1] + enabled[exit_idx] = True current_comp = fallback[0] - # else: no grid value at this exit reduced compute; leave it at the - # seed so it clamps to the disabled value below + # else: no grid value at this exit reduced compute; leave it disabled if not budget_met: warnings.warn( f"target_computation={target_computation} is unattainable: the " f"least-compute configuration found still uses {current_comp:.4f} " - "of the backbone's FLOPs on the validation set. Returning it with " - "budget_met=False; consider earlier exit points or a larger budget.", + "of the backbone's estimated FLOPs on the validation set. Returning " + "it with budget_met=False; consider earlier exit points or a larger " + "budget.", UserWarning, stacklevel=2, ) - # clamp exits still at the non-firing search seed to the in-range - # disabled value. - best = [search.no_exit if t == search.seed else t for t in best] - setattr(model.config, search.field, list(best)) - final_acc, avg_comp = cache.evaluate() + setattr(cfg, search.field, list(best)) + cfg.enabled_exits = list(enabled) + cfg.validate() + final = cache.evaluate_full() return CalibrationResult( thresholds=best, + enabled_exits=list(enabled), baseline_accuracy=baseline_acc, - final_accuracy=final_acc, - avg_computation_used=avg_comp, - fitted_temperature=fitted_t, - policy=model.config.routing_policy, + final_accuracy=final.accuracy, + avg_computation_used=final.avg_flops_fraction, + accuracy_delta=baseline_acc - final.accuracy, + exit_distribution=final.exit_distribution, + num_samples=cache.cache.num_samples, + fitted_temperature=fits[FINAL_HEAD].temperature if fits else None, + temperatures=dict(cfg.temperatures) if fits else None, + temperature_fits=fits, + policy=cfg.routing_policy, + objective="compute_budget", target_computation=target_computation, budget_met=budget_met, ) diff --git a/earlyon/core/types.py b/earlyon/core/types.py index ad2b267..f54eb92 100644 --- a/earlyon/core/types.py +++ b/earlyon/core/types.py @@ -10,10 +10,16 @@ # A (images, targets) minibatch — the element type of every earlyon DataLoader. Batch = tuple[torch.Tensor, torch.Tensor] +# Reserved key for the final classifier head in per-head mappings +# (``EarlyExitConfig.temperatures``). Exit points may not use this name. +FINAL_HEAD = "final" + +VALID_ROUTING_POLICIES = ("confidence", "entropy") + def exit_label(idx: int) -> str: """Canonical name for an exit index: ``"final"`` for -1, else ``"exit_{idx}"``.""" - return "final" if idx == -1 else f"exit_{idx}" + return FINAL_HEAD if idx == -1 else f"exit_{idx}" @dataclass(frozen=True) @@ -41,6 +47,19 @@ class EarlyExitConfig: Uses ``confidence_thresholds``. * ``"entropy"`` — exit when ``H(softmax(logits)) <= threshold`` (low entropy = high confidence). Uses ``entropy_thresholds``. + + ``enabled_exits`` is the explicit per-exit on/off switch. A disabled exit + never fires, regardless of thresholds — including the numerical edge cases + (softmax confidence saturated at exactly 1.0, entropy at exactly 0.0) that + made the old "threshold sentinel means disabled" convention unsound. + + ``temperatures`` maps each head — every ``ExitPoint.name`` plus the + reserved key ``"final"`` — to its softmax temperature (Guo et al. 2017). + Per-head temperatures matter because each exit head is a different + classifier with its own miscalibration. The legacy scalar ``temperature`` + field is a constructor convenience only: it is broadcast into + ``temperatures`` at construction time and ignored afterwards; mutate + ``temperatures`` (not ``temperature``) to change calibration. """ backbone: str @@ -49,14 +68,12 @@ class EarlyExitConfig: loss_weights: list[float] = field(default_factory=list) confidence_thresholds: list[float] = field(default_factory=list) entropy_thresholds: list[float] = field(default_factory=list) + enabled_exits: list[bool] = field(default_factory=list) routing_policy: str = "confidence" - temperature: float = 1.0 + temperature: float = 1.0 # legacy scalar; broadcast into `temperatures` + temperatures: dict[str, float] = field(default_factory=dict) def __post_init__(self) -> None: - if self.num_classes < 2: - # with a single class softmax confidence is identically 1.0, so - # every exit fires at any threshold <= 1.0 — routing is meaningless - raise ValueError(f"num_classes must be >= 2 (got {self.num_classes})") n_exits = len(self.exit_points) if not self.loss_weights: # equal weight across exits + final @@ -68,34 +85,138 @@ def __post_init__(self) -> None: # half the max possible entropy for `num_classes`-way uniform default = 0.5 * math.log(max(self.num_classes, 2)) self.entropy_thresholds = [default] * n_exits + if not self.enabled_exits: + self.enabled_exits = [True] * n_exits + if not self.temperatures: + # migrate the legacy scalar deterministically: same T on every head + self.temperatures = {ep.name: self.temperature for ep in self.exit_points} + self.temperatures[FINAL_HEAD] = self.temperature + elif self.temperature != 1.0: + raise ValueError( + "pass either the legacy scalar `temperature` or the per-head " + "`temperatures` mapping, not both" + ) + self.validate() + + # ---------------- helpers ---------------- + + @property + def head_names(self) -> list[str]: + """Every head that produces logits: exit names in order, then ``"final"``.""" + return [ep.name for ep in self.exit_points] + [FINAL_HEAD] + + def temperature_for(self, head_name: str) -> float: + """The softmax temperature for ``head_name`` (an exit name or ``"final"``).""" + return self.temperatures[head_name] + + def validate(self) -> None: + """Check every configuration invariant; raise ``ValueError`` with an + actionable message on the first violation. + + Called automatically at construction. Because this dataclass is mutable + (calibration writes fitted values back), ``EarlyExitWrapper`` and the + checkpoint loader re-invoke it at their own trust boundaries. + """ + if self.num_classes < 2: + # with a single class softmax confidence is identically 1.0, so + # every exit fires at any threshold <= 1.0 — routing is meaningless + raise ValueError(f"num_classes must be >= 2 (got {self.num_classes})") + + n_exits = len(self.exit_points) + names = [ep.name for ep in self.exit_points] + layers = [ep.layer_name for ep in self.exit_points] + if any(not n for n in names) or any(not layer for layer in layers): + raise ValueError("exit point names and layer names must be non-empty") + if len(set(names)) != n_exits: + raise ValueError(f"exit point names must be unique (got {names})") + if len(set(layers)) != n_exits: + raise ValueError(f"exit point layer names must be unique (got {layers})") + if FINAL_HEAD in names: + raise ValueError(f"exit point name {FINAL_HEAD!r} is reserved for the final head") + + if self.routing_policy not in VALID_ROUTING_POLICIES: + raise ValueError( + f"routing_policy={self.routing_policy!r} not supported " + f"(allowed: {', '.join(repr(p) for p in VALID_ROUTING_POLICIES)})" + ) + if len(self.loss_weights) != n_exits + 1: raise ValueError( f"loss_weights must have length {n_exits + 1} (got {len(self.loss_weights)})" ) + if any(not math.isfinite(w) or w < 0 for w in self.loss_weights): + raise ValueError(f"loss_weights must be finite and >= 0 (got {self.loss_weights})") + if sum(self.loss_weights) <= 0: + raise ValueError("loss_weights must not all be zero") + if len(self.confidence_thresholds) != n_exits: raise ValueError( f"confidence_thresholds must have length {n_exits} " f"(got {len(self.confidence_thresholds)})" ) + if any(not math.isfinite(t) or not 0.0 <= t <= 1.0 for t in self.confidence_thresholds): + raise ValueError( + "confidence_thresholds must be finite and in [0, 1] " + f"(got {self.confidence_thresholds}); to disable an exit set " + "enabled_exits[i] = False instead of using a sentinel threshold" + ) if len(self.entropy_thresholds) != n_exits: raise ValueError( f"entropy_thresholds must have length {n_exits} " f"(got {len(self.entropy_thresholds)})" ) - if self.routing_policy not in {"confidence", "entropy"}: + if any(not math.isfinite(t) or t < 0.0 for t in self.entropy_thresholds): raise ValueError( - f"routing_policy={self.routing_policy!r} not supported " - "(allowed: 'confidence', 'entropy')" + f"entropy_thresholds must be finite and >= 0 (got {self.entropy_thresholds}); " + "to disable an exit set enabled_exits[i] = False instead of a sentinel" ) + if len(self.enabled_exits) != n_exits: + raise ValueError( + f"enabled_exits must have length {n_exits} (got {len(self.enabled_exits)})" + ) + if any(not isinstance(e, bool) for e in self.enabled_exits): + raise ValueError(f"enabled_exits must be booleans (got {self.enabled_exits})") + + expected_keys = set(names) | {FINAL_HEAD} + if set(self.temperatures) != expected_keys: + raise ValueError( + f"temperatures must have exactly one entry per head " + f"({sorted(expected_keys)}); got keys {sorted(self.temperatures)}" + ) + for head, t in self.temperatures.items(): + if not math.isfinite(t) or t <= 0.0: + raise ValueError( + f"temperatures[{head!r}] must be finite and > 0 (got {t}). " + "A fitted temperature this bad means calibration diverged; " + "refit or fall back to 1.0 explicitly." + ) + + +def migrate_scalar_temperature(config: EarlyExitConfig, temperature: float) -> None: + """Deterministically broadcast a legacy scalar temperature onto every head. + + Used by the v1-checkpoint migration path; emits no warning itself (the + caller decides how loudly to announce the migration). + """ + config.temperatures = {ep.name: temperature for ep in config.exit_points} + config.temperatures[FINAL_HEAD] = temperature + config.validate() + @dataclass class InferenceResult: """Result of a single-sample inference pass. ``exit_taken`` is the 0-indexed exit number; -1 means the final classifier - was used. ``computation_used`` is the FLOPs fraction (not layer count) of - the network actually executed. + was used. + + ``estimated_backbone_flops_fraction`` is an *estimate* of the fraction of + the backbone's FLOPs executed, derived from a one-time static analysis of + the backbone (see ``earlyon.core.flops``). It is not a measurement of this + inference, and it excludes the exit heads' own (small) cost and all + routing overhead. ``computation_used`` remains as a read alias for + backward compatibility. ``prediction`` is produced under ``torch.inference_mode()`` (the deployment path is grad-free by contract), so it is an *inference tensor*: ready for @@ -107,7 +228,12 @@ class InferenceResult: prediction: torch.Tensor exit_taken: int confidence: float - computation_used: float + estimated_backbone_flops_fraction: float + + @property + def computation_used(self) -> float: + """Deprecated alias for ``estimated_backbone_flops_fraction``.""" + return self.estimated_backbone_flops_fraction @dataclass @@ -116,10 +242,16 @@ class BatchedInferenceResult: All samples in the batch exit together at ``exit_taken``. The ``per_sample_confidence`` tensor records each sample's confidence at the - exit point. + exit point. ``estimated_backbone_flops_fraction`` follows the same + estimator contract as :class:`InferenceResult`. """ predictions: torch.Tensor exit_taken: int per_sample_confidence: torch.Tensor - computation_used: float + estimated_backbone_flops_fraction: float + + @property + def computation_used(self) -> float: + """Deprecated alias for ``estimated_backbone_flops_fraction``.""" + return self.estimated_backbone_flops_fraction diff --git a/earlyon/core/wrappers.py b/earlyon/core/wrappers.py index e943747..e0f422f 100644 --- a/earlyon/core/wrappers.py +++ b/earlyon/core/wrappers.py @@ -10,30 +10,33 @@ import math import threading +import warnings from typing import TYPE_CHECKING, Any, Callable, Iterator, Mapping, Optional import torch import torch.nn as nn -from earlyon.core.flops import per_layer_flops -from earlyon.core.types import EarlyExitConfig, InferenceResult +from earlyon.core.flops import FlopsEstimate, estimate_layer_flops +from earlyon.core.types import FINAL_HEAD, EarlyExitConfig, InferenceResult if TYPE_CHECKING: from earlyon.core.types import BatchedInferenceResult def _safe_temperature(temperature: float) -> float: - """Clamp the softmax temperature to a strictly positive, finite value. - - A NaN/Inf temperature (e.g. from a diverged ``fit_temperature`` on degenerate - logits) would otherwise poison every softmax: ``max(nan, 1e-6)`` returns - ``nan`` because Python ``max`` keeps its first argument on a False compare, - and ``softmax(logits / nan)`` is silently all-NaN. Falling back to 1.0 - (the no-op temperature) keeps routing well-defined. + """Defense-in-depth guard for the softmax temperature at routing time. + + ``EarlyExitConfig.validate()`` rejects non-finite or non-positive + temperatures at every trust boundary (construction, wrapping, checkpoint + load), so a bad value can only appear here through direct post-hoc + mutation of ``config.temperatures``. In that case fall back to 1.0 (the + no-op temperature): a NaN would silently poison every softmax, and the + old ``max(t, 1e-6)`` clamp turned an invalid negative temperature into an + artificially razor-sharp softmax — worse than no calibration at all. """ - if not math.isfinite(temperature): + if not math.isfinite(temperature) or temperature <= 0.0: return 1.0 - return max(temperature, 1e-6) + return temperature def _is_compiling() -> bool: @@ -92,19 +95,44 @@ def __init__( final_classifier: Callable[[torch.Tensor], torch.Tensor], config: EarlyExitConfig, input_shape: tuple[int, int, int, int] = (1, 3, 224, 224), + feature_adapters: Optional[Mapping[str, Callable[[Any], torch.Tensor]]] = None, ) -> None: super().__init__() self.backbone = backbone self.exit_heads = nn.ModuleDict(exit_heads) self._final_classifier = final_classifier + # re-validate at this trust boundary: the config is mutable and may + # have been edited between construction and wrapping. + config.validate() + missing_heads = [ep.name for ep in config.exit_points if ep.name not in exit_heads] + extra_heads = [n for n in exit_heads if all(ep.name != n for ep in config.exit_points)] + if missing_heads or extra_heads: + raise ValueError( + "exit_heads must have exactly one head per configured exit point; " + f"missing: {missing_heads}, unexpected: {extra_heads}" + ) self.config = config + # optional per-exit callables mapping a layer's raw output (tuple / + # dict / token tensor / ...) to the Tensor its head consumes. Keyed by + # exit name; NOT part of the state_dict — factories must rebuild them. + unknown_adapters = [ + n for n in (feature_adapters or {}) if all(ep.name != n for ep in config.exit_points) + ] + if unknown_adapters: + raise ValueError(f"feature_adapters keyed by unknown exit name(s): {unknown_adapters}") + self._feature_adapters: dict[str, Callable[[Any], torch.Tensor]] = dict( + feature_adapters or {} + ) # ordered list of (exit_idx, exit_name, layer_name) self._exits = [(i, ep.name, ep.layer_name) for i, ep in enumerate(config.exit_points)] - # cumulative FLOPs fraction at each exit layer - layer_names = [ep.layer_name for ep in config.exit_points] - self._flops_at: dict[str, float] = per_layer_flops(backbone, layer_names, input_shape) + # cumulative estimated FLOPs fraction at each exit layer — computed + # LAZILY on first use (the analysis runs forward passes of the + # backbone; doing it here made construction of large models like ViT + # unacceptably slow). See `flops_estimate` / `_flops_at`. + self._input_shape = input_shape + self._flops_estimate: Optional["FlopsEstimate"] = None # per-thread caches: a single wrapper instance can be called from # multiple threads (DataParallel, inference servers). state lives in @@ -124,6 +152,50 @@ def __init__( # ---------------- public api ---------------- + @property + def flops_estimate(self) -> FlopsEstimate: + """The lazy per-exit FLOPs estimate, with estimator metadata. + + Computed on first access (one/two backbone forward passes) and cached. + Must be materialised *outside* an active routing pass: the analysis + runs the backbone directly, and the hooks stay no-ops only while this + thread's inference/training flags are unset — which is why the forward + paths call it before flipping those flags. + """ + if self._flops_estimate is None: + layer_names = [ep.layer_name for ep in self.config.exit_points] + # the probe input follows the backbone wherever it currently lives + device = next(self.backbone.parameters(), torch.empty(0)).device + try: + self._flops_estimate = estimate_layer_flops( + self.backbone, + layer_names, + self._input_shape, + device=str(device), + ) + except Exception as exc: # backbone can't run from zeros(input_shape) + from earlyon.core.flops import METHOD_UNIFORM + + warnings.warn( + f"FLOPs analysis failed ({type(exc).__name__}: {exc}); using a " + "low-confidence uniform estimate (FlopsEstimate.reliable=False)", + RuntimeWarning, + stacklevel=2, + ) + n = len(layer_names) + self._flops_estimate = FlopsEstimate( + fractions={nm: (i + 1) / (n + 1) for i, nm in enumerate(layer_names)}, + method=METHOD_UNIFORM, + reliable=False, + notes=(f"analysis failed: {type(exc).__name__}",), + ) + return self._flops_estimate + + @property + def _flops_at(self) -> dict[str, float]: + """Cumulative estimated FLOPs fraction per exit layer (lazy).""" + return self.flops_estimate.fractions + def forward( self, x: torch.Tensor, mode: str = "inference" ) -> InferenceResult | list[torch.Tensor]: @@ -172,6 +244,7 @@ def forward_inference_batched(self, x: torch.Tensor) -> "BatchedInferenceResult" if x.size(0) < 1: raise ValueError("empty batch") + _ = self.flops_estimate # materialise BEFORE the routing flags flip self._init_tls() self._tls.inference_mode = True self._tls.batched_mode = True @@ -195,18 +268,18 @@ def forward_inference_batched(self, x: torch.Tensor) -> "BatchedInferenceResult" predictions=sig.prediction, exit_taken=sig.exit_idx, per_sample_confidence=sig.per_sample_confidence, - computation_used=self._flops_at[layer_name], + estimated_backbone_flops_fraction=self._flops_at[layer_name], ) # no exit fired final_logits = self._final_classifier(feats) - temp = _safe_temperature(self.config.temperature) + temp = _safe_temperature(self.config.temperatures[FINAL_HEAD]) probs = torch.softmax(final_logits / temp, dim=-1) per_sample_conf = probs.max(dim=-1).values return BatchedInferenceResult( predictions=final_logits, exit_taken=-1, per_sample_confidence=per_sample_conf, - computation_used=1.0, + estimated_backbone_flops_fraction=1.0, ) finally: self._tls.inference_mode = False @@ -243,13 +316,22 @@ def hook(module: nn.Module, inputs: Any, output: torch.Tensor) -> torch.Tensor: "without exit heads for compiled inference." ) + # a disabled exit never routes — not even at softmax confidence + # exactly 1.0 or entropy exactly 0.0. In training mode the head + # still produces logits (enablement is a routing concept; disabled + # heads keep training so they can be re-enabled later). + if not in_training and not self.config.enabled_exits[exit_idx]: + return output + + adapter = self._feature_adapters.get(exit_name) + features = adapter(output) if adapter is not None else output head = self.exit_heads[exit_name] - logits = head(output) + logits = head(features) if in_training: getattr(self._tls, "training_outputs", []).append(logits) return output - temp = _safe_temperature(self.config.temperature) + temp = _safe_temperature(self.config.temperatures[exit_name]) probs = torch.softmax(logits / temp, dim=-1) policy = self.config.routing_policy @@ -309,6 +391,7 @@ def _forward_training(self, x: torch.Tensor) -> list[torch.Tensor]: self._tls.in_training = False def _forward_inference(self, x: torch.Tensor) -> InferenceResult: + _ = self.flops_estimate # materialise BEFORE the routing flags flip self._init_tls() self._tls.inference_mode = True try: @@ -323,19 +406,19 @@ def _forward_inference(self, x: torch.Tensor) -> InferenceResult: prediction=sig.prediction, exit_taken=sig.exit_idx, confidence=sig.confidence, - computation_used=self._flops_at[layer_name], + estimated_backbone_flops_fraction=self._flops_at[layer_name], ) # no exit triggered — use final classifier (inside outer try so # any exception still resets inference_mode) final_logits = self._final_classifier(feats) - temp = _safe_temperature(self.config.temperature) + temp = _safe_temperature(self.config.temperatures[FINAL_HEAD]) probs = torch.softmax(final_logits / temp, dim=-1) confidence = probs.max().item() return InferenceResult( prediction=final_logits, exit_taken=-1, confidence=confidence, - computation_used=1.0, + estimated_backbone_flops_fraction=1.0, ) finally: self._tls.inference_mode = False diff --git a/earlyon/models/custom.py b/earlyon/models/custom.py index b4a72fe..de65798 100644 --- a/earlyon/models/custom.py +++ b/earlyon/models/custom.py @@ -4,17 +4,27 @@ hooks, so callers don't hand-specify channel counts. This is the ``earlyon.models`` entry point for arbitrary user-provided backbones. -Custom-wrapped models are NOT round-trippable through ``save_wrapper`` / -``load_wrapper`` (``build_model`` cannot reconstruct an arbitrary backbone from a -string); ``build_model`` raises a clear error for ``backbone="custom"``. If you -save/load the ``state_dict`` yourself, rebuild the heads with the same -``pool_tokens`` you trained with — it is not stored in the weights, and a -mismatch loads silently while changing how 3D token features are pooled. +The dry run also *validates* the wrapping: every requested exit layer must +execute exactly once, in the order the caller listed them — a reused layer +or an out-of-order exit list would silently mis-route at inference, so both +fail loudly here instead. + +Runtime contract: the wrapper's inference/training paths call +``backbone(x)`` with a single tensor. ``example_args``/``example_kwargs`` +configure the *inspection* forward only — use them when the backbone needs +extra arguments that have defaults at runtime, or a non-default example +tensor. The example (or ``input_shape``) tensor is created on the backbone's +own device unless ``device`` says otherwise. + +Custom-wrapped models are reloadable only via a user-provided factory: +``load_wrapper(path, factory=lambda: custom_ee(...))`` with the same +structure (including ``pool_tokens`` and any ``feature_extractors`` — neither +is stored in the weights). """ from __future__ import annotations -from typing import Callable, Sequence +from typing import Any, Callable, Mapping, Sequence import torch import torch.nn as nn @@ -24,38 +34,73 @@ from earlyon.core.wrappers import EarlyExitWrapper from earlyon.models._common import identity +FeatureExtractor = Callable[[Any], torch.Tensor] + + +def _infer_backbone_device(backbone: nn.Module) -> torch.device: + for p in backbone.parameters(): + return p.device + for b in backbone.buffers(): + return b.device + return torch.device("cpu") + -def _infer_feature_widths( +def _to_device(obj: Any, device: torch.device) -> Any: + if isinstance(obj, torch.Tensor): + return obj.to(device) + return obj + + +class _DryRunReport: + def __init__(self) -> None: + self.widths: dict[str, int] = {} + self.execution_order: list[str] = [] + self.call_counts: dict[str, int] = {} + + +def _dry_run( backbone: nn.Module, layer_names: Sequence[str], - input_shape: tuple[int, int, int, int], -) -> dict[str, int]: - """Run one no-grad dry-run forward with temporary hooks to record each exit - layer's feature width (``C`` for 4D output, ``D`` for 3D/2D). - - Raises ``ValueError`` if a layer produces a non-Tensor or unsupported rank, - ``RuntimeError`` if a layer is never visited during the forward pass, and - ``ValueError`` naming the available layers for an unresolvable layer name. - Restores the backbone's prior train/eval mode on exit. + example_args: tuple[Any, ...], + example_kwargs: dict[str, Any], + extractors: Mapping[str, FeatureExtractor], +) -> _DryRunReport: + """One no-grad forward with temporary hooks: infer each exit layer's + feature width, record execution order and per-layer call counts. + + Raises ``ValueError`` for an unresolvable layer name, a non-Tensor output + without an extractor, or an unsupported rank; ``RuntimeError`` for layers + never visited, visited more than once, or visited out of order. Restores + the backbone's prior train/eval mode on exit. """ - widths: dict[str, int] = {} + report = _DryRunReport() handles: list[torch.utils.hooks.RemovableHandle] = [] def _make_probe(name: str) -> Callable[[nn.Module, object, object], None]: def hook(module: nn.Module, inputs: object, output: object) -> None: - if not isinstance(output, torch.Tensor): + report.call_counts[name] = report.call_counts.get(name, 0) + 1 + report.execution_order.append(name) + extractor = extractors.get(name) + features = extractor(output) if extractor is not None else output + if not isinstance(features, torch.Tensor): + hint = ( + "its feature_extractors entry must return a Tensor" + if extractor is not None + else "pass feature_extractors={name: fn} to convert it to a Tensor" + ) raise ValueError( - f"exit layer {name!r} produced {type(output).__name__}, not a " - "Tensor — custom_ee cannot auto-infer its feature width" + f"exit layer {name!r} produced {type(features).__name__}, not a " + f"Tensor — {hint}" ) - if output.dim() == 4: - widths[name] = output.shape[1] # (B, C, H, W) -> C - elif output.dim() in (2, 3): - widths[name] = output.shape[-1] # (B, D) / (B, N, D) -> D + if features.dim() == 4: + report.widths[name] = features.shape[1] # (B, C, H, W) -> C + elif features.dim() in (2, 3): + report.widths[name] = features.shape[-1] # (B, D) / (B, N, D) -> D else: raise ValueError( - f"exit layer {name!r} output has unsupported rank {output.dim()} " - f"(shape {tuple(output.shape)}); supported: 2D, 3D, or 4D" + f"exit layer {name!r} features have unsupported rank {features.dim()} " + f"(shape {tuple(features.shape)}); supported: 2D, 3D, or 4D. " + "Use feature_extractors to reshape." ) return hook @@ -75,19 +120,32 @@ def hook(module: nn.Module, inputs: object, output: object) -> None: backbone.eval() try: with torch.no_grad(): - backbone(torch.zeros(input_shape)) + backbone(*example_args, **example_kwargs) finally: for handle in handles: handle.remove() backbone.train(was_training) - missing = [name for name in layer_names if name not in widths] + missing = [name for name in layer_names if name not in report.widths] if missing: raise RuntimeError( f"dry-run forward never visited exit layers {missing}; verify they are " - f"reachable for input_shape={input_shape}" + "reachable for the example input" + ) + reused = [name for name, c in report.call_counts.items() if c > 1] + if reused: + raise RuntimeError( + f"exit layer(s) {reused} execute more than once per forward pass " + "(module reuse); early-exit routing at a reused layer is ambiguous " + "and unsupported" + ) + if report.execution_order != list(layer_names): + raise RuntimeError( + f"exit_layers must be listed in forward-execution order; the dry run " + f"observed {report.execution_order} but you passed {list(layer_names)}. " + "Routing short-circuits at the first confident exit, so order matters." ) - return widths + return report def custom_ee( @@ -96,6 +154,10 @@ def custom_ee( num_classes: int, *, input_shape: tuple[int, int, int, int] = (1, 3, 224, 224), + example_args: tuple[Any, ...] | None = None, + example_kwargs: dict[str, Any] | None = None, + device: str | torch.device | None = None, + feature_extractors: Mapping[str, FeatureExtractor] | None = None, hidden_dim: int = 128, dropout: float = 0.2, pool_tokens: str = "mean", @@ -104,26 +166,39 @@ def custom_ee( ) -> EarlyExitWrapper: """Wrap any ``nn.Module`` with early exits at the named submodules. - Feature widths are auto-inferred from a single dry-run forward, so you only - name the layers. ``exit_layers`` order is the exit order — the names must be - in forward-execution order (``exit_layers[0]`` runs first), since routing - short-circuits at the first confident exit. + Feature widths are auto-inferred from a single dry-run forward, which also + validates that each exit layer runs exactly once and in the listed order. The caller is responsible for making ``backbone(x)`` return - ``(B, num_classes)`` logits (e.g. replace the backbone's final head before - wrapping); ``final_classifier`` is identity. + ``(B, num_classes)`` logits at runtime (e.g. replace the backbone's final + head before wrapping); ``final_classifier`` is identity. Parameters ---------- backbone: - Any ``nn.Module`` callable on ``input_shape``-shaped input. + Any ``nn.Module``. The dry run calls it on the example input; at + runtime the wrapper calls ``backbone(x)`` with a single tensor. exit_layers: Dotted submodule names resolvable via ``backbone.get_submodule``, in - forward order. + forward-execution order (validated by the dry run). num_classes: Output classes for every exit head. input_shape: - Shape for the dry-run forward and FLOPs accounting. + Shape for the default dry-run tensor and the FLOPs estimate. Ignored + for the dry run when ``example_args`` is given. + example_args / example_kwargs: + Explicit example inputs for the dry run — for backbones needing more + than ``zeros(input_shape)`` (extra positional/keyword arguments with + runtime defaults, integer token ids, ...). Tensors are moved to the + inspection device. + device: + Where to run the dry run. Default: the backbone's own parameter + device — a CUDA model is no longer probed with a CPU tensor. + feature_extractors: + Optional per-*layer-name* callables converting a layer's raw output + (tuple, dict, unusual rank) into the 2D/3D/4D Tensor its exit head + consumes. Also applied at inference; not serialized — rebuild them in + your ``load_wrapper`` factory. pool_tokens: For 3D token features, ``"mean"`` (default, architecture-agnostic) or ``"cls"`` (token 0). Ignored for 4D conv features. @@ -135,8 +210,26 @@ def custom_ee( if len(exit_layers) == 0: raise ValueError("exit_layers must be non-empty") - widths = _infer_feature_widths(backbone, exit_layers, input_shape) - exit_points = [ExitPoint(f"e{i}", name, widths[name]) for i, name in enumerate(exit_layers)] + run_device = torch.device(device) if device is not None else _infer_backbone_device(backbone) + if example_args is None and example_kwargs: + raise ValueError("example_kwargs requires example_args (may be an empty tuple)") + if example_args is None: + example_args = (torch.zeros(input_shape, device=run_device),) + example_kwargs = {} + else: + example_args = tuple(_to_device(a, run_device) for a in example_args) + example_kwargs = {k: _to_device(v, run_device) for k, v in (example_kwargs or {}).items()} + + extractors = dict(feature_extractors or {}) + unknown = [k for k in extractors if k not in set(exit_layers)] + if unknown: + raise ValueError(f"feature_extractors keyed by unknown exit layer(s): {unknown}") + + report = _dry_run(backbone, exit_layers, example_args, example_kwargs, extractors) + + exit_points = [ + ExitPoint(f"e{i}", name, report.widths[name]) for i, name in enumerate(exit_layers) + ] heads: dict[str, nn.Module] = { ep.name: EarlyExitHead( ep.in_channels, @@ -154,4 +247,14 @@ def custom_ee( confidence_thresholds=list(confidence_thresholds) if confidence_thresholds else [], routing_policy=routing_policy, ) - return EarlyExitWrapper(backbone, heads, identity, cfg, input_shape=input_shape) + # feature extractors are keyed by layer name for the user; the wrapper + # keys adapters by exit name. + adapters = { + ep.name: extractors[ep.layer_name] for ep in exit_points if ep.layer_name in extractors + } + wrapper = EarlyExitWrapper( + backbone, heads, identity, cfg, input_shape=input_shape, feature_adapters=adapters + ) + # the fresh exit heads must live where the backbone lives, or a wrapped + # CUDA model fails at the first routed inference with a device mismatch + return wrapper.to(run_device) diff --git a/earlyon/onnx.py b/earlyon/onnx.py index 7df6c27..a91221b 100644 --- a/earlyon/onnx.py +++ b/earlyon/onnx.py @@ -120,3 +120,10 @@ def export_to_onnx( finally: model.train(was_training) return output_names + + +# Explicitly named alias: this export computes ALL exits every time (routing +# is applied by the caller to the outputs); it does not short-circuit compute. +# For an export that skips later layers, see earlyon.staged / +# docs/STAGED_DEPLOYMENT.md. +export_all_exits_to_onnx = export_to_onnx diff --git a/earlyon/staged.py b/earlyon/staged.py new file mode 100644 index 0000000..6970fec --- /dev/null +++ b/earlyon/staged.py @@ -0,0 +1,221 @@ +"""Staged deployment: split an early-exit model into per-stage modules. + +The eager :class:`EarlyExitWrapper` short-circuits with hooks and exceptions, +which no graph exporter can express. The staged contract is the deployable +alternative: the network is cut *at the exit points* into an ordered list of +stages, where stage ``k``: + +* consumes the previous stage's **continuation features** (stage 0 consumes + the input image), +* produces the next continuation features, and +* produces **exit logits** (from exit head ``k``) for every stage except the + last, which produces the final classifier's logits. + +The deployment runtime (this module's :class:`StagedModel` in PyTorch, or +your own runtime over per-stage ONNX/TensorRT engines) runs stages in order +and applies earlyon's routing rule between them — per-head temperature, +enabled flag, confidence/entropy threshold — stopping at the first firing +exit. Unlike the single static multi-output ONNX graph +(:func:`earlyon.onnx.export_to_onnx`), later stages are genuinely *not +executed* when an early exit fires. + +Scope (deliberately narrow): the reference splitter supports backbones that +are literally ``nn.Sequential`` with exit layers at top-level positions — +arbitrary graph partitioning is out of scope and fails loudly. +:func:`staged_model` verifies its own equivalence against the eager wrapper +on a probe input before returning. See ``docs/STAGED_DEPLOYMENT.md``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import cast + +import torch +import torch.nn as nn + +from earlyon.core.types import EarlyExitConfig, InferenceResult +from earlyon.core.wrappers import EarlyExitWrapper, _safe_temperature + + +@dataclass(frozen=True) +class StageSpec: + """Describes one stage of a staged deployment. + + ``exit_name`` is the exit whose logits this stage emits (``"final"`` for + the last stage). ``modules`` lists the backbone submodule names executed + by the stage, in order. + """ + + index: int + exit_name: str + modules: tuple[str, ...] + + +class Stage(nn.Module): + """One deployable unit: a trunk plus the head evaluated after it. + + ``forward`` returns ``(continuation_features, logits)``. For the last + stage the continuation features equal the logits input path's features + and ``logits`` are the final classifier's output. + """ + + def __init__(self, trunk: nn.Sequential, head: nn.Module) -> None: + super().__init__() + self.trunk = trunk + self.head = head + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + features = self.trunk(x) + return features, self.head(features) + + +class _Identity(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + +class StagedModel(nn.Module): + """Reference staged runtime: runs stages in order, applies earlyon routing + between them, and stops at the first firing enabled exit. + + Produces exactly the same :class:`InferenceResult` as the eager wrapper + for the same config (equivalence is asserted at build time by + :func:`staged_model` and pinned by tests). + """ + + def __init__( + self, + stages: list[Stage], + specs: list[StageSpec], + config: EarlyExitConfig, + flops_at: dict[str, float], + ) -> None: + super().__init__() + if len(stages) != len(config.exit_points) + 1: + raise ValueError( + f"need one stage per exit plus a final stage: got {len(stages)} stages " + f"for {len(config.exit_points)} exits" + ) + config.validate() + self.stages = nn.ModuleList(stages) + self.specs = specs + self.config = config + self._flops_at = dict(flops_at) + + @torch.inference_mode() + def infer(self, x: torch.Tensor) -> InferenceResult: + if x.size(0) != 1: + raise ValueError(f"staged inference is batch-1 (got batch {x.size(0)})") + cfg = self.config + n_exits = len(cfg.exit_points) + features = x + for idx in range(n_exits): + ep = cfg.exit_points[idx] + if not cfg.enabled_exits[idx]: + # skip the head entirely; still run the trunk to continue + stage = cast(Stage, self.stages[idx]) + features = stage.trunk(features) + continue + features, logits = self.stages[idx](features) + temp = _safe_temperature(cfg.temperatures[ep.name]) + probs = torch.softmax(logits / temp, dim=-1) + confidence = probs.max().item() + if cfg.routing_policy == "confidence": + fires = confidence >= cfg.confidence_thresholds[idx] + else: + entropy = (-(probs * probs.clamp_min(1e-12).log()).sum()).item() + fires = entropy <= cfg.entropy_thresholds[idx] + if fires: + return InferenceResult( + prediction=logits, + exit_taken=idx, + confidence=confidence, + estimated_backbone_flops_fraction=self._flops_at[ep.layer_name], + ) + _, final_logits = self.stages[-1](features) + temp = _safe_temperature(cfg.temperatures["final"]) + probs = torch.softmax(final_logits / temp, dim=-1) + return InferenceResult( + prediction=final_logits, + exit_taken=-1, + confidence=probs.max().item(), + estimated_backbone_flops_fraction=1.0, + ) + + +def _sequential_children(backbone: nn.Module) -> list[tuple[str, nn.Module]]: + if not isinstance(backbone, nn.Sequential): + raise ValueError( + "staged_model supports backbones that are nn.Sequential (exit layers " + f"at top level); got {type(backbone).__name__}. Arbitrary graph " + "partitioning is out of scope — see docs/STAGED_DEPLOYMENT.md for the " + "protocol if you want to implement stages for a custom architecture." + ) + return list(backbone.named_children()) + + +def staged_model(wrapper: EarlyExitWrapper) -> StagedModel: + """Split an eager wrapper into a :class:`StagedModel`. + + Requirements (all verified): + + * the wrapper's backbone is an ``nn.Sequential`` whose forward is exactly + its children in order, returning ``(B, num_classes)`` logits (i.e. the + wrapper's ``final_classifier`` is identity — true for ``custom_ee``); + * every exit layer is a *top-level* child of that Sequential. + + Before returning, staged inference is checked against the eager wrapper + on a probe input; a mismatch raises rather than shipping a wrong split. + """ + children = _sequential_children(wrapper.backbone) + child_names = [name for name, _ in children] + positions: list[int] = [] + for ep in wrapper.config.exit_points: + if ep.layer_name not in child_names: + raise ValueError( + f"exit layer {ep.layer_name!r} is not a top-level child of the " + f"Sequential backbone (children: {child_names}); staged splitting " + "requires top-level exit layers" + ) + positions.append(child_names.index(ep.layer_name)) + if positions != sorted(positions): + raise ValueError(f"exit layers out of order in the Sequential: {positions}") + + stages: list[Stage] = [] + specs: list[StageSpec] = [] + start = 0 + for idx, ep in enumerate(wrapper.config.exit_points): + end = positions[idx] + 1 + trunk = nn.Sequential(*(mod for _, mod in children[start:end])) + stages.append(Stage(trunk, wrapper.exit_heads[ep.name])) + specs.append(StageSpec(idx, ep.name, tuple(child_names[start:end]))) + start = end + final_trunk = nn.Sequential(*(mod for _, mod in children[start:])) + stages.append(Stage(final_trunk, _Identity())) + specs.append(StageSpec(len(positions), "final", tuple(child_names[start:]))) + + model = StagedModel(stages, specs, wrapper.config, wrapper._flops_at) + + # equivalence probe: staged and eager must agree on a deterministic input. + was_training = wrapper.training + wrapper.eval() + model.eval() + try: + probe = torch.zeros( + wrapper._input_shape, + device=next(wrapper.backbone.parameters(), torch.empty(0)).device, + ) + eager = wrapper(probe, mode="inference") + staged = model.infer(probe) + if eager.exit_taken != staged.exit_taken or not torch.allclose( + eager.prediction, staged.prediction, atol=1e-5 + ): + raise RuntimeError( + "staged split does not reproduce the eager wrapper on a probe " + "input — the backbone's forward is not exactly its Sequential " + "children in order; refusing to build a wrong split" + ) + finally: + wrapper.train(was_training) + return model diff --git a/earlyon/utils.py b/earlyon/utils.py index fc9bf1d..52b1443 100644 --- a/earlyon/utils.py +++ b/earlyon/utils.py @@ -1,15 +1,40 @@ -"""Small shared helpers: dataset loaders, checkpoint save/load.""" +"""Small shared helpers: dataset loaders, checkpoint save/load. + +Checkpoint contract (format_version 2) +-------------------------------------- +``save_wrapper`` writes a dict with: + +* ``format_version`` — integer schema version (currently 2). +* ``earlyon_version`` — the library version that wrote the file. +* ``config`` — everything needed to reconstruct routing behavior when the + factory is known: backbone id, num_classes, routing policy, both threshold + lists, ``enabled_exits``, per-head ``temperatures``, loss weights, and the + exit points (name / layer_name / in_channels). +* ``state_dict`` — the wrapper's weights. + +``load_wrapper`` reads v2 files directly and migrates unversioned (v1) files: +the legacy scalar ``temperature`` is broadcast to every head, and the legacy +"disabled" threshold sentinels (confidence ``1.0`` / entropy ``0.0``, for the +active policy only) become explicit ``enabled_exits=False`` entries, with a +``UserWarning`` describing the migration. + +Arbitrary ``custom_ee`` backbones cannot be reconstructed from a string; pass +``factory=`` (a zero-argument callable returning a structurally identical +wrapper) to load them. +""" from __future__ import annotations +import math import warnings from pathlib import Path -from typing import Tuple +from typing import Any, Callable, Tuple import torch from torch.utils.data import DataLoader, Subset -from earlyon.core.types import Batch +from earlyon import __version__ +from earlyon.core.types import FINAL_HEAD, Batch from earlyon.core.wrappers import EarlyExitWrapper from earlyon.models import ( cifar_resnet_ee, @@ -22,6 +47,8 @@ _CIFAR_PREFIX = "cifar_resnet" +CHECKPOINT_FORMAT_VERSION = 2 + FACTORIES = { "resnet18": resnet18_ee, "resnet50": resnet50_ee, @@ -58,21 +85,22 @@ def build_model(backbone: str, num_classes: int, pretrained: bool = True) -> Ear def save_wrapper(model: EarlyExitWrapper, path: str | Path) -> None: - """Save state_dict + config. Config is needed to rebuild the wrapper. + """Save state_dict + a versioned config (format_version 2). - ``routing_policy`` and ``entropy_thresholds`` are persisted too: without - them an entropy-routed model would silently reload as confidence-routed - (the wrapper default), discarding the calibrated entropy thresholds. + The config records everything needed to reconstruct routing behavior when + the model factory is known — see the module docstring for the contract. """ if model.config.backbone == "custom": warnings.warn( "saving a custom_ee model: load_wrapper cannot rebuild an arbitrary " - "backbone from this artifact. To restore it, reconstruct the backbone, " - "re-wrap with custom_ee(...), and load the saved state_dict manually.", + "backbone from a string. To restore it, rebuild the backbone and pass " + "factory=lambda: custom_ee(...) to load_wrapper.", UserWarning, stacklevel=2, ) payload = { + "format_version": CHECKPOINT_FORMAT_VERSION, + "earlyon_version": __version__, "state_dict": model.state_dict(), "config": { "backbone": model.config.backbone, @@ -80,18 +108,109 @@ def save_wrapper(model: EarlyExitWrapper, path: str | Path) -> None: "routing_policy": model.config.routing_policy, "confidence_thresholds": list(model.config.confidence_thresholds), "entropy_thresholds": list(model.config.entropy_thresholds), + "enabled_exits": list(model.config.enabled_exits), + "temperatures": dict(model.config.temperatures), "loss_weights": list(model.config.loss_weights), - "temperature": model.config.temperature, + "exit_points": [ + {"name": ep.name, "layer_name": ep.layer_name, "in_channels": ep.in_channels} + for ep in model.config.exit_points + ], }, } torch.save(payload, Path(path)) -def load_wrapper(path: str | Path, pretrained_backbone: bool = False) -> EarlyExitWrapper: +def _migrate_v1_config(cfg: dict[str, Any], model: EarlyExitWrapper, name: str) -> dict[str, Any]: + """Translate an unversioned (v1) checkpoint config into v2 shape. + + Deterministic rules: + + * the legacy scalar ``temperature`` is broadcast to every head (that is + exactly what v1 routing did); a non-finite/non-positive scalar falls + back to 1.0 with a warning instead of failing the load, matching the + v1 runtime guard. + * the legacy "disabled" sentinels — confidence threshold ``1.0`` / + entropy threshold ``0.0`` on the *active* policy — become explicit + ``enabled_exits=False``. Under v1 semantics such an exit could still + fire on a numerically saturated softmax; the documented intent was + "disabled", which the explicit flag now honors exactly. + """ + migrated = dict(cfg) + n_exits = len(model.config.exit_points) + + temperature = float(cfg.get("temperature", 1.0)) + if not math.isfinite(temperature) or temperature <= 0.0: + warnings.warn( + f"checkpoint {name}: legacy temperature {temperature} is invalid; " + "falling back to 1.0 (uncalibrated)", + UserWarning, + stacklevel=3, + ) + temperature = 1.0 + head_names = [ep.name for ep in model.config.exit_points] + [FINAL_HEAD] + migrated["temperatures"] = {h: temperature for h in head_names} + + policy = cfg.get("routing_policy", "confidence") + if policy == "confidence": + active = [float(t) for t in cfg.get("confidence_thresholds", [])] + sentinel = 1.0 + else: + active = [float(t) for t in cfg.get("entropy_thresholds", [0.0] * n_exits)] + sentinel = 0.0 + enabled = [t != sentinel for t in active] + migrated["enabled_exits"] = enabled + if not all(enabled): + disabled_idx = [i for i, e in enumerate(enabled) if not e] + warnings.warn( + f"checkpoint {name}: migrated v1 sentinel threshold(s) at exit(s) " + f"{disabled_idx} to explicit enabled_exits=False. Under v1 these " + "exits could still fire on a saturated softmax; they are now " + "strictly disabled.", + UserWarning, + stacklevel=3, + ) + return migrated + + +def load_wrapper( + path: str | Path, + pretrained_backbone: bool = False, + factory: Callable[[], EarlyExitWrapper] | None = None, +) -> EarlyExitWrapper: + """Load a checkpoint written by :func:`save_wrapper`. + + ``factory`` — required for ``custom_ee`` models — must return a wrapper + structurally identical to the saved one; the checkpoint's routing config + and weights are then applied to it. Built-in backbones are reconstructed + automatically via :func:`build_model`. + """ # weights_only=True prevents pickle code-exec from malicious .pth files payload = torch.load(Path(path), map_location="cpu", weights_only=True) + name = Path(path).name + if "config" not in payload or "state_dict" not in payload: + raise ValueError(f"checkpoint {name}: missing 'config'/'state_dict'; not an earlyon file") + version = int(payload.get("format_version", 1)) + if version > CHECKPOINT_FORMAT_VERSION: + raise ValueError( + f"checkpoint {name}: format_version {version} is newer than this " + f"earlyon supports ({CHECKPOINT_FORMAT_VERSION}); upgrade earlyon" + ) cfg = payload["config"] - model = build_model(cfg["backbone"], cfg["num_classes"], pretrained=pretrained_backbone) + + if factory is not None: + model = factory() + elif cfg["backbone"] == "custom": + raise ValueError( + f"checkpoint {name}: backbone is 'custom', which cannot be rebuilt " + "from a string. Reconstruct the backbone and pass " + "factory=lambda: custom_ee(backbone, exit_layers=[...], ...) with the " + "same structure it was saved with." + ) + else: + model = build_model(cfg["backbone"], cfg["num_classes"], pretrained=pretrained_backbone) + + if version < 2: + cfg = _migrate_v1_config(cfg, model, name) # Validate config shape BEFORE mutating model.config. A checkpoint saved # against a different exit count would otherwise leave model.config in a @@ -101,46 +220,54 @@ def load_wrapper(path: str | Path, pretrained_backbone: bool = False) -> EarlyEx weight_len = len(cfg["loss_weights"]) if thr_len != n_exits: raise ValueError( - f"checkpoint {Path(path).name}: confidence_thresholds has " + f"checkpoint {name}: confidence_thresholds has " f"length {thr_len}, but backbone {cfg['backbone']!r} has " f"{n_exits} exit points" ) if weight_len != n_exits + 1: raise ValueError( - f"checkpoint {Path(path).name}: loss_weights has length " + f"checkpoint {name}: loss_weights has length " f"{weight_len}, expected {n_exits + 1} (one per exit plus final)" ) - # entropy_thresholds + routing_policy are optional for back-compat with - # pre-0.2 checkpoints; validate them only when present. load_wrapper - # bypasses EarlyExitConfig.__post_init__, so re-validate the persisted policy - # here — otherwise a corrupted checkpoint would silently mis-route. + # entropy_thresholds is optional in v1 files; validate when present. if "entropy_thresholds" in cfg and len(cfg["entropy_thresholds"]) != n_exits: raise ValueError( - f"checkpoint {Path(path).name}: entropy_thresholds has length " + f"checkpoint {name}: entropy_thresholds has length " f"{len(cfg['entropy_thresholds'])}, but backbone {cfg['backbone']!r} " f"has {n_exits} exit points" ) policy = cfg.get("routing_policy", model.config.routing_policy) if policy not in {"confidence", "entropy"}: raise ValueError( - f"checkpoint {Path(path).name}: unsupported routing_policy {policy!r} " + f"checkpoint {name}: unsupported routing_policy {policy!r} " "(allowed: 'confidence', 'entropy')" ) if policy == "entropy" and "entropy_thresholds" not in cfg: raise ValueError( - f"checkpoint {Path(path).name}: routing_policy='entropy' but no " + f"checkpoint {name}: routing_policy='entropy' but no " "entropy_thresholds were persisted — the model would silently route " "on uncalibrated defaults" ) + # v2 files carry the exit points; cross-check them against the rebuilt + # model so a checkpoint from a different factory revision fails loudly. + if "exit_points" in cfg: + saved = [(p["name"], p["layer_name"], int(p["in_channels"])) for p in cfg["exit_points"]] + built = [(ep.name, ep.layer_name, ep.in_channels) for ep in model.config.exit_points] + if saved != built: + raise ValueError( + f"checkpoint {name}: saved exit points {saved} do not match the " + f"reconstructed model's exit points {built}; the factory or " + "earlyon version that wrote this file placed exits differently" + ) - model.config.confidence_thresholds = list(cfg["confidence_thresholds"]) - model.config.loss_weights = list(cfg["loss_weights"]) - model.config.temperature = float(cfg["temperature"]) - # back-compat: older checkpoints predate these fields; keep the fresh - # model's defaults when the key is absent. + model.config.confidence_thresholds = [float(t) for t in cfg["confidence_thresholds"]] + model.config.loss_weights = [float(w) for w in cfg["loss_weights"]] model.config.routing_policy = policy if "entropy_thresholds" in cfg: - model.config.entropy_thresholds = list(cfg["entropy_thresholds"]) + model.config.entropy_thresholds = [float(t) for t in cfg["entropy_thresholds"]] + model.config.enabled_exits = [bool(e) for e in cfg["enabled_exits"]] + model.config.temperatures = {k: float(v) for k, v in cfg["temperatures"].items()} + model.config.validate() model.load_state_dict(payload["state_dict"]) return model diff --git a/examples/01_train_resnet50_cifar10.py b/examples/01_train_resnet50_cifar10.py index e1b1152..5e026fd 100644 --- a/examples/01_train_resnet50_cifar10.py +++ b/examples/01_train_resnet50_cifar10.py @@ -36,23 +36,39 @@ def main(): print("stage 1: training backbone") stage1_train_backbone( - model.backbone, train_loader, epochs=args.epochs_stage1, - lr=0.01, device=device, # lower lr because we start from pretrained + model.backbone, + train_loader, + epochs=args.epochs_stage1, + lr=0.01, + device=device, # lower lr because we start from pretrained ) print("stage 2: training exit heads") stage2_train_exits( - model, train_loader, epochs=args.epochs_stage2, lr=1e-3, device=device, + model, + train_loader, + epochs=args.epochs_stage2, + lr=1e-3, + device=device, ) + # Thresholds are chosen on the VALIDATION split only; the accuracies + # printed below are validation numbers. Report the deployable accuracy + # from the held-out test loader (earlyon.benchmarking.evaluate) — the + # threshold search must never consume test labels. print("calibrating thresholds on validation set") result = calibrate_thresholds( - model, val_loader, target_accuracy_drop=0.01, device=device, + model, + val_loader, + target_accuracy_drop=0.01, + device=device, + ) + print(f"thresholds: {result.thresholds} enabled: {result.enabled_exits}") + print( + f"val baseline acc: {result.baseline_accuracy:.4f} " + f"val routed acc: {result.final_accuracy:.4f} " + f"est. FLOPs fraction: {result.avg_computation_used:.4f}" ) - print(f"thresholds: {result.thresholds}") - print(f"baseline acc: {result.baseline_accuracy:.4f} " - f"final acc: {result.final_accuracy:.4f} " - f"avg compute: {result.avg_computation_used:.4f}") save_wrapper(model, args.output) print(f"saved {args.output}") diff --git a/examples/02_jetson_deployment.py b/examples/02_jetson_deployment.py index 446d3cc..9f7d6b7 100644 --- a/examples/02_jetson_deployment.py +++ b/examples/02_jetson_deployment.py @@ -4,15 +4,19 @@ python examples/02_jetson_deployment.py --model resnet50_ee_cifar10.pth -Outputs per-run latency, GPU utilization, power draw, and temperature. -On a non-Jetson host the script still runs but power/temp fields are zero. +Outputs a fair wrapper-vs-backbone throughput comparison (identical noise +samples and boundaries for both models — a best-case bound, since trained +heads may fire spuriously on noise), then per-run latency with instantaneous +power/temperature and the energy integrated over the timed window. +On a non-Jetson host the script still runs; missing telemetry is reported as +null, never as zero. """ import argparse import json import statistics -from earlyon.benchmarking import JetsonProfiler, benchmark_backbone, benchmark_wrapper +from earlyon.benchmarking import JetsonProfiler, benchmark_models from earlyon.utils import load_wrapper @@ -27,29 +31,53 @@ def main(): model = load_wrapper(args.model) shape = (1, 3, args.input_size, args.input_size) - print("throughput: wrapper vs backbone (cuda)") - wrap_r = benchmark_wrapper(model, input_shape=shape, device="cuda", - num_warmup=args.warmup, num_runs=args.runs) - bb_r = benchmark_backbone(model.backbone, input_shape=shape, device="cuda", - num_warmup=args.warmup, num_runs=args.runs) - speedup = wrap_r.throughput_ips / bb_r.throughput_ips - print(f" wrapper: {wrap_r.throughput_ips:.1f} ips " - f"(p50 {wrap_r.latency_p50_ms:.2f}ms, p95 {wrap_r.latency_p95_ms:.2f}ms)") - print(f" backbone: {bb_r.throughput_ips:.1f} ips " - f"(p50 {bb_r.latency_p50_ms:.2f}ms, p95 {bb_r.latency_p95_ms:.2f}ms)") - print(f" speedup: {speedup:.2f}x") - print(f" avg compute used: {wrap_r.avg_computation_used:.2%}") + print("fair throughput: wrapper vs backbone on identical samples (cuda)") + cmp_r = benchmark_models( + {"early_exit": model, "backbone": model.backbone}, + input_shape=shape, + device="cuda", + num_warmup=args.warmup, + num_runs=args.runs, + ) + wrap_r = cmp_r.results["early_exit"] + bb_r = cmp_r.results["backbone"] + print( + f" wrapper: {wrap_r.throughput_ips:.1f} ips " + f"(p50 {wrap_r.latency_p50_ms:.2f}ms, p95 {wrap_r.latency_p95_ms:.2f}ms)" + ) + print( + f" backbone: {bb_r.throughput_ips:.1f} ips " + f"(p50 {bb_r.latency_p50_ms:.2f}ms, p95 {bb_r.latency_p95_ms:.2f}ms)" + ) + print( + f" speedup: {cmp_r.speedup_vs('early_exit', 'backbone'):.2f}x " + f"(noise input: best-case bound)" + ) + print(f" est. backbone FLOPs fraction: {wrap_r.avg_estimated_flops_fraction:.2%}") print(f" exit distribution: {wrap_r.exit_distribution}") - print("\njetson power/thermal profile") + print("\njetson power/thermal/energy profile") profiler = JetsonProfiler() - runs = profiler.profile(model, input_shape=shape, num_runs=args.runs) - print(json.dumps({ - "tegrastats_available": profiler.monitor.available, - "latency_median_ms": statistics.median(r.latency_ms for r in runs), - "power_median_mw": statistics.median(r.power_mw for r in runs), - "temp_median_c": statistics.median(r.temp_c for r in runs), - }, indent=2)) + runs, energy = profiler.profile_with_energy( + model, input_shape=shape, num_warmup=args.warmup, num_runs=args.runs + ) + power = [r.power_mw for r in runs if r.power_mw is not None] + temp = [r.temp_c for r in runs if r.temp_c is not None] + print( + json.dumps( + { + "tegrastats_available": profiler.monitor.available, + "latency_median_ms": statistics.median(r.latency_ms for r in runs), + "instantaneous_power_median_mw": statistics.median(power) if power else None, + "temp_median_c": statistics.median(temp) if temp else None, + "avg_power_mw": energy.avg_power_mw, + "energy_mj": energy.energy_mj, + "energy_per_inference_mj": energy.energy_per_inference_mj, + "num_power_samples": energy.num_power_samples, + }, + indent=2, + ) + ) if __name__ == "__main__": diff --git a/pyproject.toml b/pyproject.toml index 8cc8e47..ac19d7a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta" [project] name = "earlyon" -version = "0.2.0" -description = "Production-ready early exit for PyTorch CV models." +version = "0.3.0" +description = "Deployment-oriented early-exit toolkit for PyTorch CV models." readme = "README.md" requires-python = ">=3.10" license = { file = "LICENSE" } diff --git a/scripts/evidence_run.py b/scripts/evidence_run.py new file mode 100644 index 0000000..ccd22cb --- /dev/null +++ b/scripts/evidence_run.py @@ -0,0 +1,298 @@ +"""Bounded, reproducible CUDA evidence run for release readiness. + +One seeded experiment on CIFAR-10 (public, auto-downloaded) intended to +*validate the system end-to-end*, not to maximize a headline number: + +- model: ResNet-18 early-exit wrapper, fine-tuned from ImageNet-pretrained + weights for a small, explicit number of epochs; +- static smaller baseline: MobileNetV2 (ImageNet-pretrained), fine-tuned the + same way — a legitimately smaller static model trained on the same data; +- explicit disjoint splits: train (45k), temperature (2.5k), calibration + (2.5k), official test (10k). The threshold search never sees test labels. +- fair benchmark: all three models measured by benchmark_models on the exact + same test samples with identical boundaries; a noise-input variant is + recorded separately as a best-case bound. + +Hard bounds: explicit epoch caps, a wall-clock budget check, one fixed seed. +Note: full bitwise determinism is not guaranteed on GPU (cuDNN autotuning); +the seed fixes data order and initialisation. + +Output: docs/evidence/cuda_evidence.json + docs/evidence/CUDA_EVIDENCE.md. +Negative results (e.g. wrapper slower than backbone) are recorded as-is. +""" + +from __future__ import annotations + +import argparse +import json +import platform +import sys +import time +from pathlib import Path + +import torch +import torchvision +from torch.utils.data import DataLoader, Subset +from torchvision import transforms + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT)) + +from earlyon import __version__ # noqa: E402 +from earlyon.benchmarking import benchmark_models, evaluate # noqa: E402 +from earlyon.core.thresholds import calibrate_thresholds # noqa: E402 +from earlyon.training import stage1_train_backbone, stage2_train_exits # noqa: E402 +from earlyon.utils import build_model, save_wrapper # noqa: E402 + +SEED = 42 +WALL_BUDGET_S = 3 * 3600 +OUT_DIR = ROOT / "docs" / "evidence" + +MEAN = (0.4914, 0.4822, 0.4465) +STD = (0.2470, 0.2435, 0.2616) + + +def loaders(root: str, image_size: int = 224, batch_size: int = 128): + train_tf = transforms.Compose( + [ + transforms.Resize(image_size), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + transforms.Normalize(MEAN, STD), + ] + ) + eval_tf = transforms.Compose( + [ + transforms.Resize(image_size), + transforms.ToTensor(), + transforms.Normalize(MEAN, STD), + ] + ) + train_full = torchvision.datasets.CIFAR10(root, train=True, download=True, transform=train_tf) + eval_train = torchvision.datasets.CIFAR10(root, train=True, download=False, transform=eval_tf) + test = torchvision.datasets.CIFAR10(root, train=False, download=True, transform=eval_tf) + + # explicit disjoint splits over the 50k train images + n = len(train_full) + temp_idx = list(range(0, 2500)) + calib_idx = list(range(2500, 5000)) + train_idx = list(range(5000, n)) + + mk = lambda ds, idx, bs, shuffle: DataLoader( # noqa: E731 + Subset(ds, idx), batch_size=bs, shuffle=shuffle, num_workers=2, pin_memory=True + ) + return { + "train": mk(train_full, train_idx, batch_size, True), + # MobileNetV2 activations at 224px don't fit a 6 GB GPU at batch 128; + # the baseline trains with a smaller batch (recorded in the output) + "train_small_batch": mk(train_full, train_idx, 32, True), + "temperature": mk(eval_train, temp_idx, batch_size, False), + "calibration": mk(eval_train, calib_idx, batch_size, False), + "test_batched": DataLoader(test, batch_size=64, num_workers=2), + "test_b1": DataLoader(test, batch_size=1, num_workers=2, pin_memory=True), + "counts": { + "train": len(train_idx), + "temperature": len(temp_idx), + "calibration": len(calib_idx), + "test": len(test), + }, + } + + +def batched_accuracy(module: torch.nn.Module, loader: DataLoader, device: str) -> float: + module = module.to(device).eval() + correct = total = 0 + with torch.inference_mode(): + for x, y in loader: + pred = module(x.to(device)).argmax(dim=-1).cpu() + correct += int((pred == y).sum()) + total += y.numel() + return correct / total + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--data", default=str(ROOT / "data")) + parser.add_argument("--epochs-stage1", type=int, default=2) + parser.add_argument("--epochs-stage2", type=int, default=2) + parser.add_argument("--epochs-baseline", type=int, default=2) + parser.add_argument("--bench-runs", type=int, default=300) + args = parser.parse_args() + + if not torch.cuda.is_available(): + sys.exit("CUDA is required for the evidence run") + device = "cuda" + torch.manual_seed(SEED) + torch.cuda.manual_seed_all(SEED) + t_start = time.time() + + def budget_check(stage: str) -> None: + spent = time.time() - t_start + if spent > WALL_BUDGET_S: + sys.exit(f"wall-clock budget exceeded before {stage} ({spent:.0f}s)") + + data = loaders(args.data) + log = {"stages": {}} + + # ---- early-exit ResNet-18 ---- + model = build_model("resnet18", num_classes=10, pretrained=True) + t0 = time.time() + stage1_train_backbone( + model.backbone, data["train"], epochs=args.epochs_stage1, lr=0.005, device=device + ) + log["stages"]["stage1_seconds"] = round(time.time() - t0, 1) + budget_check("stage2") + t0 = time.time() + stage2_train_exits(model, data["train"], epochs=args.epochs_stage2, lr=1e-3, device=device) + log["stages"]["stage2_seconds"] = round(time.time() - t0, 1) + budget_check("baseline") + + # ---- static smaller baseline: MobileNetV2 backbone ---- + # free the early-exit model's GPU residency first (6 GB card) + model = model.cpu() + torch.cuda.empty_cache() + baseline = build_model("mobilenetv2", num_classes=10, pretrained=True) + t0 = time.time() + stage1_train_backbone( + baseline.backbone, + data["train_small_batch"], + epochs=args.epochs_baseline, + lr=0.002, + device=device, + ) + log["stages"]["baseline_seconds"] = round(time.time() - t0, 1) + log["stages"]["baseline_batch_size"] = 32 + baseline = baseline.cpu() + torch.cuda.empty_cache() + budget_check("calibration") + + # ---- calibration: temperatures on the temperature split, thresholds on + # the calibration split; the test set is never touched here ---- + t0 = time.time() + calib = calibrate_thresholds( + model, + data["calibration"], + target_accuracy_drop=0.01, + device=device, + fit_temperature=True, + temperature_loader=data["temperature"], + ) + log["stages"]["calibrate_seconds"] = round(time.time() - t0, 1) + budget_check("evaluation") + + # ---- held-out test evaluation ---- + t0 = time.time() + report = evaluate(model, data["test_b1"], device=device) + log["stages"]["test_eval_seconds"] = round(time.time() - t0, 1) + backbone_test_acc = batched_accuracy(model.backbone, data["test_batched"], device) + baseline_test_acc = batched_accuracy(baseline.backbone, data["test_batched"], device) + budget_check("benchmark") + + # ---- fair benchmark: identical real test samples for all three models ---- + torch.cuda.reset_peak_memory_stats() + cmp_real = benchmark_models( + { + "backbone_resnet18": model.backbone, + "early_exit_resnet18": model, + "static_mobilenetv2": baseline.backbone, + }, + loader=data["test_b1"], + device=device, + num_warmup=50, + num_runs=args.bench_runs, + ) + peak_mem_mb = torch.cuda.max_memory_allocated() / 1e6 + cmp_noise = benchmark_models( + { + "backbone_resnet18": model.backbone, + "early_exit_resnet18": model, + "static_mobilenetv2": baseline.backbone, + }, + input_shape=(1, 3, 224, 224), + device=device, + num_warmup=50, + num_runs=args.bench_runs, + ) + + OUT_DIR.mkdir(parents=True, exist_ok=True) + save_wrapper(model, OUT_DIR / "resnet18_ee_evidence.pth") + + result = { + "purpose": "bounded system-validation run; not a tuned headline benchmark", + "seed": SEED, + "determinism_note": ( + "seed fixes init/data order; bitwise GPU determinism not enforced (cuDNN)" + ), + "dataset": {"name": "CIFAR-10 (torchvision, images upsampled to 224)", **data["counts"]}, + "weights": "ImageNet-pretrained torchvision weights, fine-tuned as recorded below", + "epochs": { + "stage1_backbone": args.epochs_stage1, + "stage2_exits": args.epochs_stage2, + "static_baseline": args.epochs_baseline, + }, + "model": { + "backbone": "resnet18", + "exit_points": [ + {"name": ep.name, "layer": ep.layer_name, "in_channels": ep.in_channels} + for ep in model.config.exit_points + ], + "head_config": "EarlyExitHead(hidden_dim=128, dropout=0.2)", + }, + "calibration": { + "objective": "accuracy_budget", + "target_accuracy_drop": 0.01, + "method": calib.method, + "thresholds": calib.thresholds, + "enabled_exits": calib.enabled_exits, + "temperatures": calib.temperatures, + "temperature_fit_converged": ( + {k: v.converged for k, v in calib.temperature_fits.items()} + if calib.temperature_fits + else None + ), + "calibration_split_baseline_acc": round(calib.baseline_accuracy, 4), + "calibration_split_routed_acc": round(calib.final_accuracy, 4), + }, + "test": { + "early_exit_routed_accuracy": round(report.overall_accuracy, 4), + "backbone_accuracy": round(backbone_test_acc, 4), + "static_mobilenetv2_accuracy": round(baseline_test_acc, 4), + "early_exit_estimated_flops_fraction": round(report.avg_computation_used, 4), + "exit_distribution": report.exit_distribution, + }, + "benchmark_real_input": {name: r.to_dict() for name, r in cmp_real.results.items()}, + "benchmark_real_speedups": { + "early_exit_vs_backbone": round( + cmp_real.speedup_vs("early_exit_resnet18", "backbone_resnet18"), 3 + ), + "static_vs_backbone": round( + cmp_real.speedup_vs("static_mobilenetv2", "backbone_resnet18"), 3 + ), + }, + "benchmark_noise_input": { + "note": "best-case bound: trained heads may fire spuriously on noise", + **{name: r.to_dict() for name, r in cmp_noise.results.items()}, + "early_exit_vs_backbone_speedup": round( + cmp_noise.speedup_vs("early_exit_resnet18", "backbone_resnet18"), 3 + ), + }, + "peak_cuda_memory_mb_during_real_benchmark": round(peak_mem_mb, 1), + "stage_seconds": log["stages"], + "total_wall_seconds": round(time.time() - t_start, 1), + "environment": { + "python": platform.python_version(), + "torch": torch.__version__, + "torchvision": torchvision.__version__, + "cuda": torch.version.cuda, + "gpu": torch.cuda.get_device_name(0), + "earlyon": __version__, + "os": platform.platform(), + }, + } + (OUT_DIR / "cuda_evidence.json").write_text(json.dumps(result, indent=2)) + print(json.dumps(result, indent=2)) + print(f"\nwrote {OUT_DIR / 'cuda_evidence.json'}") + + +if __name__ == "__main__": + main() diff --git a/scripts/re_evaluate.py b/scripts/re_evaluate.py index 04f72b7..4d46c2f 100644 --- a/scripts/re_evaluate.py +++ b/scripts/re_evaluate.py @@ -1,12 +1,15 @@ -"""Re-evaluate already-trained checkpoints with fixed FLOPs accounting. +"""Re-evaluate already-trained checkpoints with the v0.3 fair benchmark runner. -Skips training; just runs test-eval + throughput bench and rewrites -docs/benchmarks.json with the corrected avg_computation_used. +Skips training; runs test-set evaluation plus a fair throughput comparison +(wrapper and backbone measured on the exact same sample sequence and +boundary) and writes records under the "runs" key of docs/benchmarks.json. + +Pre-v0.3 records are preserved under "legacy_v0_2" and are NOT comparable: +the old runner benchmarked the wrapper and backbone on different inputs. """ import json import sys -import time from pathlib import Path import torch @@ -14,20 +17,17 @@ ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) -from earlyon.benchmarking import ( - benchmark_backbone, - benchmark_wrapper, - benchmark_wrapper_on_loader, - evaluate, -) -from earlyon.utils import cifar10_loaders, load_wrapper +from earlyon.benchmarking import benchmark_models, evaluate # noqa: E402 +from earlyon.utils import cifar10_loaders, load_wrapper # noqa: E402 OUT_JSON = ROOT / "docs" / "benchmarks.json" +METHODOLOGY = "v0.3-fair-runner" -def main(): +def main() -> None: device = "cuda" if torch.cuda.is_available() else "cpu" - results = json.loads(OUT_JSON.read_text()) if OUT_JSON.exists() else {} + doc = json.loads(OUT_JSON.read_text()) if OUT_JSON.exists() else {} + runs = doc.setdefault("runs", {}) _, _, test_loader = cifar10_loaders(batch_size=1, num_workers=2) @@ -38,57 +38,56 @@ def main(): continue print(f"[{backbone}] loading {ckpt}") model = load_wrapper(ckpt) - print(f"[{backbone}] _flops_at = {model._flops_at}") - # test-set evaluation with FIXED flops print(f"[{backbone}] test eval", flush=True) report = evaluate(model, test_loader, device=device) - # throughput bench — noise input (legacy upper bound) - print(f"[{backbone}] throughput (noise)", flush=True) - shape = (1, 3, 224, 224) - wrap_r = benchmark_wrapper( - model, input_shape=shape, device=device, num_warmup=50, num_runs=300, - ) - bb_r = benchmark_backbone( - model.backbone, input_shape=shape, device=device, num_warmup=50, num_runs=300, + # fair comparison: identical real-input samples for wrapper AND backbone + print(f"[{backbone}] fair throughput (real input)", flush=True) + cmp_real = benchmark_models( + {"early_exit": model, "backbone": model.backbone}, + loader=test_loader, + device=device, + num_warmup=50, + num_runs=300, ) - speedup = wrap_r.throughput_ips / bb_r.throughput_ips - - # throughput bench — real CIFAR-10 test samples (honest signal) - print(f"[{backbone}] throughput (real input)", flush=True) - real_r = benchmark_wrapper_on_loader( - model, test_loader, device=device, num_warmup=50, num_runs=300, + # noise input: best-case bound, same samples for both models + print(f"[{backbone}] fair throughput (noise input)", flush=True) + cmp_noise = benchmark_models( + {"early_exit": model, "backbone": model.backbone}, + input_shape=(1, 3, 224, 224), + device=device, + num_warmup=50, + num_runs=300, ) - real_speedup = real_r.throughput_ips / bb_r.throughput_ips - # merge into existing record - existing = results.get(backbone, {}) - existing.update({ + runs[backbone] = { "backbone": backbone, "dataset": "cifar10", + "methodology": METHODOLOGY, "device": torch.cuda.get_device_name(0) if device == "cuda" else "cpu", "test_accuracy": round(report.overall_accuracy, 4), - "test_avg_computation_used": round(report.avg_computation_used, 4), + "test_avg_estimated_flops_fraction": round(report.avg_computation_used, 4), "test_exit_distribution": report.exit_distribution, - "throughput_backbone_ips": round(bb_r.throughput_ips, 1), - "throughput_wrapper_ips_noise_input": round(wrap_r.throughput_ips, 1), - "throughput_wrapper_ips_real_input": round(real_r.throughput_ips, 1), - "throughput_speedup_noise_input": round(speedup, 3), - "throughput_speedup_real_input": round(real_speedup, 3), - "latency_backbone_p50_ms": round(bb_r.latency_p50_ms, 3), - "latency_wrapper_p50_ms_noise_input": round(wrap_r.latency_p50_ms, 3), - "latency_wrapper_p50_ms_real_input": round(real_r.latency_p50_ms, 3), - "real_input_avg_computation_used": round(real_r.avg_computation_used, 4), - "real_input_exit_distribution": real_r.exit_distribution, - "note": "noise-input numbers are best-case (trained heads fire spuriously on noise); real-input numbers are the honest input-distribution signal", - }) - results[backbone] = existing - OUT_JSON.write_text(json.dumps(results, indent=2)) - print(f"[{backbone}] acc={report.overall_accuracy:.4f} " - f"avg_comp={report.avg_computation_used:.4f} " - f"real_speedup={real_speedup:.2f}x " - f"dist={report.exit_distribution}", flush=True) + "real_input": { + "early_exit": cmp_real.results["early_exit"].to_dict(), + "backbone": cmp_real.results["backbone"].to_dict(), + "speedup": round(cmp_real.speedup_vs("early_exit", "backbone"), 3), + }, + "noise_input": { + "early_exit": cmp_noise.results["early_exit"].to_dict(), + "backbone": cmp_noise.results["backbone"].to_dict(), + "speedup": round(cmp_noise.speedup_vs("early_exit", "backbone"), 3), + "note": "best-case bound: trained heads may fire spuriously on noise", + }, + } + OUT_JSON.write_text(json.dumps(doc, indent=2)) + print( + f"[{backbone}] acc={report.overall_accuracy:.4f} " + f"est_flops={report.avg_computation_used:.4f} " + f"real_speedup={cmp_real.speedup_vs('early_exit', 'backbone'):.2f}x", + flush=True, + ) print("all done") diff --git a/scripts/run_benchmarks.py b/scripts/run_benchmarks.py index b6e2cb1..ba32dd1 100644 --- a/scripts/run_benchmarks.py +++ b/scripts/run_benchmarks.py @@ -9,7 +9,6 @@ import argparse import json -import os import sys import time import traceback @@ -20,10 +19,12 @@ ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(ROOT)) -from earlyon.benchmarking import benchmark_backbone, benchmark_wrapper, evaluate -from earlyon.core.thresholds import calibrate_thresholds -from earlyon.training import stage1_train_backbone, stage2_train_exits -from earlyon.utils import build_model, cifar10_loaders, save_wrapper +from earlyon.benchmarking import benchmark_models, evaluate # noqa: E402 +from earlyon.core.thresholds import calibrate_thresholds # noqa: E402 +from earlyon.training import stage1_train_backbone, stage2_train_exits # noqa: E402 +from earlyon.utils import build_model, cifar10_loaders, save_wrapper # noqa: E402 + +METHODOLOGY = "v0.3-fair-runner" OUT_DIR = ROOT / "docs" OUT_DIR.mkdir(exist_ok=True) @@ -31,9 +32,9 @@ # fine-tuning from imagenet pretrained -- a few epochs is enough for cifar10 PLAN = { - "resnet18": {"batch": 128, "stage1_epochs": 4, "stage2_epochs": 4, "lr1": 0.005}, - "resnet50": {"batch": 32, "stage1_epochs": 3, "stage2_epochs": 3, "lr1": 0.005}, - "mobilenetv2": {"batch": 32, "stage1_epochs": 4, "stage2_epochs": 4, "lr1": 0.002}, + "resnet18": {"batch": 128, "stage1_epochs": 4, "stage2_epochs": 4, "lr1": 0.005}, + "resnet50": {"batch": 32, "stage1_epochs": 3, "stage2_epochs": 3, "lr1": 0.005}, + "mobilenetv2": {"batch": 32, "stage1_epochs": 4, "stage2_epochs": 4, "lr1": 0.002}, } @@ -50,17 +51,21 @@ def run_one(backbone: str) -> dict: print(f"{log_prefix} starting on {device_name()}", flush=True) train_loader, val_loader, test_loader = cifar10_loaders( - batch_size=cfg["batch"], num_workers=2, + batch_size=cfg["batch"], + num_workers=2, ) model = build_model(backbone, num_classes=10, pretrained=True) t0 = time.time() print(f"{log_prefix} stage1: {cfg['stage1_epochs']} epochs", flush=True) stage1_train_backbone( - model.backbone, train_loader, - epochs=cfg["stage1_epochs"], lr=cfg["lr1"], device=device, - on_epoch_end=lambda l: print( - f"{log_prefix} s1 e{l.epoch}: loss={l.loss:.4f} acc={l.accuracy:.4f}", + model.backbone, + train_loader, + epochs=cfg["stage1_epochs"], + lr=cfg["lr1"], + device=device, + on_epoch_end=lambda log: print( + f"{log_prefix} s1 e{log.epoch}: loss={log.loss:.4f} acc={log.accuracy:.4f}", flush=True, ), ) @@ -69,11 +74,14 @@ def run_one(backbone: str) -> dict: t0 = time.time() print(f"{log_prefix} stage2: {cfg['stage2_epochs']} epochs", flush=True) stage2_train_exits( - model, train_loader, - epochs=cfg["stage2_epochs"], lr=1e-3, device=device, - on_epoch_end=lambda l: print( - f"{log_prefix} s2 e{l.epoch}: loss={l.loss:.4f} mean_acc={l.accuracy:.4f} " - f"per_exit={l.per_exit_accuracy}", + model, + train_loader, + epochs=cfg["stage2_epochs"], + lr=1e-3, + device=device, + on_epoch_end=lambda log: print( + f"{log_prefix} s2 e{log.epoch}: loss={log.loss:.4f} mean_acc={log.accuracy:.4f} " + f"per_exit={log.per_exit_accuracy}", flush=True, ), ) @@ -82,7 +90,10 @@ def run_one(backbone: str) -> dict: print(f"{log_prefix} calibrating thresholds", flush=True) t0 = time.time() calib = calibrate_thresholds( - model, val_loader, target_accuracy_drop=0.01, device=device, + model, + val_loader, + target_accuracy_drop=0.01, + device=device, ) t_calib = time.time() - t0 print( @@ -102,36 +113,52 @@ def run_one(backbone: str) -> dict: flush=True, ) - print(f"{log_prefix} throughput (wrapper vs backbone)", flush=True) - shape = (1, 3, 224, 224) - wrap_r = benchmark_wrapper( - model, input_shape=shape, device=device, num_warmup=50, num_runs=300, + # fair comparison: identical samples + boundaries for wrapper and backbone. + # real-input numbers are the honest signal; noise input is a best-case bound. + print(f"{log_prefix} fair throughput (real input)", flush=True) + cmp_real = benchmark_models( + {"early_exit": model, "backbone": model.backbone}, + loader=test_loader, + device=device, + num_warmup=50, + num_runs=300, ) - bb_r = benchmark_backbone( - model.backbone, input_shape=shape, device=device, num_warmup=50, num_runs=300, + print(f"{log_prefix} fair throughput (noise input)", flush=True) + cmp_noise = benchmark_models( + {"early_exit": model, "backbone": model.backbone}, + input_shape=(1, 3, 224, 224), + device=device, + num_warmup=50, + num_runs=300, ) - speedup = wrap_r.throughput_ips / bb_r.throughput_ips save_wrapper(model, OUT_DIR / f"{backbone}_cifar10.pth") return { "backbone": backbone, "dataset": "cifar10", + "methodology": METHODOLOGY, "device": device_name(), "stage1_seconds": round(t_stage1, 1), "stage2_seconds": round(t_stage2, 1), "calibrate_seconds": round(t_calib, 1), "thresholds": calib.thresholds, + "enabled_exits": calib.enabled_exits, "baseline_accuracy_val": round(calib.baseline_accuracy, 4), "test_accuracy": round(report.overall_accuracy, 4), - "test_avg_computation_used": round(report.avg_computation_used, 4), + "test_avg_estimated_flops_fraction": round(report.avg_computation_used, 4), "test_exit_distribution": report.exit_distribution, - "throughput_backbone_ips": round(bb_r.throughput_ips, 1), - "throughput_wrapper_ips": round(wrap_r.throughput_ips, 1), - "speedup": round(speedup, 3), - "latency_backbone_p50_ms": round(bb_r.latency_p50_ms, 3), - "latency_wrapper_p50_ms": round(wrap_r.latency_p50_ms, 3), - "latency_wrapper_p95_ms": round(wrap_r.latency_p95_ms, 3), + "real_input": { + "early_exit": cmp_real.results["early_exit"].to_dict(), + "backbone": cmp_real.results["backbone"].to_dict(), + "speedup": round(cmp_real.speedup_vs("early_exit", "backbone"), 3), + }, + "noise_input": { + "early_exit": cmp_noise.results["early_exit"].to_dict(), + "backbone": cmp_noise.results["backbone"].to_dict(), + "speedup": round(cmp_noise.speedup_vs("early_exit", "backbone"), 3), + "note": "best-case bound: trained heads may fire spuriously on noise", + }, } @@ -140,12 +167,13 @@ def main(): parser.add_argument("backbones", nargs="*", default=list(PLAN.keys())) args = parser.parse_args() - results = {} + doc = {} if OUT_JSON.exists(): try: - results = json.loads(OUT_JSON.read_text()) + doc = json.loads(OUT_JSON.read_text()) except Exception: - results = {} + doc = {} + results = doc.setdefault("runs", {}) for backbone in args.backbones: if backbone not in PLAN: @@ -153,7 +181,7 @@ def main(): continue try: results[backbone] = run_one(backbone) - OUT_JSON.write_text(json.dumps(results, indent=2)) + OUT_JSON.write_text(json.dumps(doc, indent=2)) print(f"[{backbone}] DONE; wrote {OUT_JSON}", flush=True) except Exception as exc: print(f"[{backbone}] FAILED: {exc}", flush=True) diff --git a/scripts/smoke_test.py b/scripts/smoke_test.py new file mode 100644 index 0000000..0f5283b --- /dev/null +++ b/scripts/smoke_test.py @@ -0,0 +1,155 @@ +"""Post-install smoke test: exercises the public API of an *installed* earlyon. + +Run from OUTSIDE the repository against a clean environment: + + pip install + python smoke_test.py + +Covers: import + version, built-in model construction, training-mode forward, +routed inference, deterministic calibration on a synthetic fixture, v2 +checkpoint save/load, v1 checkpoint migration, staged runtime equivalence, +and a benchmark smoke. Uses synthetic data only — nothing here is an accuracy +or speed claim. Exits non-zero on the first failure. +""" + +from __future__ import annotations + +import sys +import tempfile +from pathlib import Path + +import torch +from torch.utils.data import DataLoader, TensorDataset + + +def check(name: str, ok: bool, detail: str = "") -> None: + print(f"[{'PASS' if ok else 'FAIL'}] {name}" + (f" — {detail}" if detail else "")) + if not ok: + sys.exit(1) + + +def main() -> None: + torch.manual_seed(0) + + import earlyon + + check("import + version", bool(earlyon.__version__), earlyon.__version__) + + # built-in model (no pretrained download), training + inference forwards + from earlyon.models import cifar_resnet_ee + + model = cifar_resnet_ee(num_classes=10, depth=20) + outputs = model(torch.randn(2, 3, 32, 32), mode="training") + check("training forward", len(outputs) == len(model.config.exit_points) + 1) + model.eval() + result = model(torch.randn(1, 3, 32, 32), mode="inference") + check( + "routed inference", + result.prediction.shape == (1, 10) + and 0.0 <= result.estimated_backbone_flops_fraction <= 1.0, + f"exit_taken={result.exit_taken}", + ) + + # deterministic calibration on a synthetic fixture + from earlyon.core.thresholds import calibrate_thresholds + + x = torch.randn(32, 3, 32, 32) + y = torch.randint(0, 10, (32,)) + loader = DataLoader(TensorDataset(x, y), batch_size=8) + calib = calibrate_thresholds(model, loader, target_accuracy_drop=0.05, device="cpu") + check( + "calibration", + len(calib.thresholds) == len(model.config.exit_points) + and calib.num_samples == 32 + and calib.schema_version == 2, + f"enabled={calib.enabled_exits}", + ) + + from earlyon.utils import load_wrapper, save_wrapper + + with tempfile.TemporaryDirectory() as tmp: + # v2 round-trip + v2_path = Path(tmp) / "model_v2.pth" + save_wrapper(model, v2_path) + loaded = load_wrapper(v2_path) + check( + "v2 save/load round-trip", + loaded.config.enabled_exits == model.config.enabled_exits + and loaded.config.temperatures == model.config.temperatures, + ) + + # v1 migration: the exact payload shape earlyon <= 0.2 wrote + v1_path = Path(tmp) / "model_v1.pth" + torch.save( + { + "state_dict": model.state_dict(), + "config": { + "backbone": "cifar_resnet20", + "num_classes": 10, + "routing_policy": "confidence", + "confidence_thresholds": [0.7, 1.0, 0.85], # 1.0 = v1 sentinel + "entropy_thresholds": [0.5, 0.5, 0.5], + "loss_weights": [0.1, 0.2, 0.3, 0.4], + "temperature": 1.5, + }, + }, + v1_path, + ) + import warnings + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + migrated = load_wrapper(v1_path) + check( + "v1 checkpoint migration", + migrated.config.enabled_exits == [True, False, True] + and migrated.config.temperatures["final"] == 1.5, + ) + + # staged runtime equivalence on a Sequential backbone + import torch.nn as nn + + from earlyon.models import custom_ee + from earlyon.staged import staged_model + + seq = nn.Sequential( + nn.Conv2d(3, 8, 3, padding=1), + nn.ReLU(), + nn.AdaptiveAvgPool2d(1), + nn.Flatten(), + nn.Linear(8, 10), + ) + wrapper = custom_ee(seq, exit_layers=["0"], num_classes=10, input_shape=(1, 3, 16, 16)) + staged = staged_model(wrapper) + wrapper.eval() + staged.eval() + probe = torch.randn(1, 3, 16, 16) + eager_r = wrapper(probe, mode="inference") + staged_r = staged.infer(probe) + check( + "staged == eager", + staged_r.exit_taken == eager_r.exit_taken + and torch.allclose(staged_r.prediction, eager_r.prediction, atol=1e-6), + ) + + # benchmark smoke: identical samples, tiny run + from earlyon.benchmarking import benchmark_models + + cmp_r = benchmark_models( + {"early_exit": model, "backbone": model.backbone}, + input_shape=(1, 3, 32, 32), + num_warmup=2, + num_runs=5, + ) + check( + "benchmark smoke", + cmp_r.results["early_exit"].num_runs == 5 + and cmp_r.results["early_exit"].boundary == "model-only", + f"speedup={cmp_r.speedup_vs('early_exit', 'backbone'):.2f}x (synthetic; not a claim)", + ) + + print("SMOKE OK") + + +if __name__ == "__main__": + main() diff --git a/tests/fixtures/README.md b/tests/fixtures/README.md new file mode 100644 index 0000000..195d2c8 --- /dev/null +++ b/tests/fixtures/README.md @@ -0,0 +1,11 @@ +# Test fixtures + +- `tiny_models.py` — tiny CNN/transformer backbones for fast tests. +- `v1_cifar_resnet20.pth` (~1.2 MB) — a genuine pre-v0.3 (format v1) + checkpoint written by the actual historical `save_wrapper` (git `main`, + commit a5f8414) on a seeded, untrained `cifar_resnet_ee(depth=20)`. + Contains the legacy scalar `temperature` (1.6) and the v1 "disabled" + sentinel threshold (confidence 1.0 at exit 1). Used by + `tests/test_release_audit.py::test_real_v1_fixture_checkpoint_migrates` + to prove the v1→v2 migration against a real prior-version file, not a + hand-built dict. Deliberately force-added past the `*.pth` gitignore. diff --git a/tests/fixtures/v1_cifar_resnet20.pth b/tests/fixtures/v1_cifar_resnet20.pth new file mode 100644 index 0000000..b1e1554 Binary files /dev/null and b/tests/fixtures/v1_cifar_resnet20.pth differ diff --git a/tests/test_benchmarks.py b/tests/test_benchmarks.py index 76e6ac6..7901b70 100644 --- a/tests/test_benchmarks.py +++ b/tests/test_benchmarks.py @@ -132,3 +132,126 @@ def test_benchmark_wrapper_on_loader_cycles_loader_when_short(): r = benchmark_wrapper_on_loader(model, loader, device="cpu", num_warmup=2, num_runs=20) assert r.num_runs == 20 + + +# ---------------- fair comparison runner (v0.3) ---------------- + + +class _SpyModule(torch.nn.Module): + """Records every input it sees so tests can prove identical sample feeds.""" + + def __init__(self): + super().__init__() + self.seen = [] + self.linear = torch.nn.Linear(3 * 32 * 32, 10) + + def forward(self, x): + self.seen.append(x.detach().clone()) + return self.linear(x.flatten(1)) + + +def test_benchmark_models_feeds_identical_samples_to_every_model(): + from earlyon.benchmarking import benchmark_models + + a, b = _SpyModule(), _SpyModule() + benchmark_models({"a": a, "b": b}, input_shape=(1, 3, 32, 32), num_warmup=2, num_runs=5) + assert len(a.seen) == len(b.seen) == 7 + for xa, xb in zip(a.seen, b.seen): + assert torch.equal(xa, xb), "models must see the exact same sample sequence" + + +def test_benchmark_models_reports_accuracy_from_labelled_loader(): + from earlyon.benchmarking import benchmark_models + + torch.manual_seed(0) + model = _build(thresholds=(0.0, 0.0)) # always exits at e0 + images = torch.randn(16, 3, 32, 32) + labels = torch.randint(0, 10, (16,)) + loader = DataLoader(TensorDataset(images, labels), batch_size=1) + cmp = benchmark_models( + {"wrapper": model, "backbone": model.backbone}, + loader=loader, + num_warmup=2, + num_runs=8, + ) + for r in cmp.results.values(): + assert r.accuracy is not None and 0.0 <= r.accuracy <= 1.0 + assert r.num_labelled_samples == 8 + assert r.input_source == "loader" + assert cmp.results["wrapper"].exit_distribution == {"exit_0": 1.0} + assert cmp.results["backbone"].exit_distribution == {"final": 1.0} + assert cmp.speedup_vs("wrapper", "backbone") > 0 + + +def test_benchmark_result_carries_measurement_metadata(): + from earlyon.benchmarking import benchmark_models + + model = _build() + cmp = benchmark_models({"m": model}, input_shape=(1, 3, 32, 32), num_warmup=1, num_runs=3) + r = cmp.results["m"] + assert r.boundary == "model-only" + assert r.sync == "none" # cpu + assert r.batch_size == 1 + assert r.input_shape == (1, 3, 32, 32) + assert r.dtype == "torch.float32" + assert r.num_warmup == 1 + assert r.input_source == "random-noise" + assert r.accuracy is None + d = r.to_dict() + assert d["boundary"] == "model-only" + # deprecated alias still reads + assert r.avg_computation_used == r.avg_estimated_flops_fraction + + +def test_benchmark_models_end_to_end_boundary_is_labelled(): + from earlyon.benchmarking import benchmark_models + + model = _build() + cmp = benchmark_models( + {"m": model}, + input_shape=(1, 3, 32, 32), + num_warmup=1, + num_runs=3, + boundary="end-to-end", + ) + assert cmp.results["m"].boundary == "end-to-end" + + +def test_benchmark_models_validates_inputs(): + import pytest + + from earlyon.benchmarking import benchmark_models + + model = _build() + with pytest.raises(ValueError, match="num_warmup"): + benchmark_models({"m": model}, input_shape=(1, 3, 32, 32), num_warmup=-1) + with pytest.raises(ValueError, match="num_runs"): + benchmark_models({"m": model}, input_shape=(1, 3, 32, 32), num_runs=0) + with pytest.raises(ValueError, match="models"): + benchmark_models({}, input_shape=(1, 3, 32, 32)) + with pytest.raises(ValueError, match="loader or an input_shape"): + benchmark_models({"m": model}) + with pytest.raises(ValueError, match="boundary"): + benchmark_models( + {"m": model}, + input_shape=(1, 3, 32, 32), + num_warmup=0, + num_runs=1, + boundary="sideways", + ) + empty = DataLoader( + TensorDataset(torch.empty(0, 3, 32, 32), torch.empty(0, dtype=torch.long)), + batch_size=1, + ) + with pytest.raises(ValueError, match="non-empty"): + benchmark_models({"m": model}, loader=empty, num_warmup=0, num_runs=1) + + +def test_benchmark_models_noise_is_deterministic_per_seed(): + from earlyon.benchmarking import benchmark_models + + a, b = _SpyModule(), _SpyModule() + benchmark_models({"a": a}, input_shape=(1, 3, 32, 32), num_warmup=0, num_runs=3, seed=7) + benchmark_models({"b": b}, input_shape=(1, 3, 32, 32), num_warmup=0, num_runs=3, seed=7) + for xa, xb in zip(a.seen, b.seen): + assert torch.equal(xa, xb) diff --git a/tests/test_cli.py b/tests/test_cli.py index 87a82b2..ea86b9c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -122,7 +122,8 @@ def test_wrap_creates_checkpoint(tmp_path): def test_save_load_round_trip(tmp_path): - """save_wrapper + load_wrapper must preserve thresholds, weights, temperature.""" + """save_wrapper + load_wrapper must preserve thresholds, weights, + per-head temperatures, and enabled_exits.""" import torch from earlyon.utils import build_model, load_wrapper, save_wrapper @@ -130,7 +131,8 @@ def test_save_load_round_trip(tmp_path): model = build_model("resnet18", num_classes=10, pretrained=False) model.config.confidence_thresholds = [0.42, 0.73] model.config.loss_weights = [0.1, 0.2, 0.7] - model.config.temperature = 1.5 + model.config.temperatures = {"e0": 1.5, "e1": 0.9, "final": 1.2} + model.config.enabled_exits = [True, False] # mutate an exit head weight so we can verify state_dict round-trip with torch.no_grad(): first_param = next(iter(model.exit_heads["e0"].parameters())) @@ -142,7 +144,8 @@ def test_save_load_round_trip(tmp_path): assert loaded.config.confidence_thresholds == [0.42, 0.73] assert loaded.config.loss_weights == [0.1, 0.2, 0.7] - assert loaded.config.temperature == 1.5 + assert loaded.config.temperatures == {"e0": 1.5, "e1": 0.9, "final": 1.2} + assert loaded.config.enabled_exits == [True, False] loaded_first = next(iter(loaded.exit_heads["e0"].parameters())) assert torch.allclose(loaded_first, torch.full_like(loaded_first, 0.123)) diff --git a/tests/test_custom_ee.py b/tests/test_custom_ee.py index 7bde6e7..b5a89e9 100644 --- a/tests/test_custom_ee.py +++ b/tests/test_custom_ee.py @@ -146,3 +146,156 @@ def test_custom_ee_bad_layer_name_raises_value_error_listing_layers(): backbone = TinyBackbone(num_classes=10) with pytest.raises(ValueError, match="stage1"): custom_ee(backbone, ["stage_one_typo"], num_classes=10, input_shape=(1, 3, 32, 32)) + + +# ---------------- v0.3 interface: examples, devices, adapters, validation ---------------- + + +class _KwargBackbone(torch.nn.Module): + """Backbone whose forward takes an extra kwarg with a runtime default.""" + + def __init__(self, num_classes=10): + super().__init__() + self.stem = torch.nn.Conv2d(3, 8, 3, padding=1) + self.head = torch.nn.Linear(8, num_classes) + + def forward(self, x, scale: float = 1.0): + f = self.stem(x) * scale + return self.head(f.mean(dim=(2, 3))) + + +class _TupleOutputBackbone(torch.nn.Module): + """The exit layer emits a (features, aux) tuple — needs an extractor.""" + + def __init__(self, num_classes=10): + super().__init__() + self.block = _TupleBlock() + self.head = torch.nn.Linear(8, num_classes) + + def forward(self, x): + f, _aux = self.block(x) + return self.head(f.mean(dim=(2, 3))) + + +class _TupleBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3, padding=1) + + def forward(self, x): + f = self.conv(x) + return f, {"aux": f.sum()} + + +class _ReusedLayerBackbone(torch.nn.Module): + def __init__(self, num_classes=10): + super().__init__() + self.conv = torch.nn.Conv2d(3, 3, 3, padding=1) + self.head = torch.nn.Linear(3, num_classes) + + def forward(self, x): + x = self.conv(x) + x = self.conv(x) # same module twice + return self.head(x.mean(dim=(2, 3))) + + +def test_custom_ee_accepts_example_args_and_kwargs(): + backbone = _KwargBackbone() + model = custom_ee( + backbone, + exit_layers=["stem"], + num_classes=10, + input_shape=(1, 3, 16, 16), + example_args=(torch.zeros(1, 3, 16, 16),), + example_kwargs={"scale": 2.0}, + ) + model.eval() + result = model(torch.randn(1, 3, 16, 16), mode="inference") + assert result.prediction.shape == (1, 10) + + +def test_custom_ee_example_kwargs_without_args_rejected(): + with pytest.raises(ValueError, match="example_kwargs"): + custom_ee( + _KwargBackbone(), + exit_layers=["stem"], + num_classes=10, + example_kwargs={"scale": 2.0}, + ) + + +def test_custom_ee_tuple_output_needs_and_uses_feature_extractor(): + backbone = _TupleOutputBackbone() + # without an extractor: clear error naming the fix + with pytest.raises(ValueError, match="feature_extractors"): + custom_ee(backbone, exit_layers=["block"], num_classes=10, input_shape=(1, 3, 16, 16)) + # with one: wraps and routes + model = custom_ee( + _TupleOutputBackbone(), + exit_layers=["block"], + num_classes=10, + input_shape=(1, 3, 16, 16), + feature_extractors={"block": lambda out: out[0]}, + ) + model.eval() + model.config.confidence_thresholds = [0.0] # force the exit to fire + result = model(torch.randn(1, 3, 16, 16), mode="inference") + assert result.exit_taken == 0 # the adapter ran inside the routing hook + assert result.prediction.shape == (1, 10) + + +def test_custom_ee_rejects_extractor_for_unknown_layer(): + with pytest.raises(ValueError, match="unknown exit layer"): + custom_ee( + _TupleOutputBackbone(), + exit_layers=["block"], + num_classes=10, + input_shape=(1, 3, 16, 16), + feature_extractors={"block": lambda o: o[0], "nope": lambda o: o}, + ) + + +def test_custom_ee_rejects_reused_exit_layer(): + with pytest.raises(RuntimeError, match="more than once"): + custom_ee( + _ReusedLayerBackbone(), + exit_layers=["conv"], + num_classes=10, + input_shape=(1, 3, 16, 16), + ) + + +def test_custom_ee_rejects_out_of_order_exit_layers(): + backbone = TinyBackbone(num_classes=10) + with pytest.raises(RuntimeError, match="forward-execution order"): + custom_ee( + backbone, + exit_layers=["stage2", "stage1"], # reversed + num_classes=10, + input_shape=(1, 3, 32, 32), + ) + + +def test_custom_ee_restores_training_mode_after_dry_run(): + backbone = TinyBackbone(num_classes=10) + backbone.train() + custom_ee(backbone, exit_layers=["stage1"], num_classes=10, input_shape=(1, 3, 32, 32)) + assert backbone.training is True + + +def test_custom_ee_infers_device_from_backbone(): + """A meta-level check that stays CPU-only: the dry-run tensor must follow + the backbone's parameter device rather than defaulting to CPU.""" + from earlyon.models.custom import _infer_backbone_device + + backbone = TinyBackbone(num_classes=10) + assert _infer_backbone_device(backbone).type == "cpu" + + +@pytest.mark.gpu +def test_custom_ee_wraps_cuda_backbone_without_device_error(): + backbone = TinyBackbone(num_classes=10).cuda() + model = custom_ee(backbone, exit_layers=["stage1"], num_classes=10, input_shape=(1, 3, 32, 32)) + model.eval() + result = model(torch.randn(1, 3, 32, 32, device="cuda"), mode="inference") + assert result.prediction.shape == (1, 10) diff --git a/tests/test_flops.py b/tests/test_flops.py index b3993f9..b93c598 100644 --- a/tests/test_flops.py +++ b/tests/test_flops.py @@ -116,3 +116,94 @@ def test_uniform_fallback_for_layer_name_with_no_leaves(monkeypatch): backbone = _toy_backbone() result = per_layer_flops(backbone, ["nonexistent"], input_shape=(1, 3, 8, 8)) assert result["nonexistent"] == pytest.approx(1 / 2) + + +# ---------------- estimator metadata + reuse detection (v0.3) ---------------- + + +def test_estimate_layer_flops_reports_method_and_reliability(): + from earlyon.core.flops import METHOD_FVCORE, estimate_layer_flops + + backbone = _toy_backbone() + est = estimate_layer_flops(backbone, ["0", "1"], input_shape=(1, 3, 8, 8)) + assert est.method == METHOD_FVCORE + assert est.reliable is True + assert est.excludes_exit_heads is True + assert set(est.fractions) == {"0", "1"} + assert est.fractions["0"] <= est.fractions["1"] <= 1.0 + + +def test_reused_module_falls_back_to_low_confidence_uniform(): + """A backbone that calls the same leaf twice breaks the module-order + assumption; the estimator must warn and return reliable=False instead of a + precise-looking but wrong fraction.""" + + from earlyon.core.flops import METHOD_UNIFORM, estimate_layer_flops + + class Reuser(nn.Module): + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(3, 3, 3, padding=1) + self.head = nn.Conv2d(3, 8, 3, padding=1) + + def forward(self, x): + x = self.conv(x) + x = self.conv(x) # same module twice + return self.head(x) + + backbone = Reuser() + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + est = estimate_layer_flops(backbone, ["conv"], input_shape=(1, 3, 8, 8)) + assert est.method == METHOD_UNIFORM + assert est.reliable is False + assert any("reuses module" in str(w.message) for w in caught) + + +def test_wrapper_flops_analysis_is_lazy(): + """Constructing a wrapper must NOT run the FLOPs analysis (it made ViT + construction unacceptably slow); the first inference materialises it.""" + from unittest.mock import patch + + import torch + + from earlyon.core.exit_head import EarlyExitHead + from earlyon.core.types import EarlyExitConfig, ExitPoint + from earlyon.core.wrappers import EarlyExitWrapper + from tests.fixtures.tiny_models import STAGE_CHANNELS, TinyBackbone + + with patch("earlyon.core.wrappers.estimate_layer_flops", autospec=True) as mocked: + from earlyon.core.flops import estimate_layer_flops as real + + mocked.side_effect = real + backbone = TinyBackbone(num_classes=10) + exits = [ExitPoint("e0", "stage1", STAGE_CHANNELS["stage1"])] + heads = {"e0": EarlyExitHead(STAGE_CHANNELS["stage1"], 10)} + cfg = EarlyExitConfig(backbone="tiny", num_classes=10, exit_points=exits) + wrapper = EarlyExitWrapper(backbone, heads, lambda t: t, cfg, input_shape=(1, 3, 32, 32)) + assert mocked.call_count == 0, "construction must not run the FLOPs analysis" + + wrapper.eval() + wrapper(torch.randn(1, 3, 32, 32), mode="inference") + assert mocked.call_count == 1 + + # cached: further inferences do not re-run the analysis + wrapper(torch.randn(1, 3, 32, 32), mode="inference") + assert mocked.call_count == 1 + + +def test_wrapper_exposes_flops_estimate_metadata(): + + from earlyon.core.exit_head import EarlyExitHead + from earlyon.core.types import EarlyExitConfig, ExitPoint + from earlyon.core.wrappers import EarlyExitWrapper + from tests.fixtures.tiny_models import STAGE_CHANNELS, TinyBackbone + + backbone = TinyBackbone(num_classes=10) + exits = [ExitPoint("e0", "stage1", STAGE_CHANNELS["stage1"])] + heads = {"e0": EarlyExitHead(STAGE_CHANNELS["stage1"], 10)} + cfg = EarlyExitConfig(backbone="tiny", num_classes=10, exit_points=exits) + wrapper = EarlyExitWrapper(backbone, heads, lambda t: t, cfg, input_shape=(1, 3, 32, 32)) + est = wrapper.flops_estimate + assert est.excludes_exit_heads is True + assert 0.0 < est.fractions["stage1"] <= 1.0 diff --git a/tests/test_integration.py b/tests/test_integration.py index 51d4790..8798a9f 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -98,9 +98,9 @@ def test_pipeline_custom_token_entropy_joint_calibrate_onnx(tmp_path): 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 + # custom models save the state_dict but can't be rebuilt without a factory save_wrapper(model, tmp_path / "c.pth") - with pytest.raises(NotImplementedError, match="custom_ee"): + with pytest.raises(ValueError, match="factory"): load_wrapper(tmp_path / "c.pth") # ONNX export still works (the token/3D head path) diff --git a/tests/test_jetson_profiler.py b/tests/test_jetson_profiler.py index 0bf8a4f..58f3a5a 100644 --- a/tests/test_jetson_profiler.py +++ b/tests/test_jetson_profiler.py @@ -1,11 +1,18 @@ """Jetson profiler tests. These run on a non-Jetson host: tegrastats is -absent so power/temp fields will be zero, but the harness must still execute -end-to-end.""" +absent so power/temp fields are None (missing, not invented), but the +harness must still execute end-to-end. Parser, lifecycle, and energy +integration are tested without any hardware.""" + +import time + +import pytest from earlyon.benchmarking.jetson_profiler import ( JetsonProfiler, + JetsonSample, TegrastatsMonitor, _parse_line, + integrate_energy, ) from earlyon.core.exit_head import EarlyExitHead from earlyon.core.types import EarlyExitConfig, ExitPoint @@ -17,6 +24,22 @@ "GR3D_FREQ 75% GR3D2_FREQ 0% tj@42.5C VDD_IN 4500 VDD_CPU_GPU_CV 1200" ) +# Orin-style format variant: different power rail name, PLL temp +ORIN_LINE = "RAM 2000/16000MB GR3D_FREQ 12% pll@38C VDD_GPU 800" + + +def _wrapper(): + backbone = TinyBackbone(num_classes=10) + exits = [ExitPoint("e0", "stage1", STAGE_CHANNELS["stage1"])] + heads = {ep.name: EarlyExitHead(ep.in_channels, 10) for ep in exits} + cfg = EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=exits, + confidence_thresholds=[0.8], + ) + return EarlyExitWrapper(backbone, heads, lambda x: x, cfg, input_shape=(1, 3, 32, 32)) + def test_parse_line_extracts_gpu_temp_power(): s = _parse_line(SAMPLE_LINE) @@ -25,41 +48,142 @@ def test_parse_line_extracts_gpu_temp_power(): assert s.power_mw == 4500.0 -def test_parse_line_missing_fields_returns_zero(): +def test_parse_line_orin_variant(): + s = _parse_line(ORIN_LINE) + assert s.gpu_util_pct == 12.0 + assert s.temp_c == 38.0 + assert s.power_mw == 800.0 + + +def test_parse_line_missing_fields_are_none_not_zero(): + """A line without a field must report None — 0 mW is a real (wrong) + measurement, None is honestly missing.""" s = _parse_line("nothing useful here") - assert s.gpu_util_pct == 0.0 - assert s.temp_c == 0.0 - assert s.power_mw == 0.0 + assert s.gpu_util_pct is None + assert s.temp_c is None + assert s.power_mw is None def test_monitor_safe_to_start_stop_without_tegrastats(): """On a host without tegrastats the monitor must not crash.""" m = TegrastatsMonitor() - # We don't know the host; just ensure the lifecycle works regardless. m.start() sample = m.latest() - assert sample.gpu_util_pct >= 0.0 + assert sample.timestamp > 0 + m.stop() + + +def test_monitor_is_restartable(): + """Regression: stop() set a threading.Event that start() never cleared, so + a restarted monitor's reader thread exited immediately and silently + collected nothing. A stop/start cycle must yield a working session.""" + m = TegrastatsMonitor() + m.start() + m.stop() + assert not m._stop.is_set() or not m.available # unavailable hosts no-op + m.start() + if m.available: + assert m.running + assert not m._stop.is_set() m.stop() + assert m._proc is None and m._thread is None + + +def test_monitor_double_start_is_noop_not_duplicate(): + m = TegrastatsMonitor() + m.start() + proc = m._proc + m.start() # must not spawn a second process/thread + assert m._proc is proc + m.stop() + + +def test_monitor_validates_arguments(): + with pytest.raises(ValueError, match="interval_ms"): + TegrastatsMonitor(interval_ms=0) + with pytest.raises(ValueError, match="max_samples"): + TegrastatsMonitor(max_samples=0) def test_profiler_runs_on_cpu_without_tegrastats(): - backbone = TinyBackbone(num_classes=10) - exits = [ExitPoint("e0", "stage1", STAGE_CHANNELS["stage1"])] - heads = {ep.name: EarlyExitHead(ep.in_channels, 10) for ep in exits} - cfg = EarlyExitConfig( - backbone="tiny", - num_classes=10, - exit_points=exits, - confidence_thresholds=[0.8], - ) - wrapper = EarlyExitWrapper(backbone, heads, lambda x: x, cfg, input_shape=(1, 3, 32, 32)) profiler = JetsonProfiler() runs = profiler.profile( - wrapper, input_shape=(1, 3, 32, 32), num_warmup=2, num_runs=5, device="cpu" + _wrapper(), input_shape=(1, 3, 32, 32), num_warmup=2, num_runs=5, device="cpu" ) assert len(runs) == 5 for r in runs: assert r.latency_ms > 0 - assert r.gpu_util_pct >= 0.0 - assert r.temp_c >= 0.0 - assert r.power_mw >= 0.0 + # no tegrastats on this host -> telemetry is missing, not zero + if not profiler.monitor.available: + assert r.gpu_util_pct is None + assert r.power_mw is None + + +def test_profiler_reports_energy_summary_without_fabrication(): + profiler = JetsonProfiler() + runs, summary = profiler.profile_with_energy( + _wrapper(), input_shape=(1, 3, 32, 32), num_warmup=1, num_runs=4, device="cpu" + ) + assert len(runs) == 4 + assert summary.num_inferences == 4 + assert summary.window_seconds > 0 + if not profiler.monitor.available: + assert summary.num_power_samples == 0 + assert summary.avg_power_mw is None + assert summary.energy_mj is None + assert summary.energy_per_inference_mj is None + + +def test_profiler_validates_arguments(): + profiler = JetsonProfiler() + with pytest.raises(ValueError, match="num_runs"): + profiler.profile(_wrapper(), input_shape=(1, 3, 32, 32), num_runs=0, device="cpu") + with pytest.raises(ValueError, match="num_warmup"): + profiler.profile(_wrapper(), input_shape=(1, 3, 32, 32), num_warmup=-1, device="cpu") + + +# ---------------- energy integration (pure, no hardware) ---------------- + + +def _sample(t: float, mw: float | None) -> JetsonSample: + return JetsonSample(timestamp=t, gpu_util_pct=None, temp_c=None, power_mw=mw) + + +def test_integrate_energy_trapezoid(): + # constant 1000 mW over 2 s -> 2000 mJ; 10 inferences -> 200 mJ each + samples = [_sample(10.0, 1000.0), _sample(11.0, 1000.0), _sample(12.0, 1000.0)] + s = integrate_energy(samples, 10.0, 12.0, num_inferences=10, sampling_interval_ms=1000) + assert s.num_power_samples == 3 + assert s.avg_power_mw == 1000.0 + assert s.energy_mj == pytest.approx(2000.0) + assert s.energy_per_inference_mj == pytest.approx(200.0) + assert s.window_seconds == pytest.approx(2.0) + + +def test_integrate_energy_single_sample_yields_no_energy(): + """One instantaneous sample cannot be integrated; avg power is reported, + energy is honestly None.""" + s = integrate_energy([_sample(10.0, 500.0)], 9.0, 11.0, 5, 100) + assert s.avg_power_mw == 500.0 + assert s.energy_mj is None + assert s.energy_per_inference_mj is None + + +def test_integrate_energy_ignores_samples_outside_window_and_missing_power(): + samples = [ + _sample(5.0, 9999.0), # before window + _sample(10.0, 1000.0), + _sample(10.5, None), # power missing on this line + _sample(11.0, 1000.0), + _sample(20.0, 9999.0), # after window + ] + s = integrate_energy(samples, 10.0, 11.0, 2, 500) + assert s.num_power_samples == 2 + assert s.energy_mj == pytest.approx(1000.0) # 1000 mW * 1 s + assert s.energy_per_inference_mj == pytest.approx(500.0) + + +def test_integrate_energy_no_samples(): + s = integrate_energy([], time.monotonic(), time.monotonic() + 1, 3, 100) + assert s.num_power_samples == 0 + assert s.avg_power_mw is None and s.energy_mj is None diff --git a/tests/test_release_audit.py b/tests/test_release_audit.py new file mode 100644 index 0000000..f0c4d2d --- /dev/null +++ b/tests/test_release_audit.py @@ -0,0 +1,305 @@ +"""Independent release-readiness audit tests. + +Adversarial checks of the hardening claims, covering gaps the original +suites left: batched-path per-head temperatures, a decisive +global-vs-per-head temperature counterfactual, NaN-emitting heads during +calibration, calibration determinism, and proof that the staged runtime +genuinely skips later computation. +""" + +import pytest +import torch +import torch.nn as nn +from torch.utils.data import DataLoader, TensorDataset + +from earlyon.core.exit_head import EarlyExitHead +from earlyon.core.thresholds import ( + calibrate_thresholds, + collect_head_logits, + fit_head_temperatures, +) +from earlyon.core.types import EarlyExitConfig, ExitPoint +from earlyon.core.wrappers import EarlyExitWrapper +from earlyon.models import custom_ee +from earlyon.staged import staged_model +from tests.fixtures.tiny_models import STAGE_CHANNELS, TinyBackbone + + +def _build(thresholds=(0.95, 0.95)): + torch.manual_seed(0) + backbone = TinyBackbone(num_classes=10) + exits = [ + ExitPoint("e0", "stage1", STAGE_CHANNELS["stage1"]), + ExitPoint("e1", "stage2", STAGE_CHANNELS["stage2"]), + ] + heads = {ep.name: EarlyExitHead(ep.in_channels, 10) for ep in exits} + cfg = EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=exits, + confidence_thresholds=list(thresholds), + ) + return EarlyExitWrapper(backbone, heads, lambda x: x, cfg, input_shape=(1, 3, 32, 32)) + + +# ---------------- A. per-head temperatures ---------------- + + +def test_single_global_temperature_would_misroute_where_per_head_does_not(): + """The decisive counterfactual for per-head temperature scaling. + + Exit e0 is an overconfident junk head (huge constant logits for class 0); + the labels are the full network's own predictions, so routing to the + final classifier is always correct and routing through e0 is ~10%. + + A per-head fit assigns e0 a huge temperature, deflating its confidence + below the threshold: routed accuracy stays perfect. Reusing the final + head's (near-1.0) temperature for e0 — the pre-hardening behavior — + leaves e0 saturated, so it fires and accuracy collapses. One global + temperature cannot be right for both heads. + """ + torch.manual_seed(0) + model = _build(thresholds=(0.95, 0.95)) + with torch.no_grad(): + model.exit_heads["e0"].classifier[-1].weight.zero_() + model.exit_heads["e0"].classifier[-1].bias.zero_() + model.exit_heads["e0"].classifier[-1].bias[0] = 50.0 # saturated junk + # e1 flat so it never interferes + model.exit_heads["e1"].classifier[-1].weight.zero_() + model.exit_heads["e1"].classifier[-1].bias.zero_() + model.eval() + x = torch.randn(40, 3, 32, 32) + with torch.no_grad(): + y = model.backbone(x).argmax(dim=-1) + loader = DataLoader(TensorDataset(x, y), batch_size=8) + + cache = collect_head_logits(model, loader, "cpu") + fits = fit_head_temperatures(cache) + t_e0 = fits["e0"].temperature + t_final = fits["final"].temperature + assert t_e0 > 10 * t_final, (t_e0, t_final) # radically different calibration needs + + def routed_accuracy() -> float: + correct = 0 + for i in range(x.shape[0]): + result = model(x[i : i + 1], mode="inference") + correct += int(result.prediction.argmax().item() == int(y[i])) + return correct / x.shape[0] + + # per-head temperatures: e0 deflated below 0.95, never fires -> perfect + model.config.temperatures = {k: v.temperature for k, v in fits.items()} + assert routed_accuracy() == 1.0 + + # counterfactual: one global temperature (the final head's fit) for all + # heads — e0 stays saturated, fires its junk prediction, accuracy tanks + model.config.temperatures = {"e0": t_final, "e1": t_final, "final": t_final} + assert routed_accuracy() < 0.5 + + +def test_batched_routing_uses_per_head_temperatures(): + """The batched path shares the hook, but nothing pinned it: flattening + only e0 via its own temperature must stop the batch from exiting at e0.""" + model = _build(thresholds=(0.9, 0.9)) + with torch.no_grad(): + model.exit_heads["e0"].classifier[-1].weight.zero_() + model.exit_heads["e0"].classifier[-1].bias.zero_() + model.exit_heads["e0"].classifier[-1].bias[0] = 20.0 + model.exit_heads["e1"].classifier[-1].weight.zero_() + model.exit_heads["e1"].classifier[-1].bias.zero_() + model.eval() + x = torch.randn(4, 3, 32, 32) + + model.config.temperatures = {"e0": 1.0, "e1": 1.0, "final": 1.0} + assert model.forward_inference_batched(x).exit_taken == 0 + + model.config.temperatures = {"e0": 100.0, "e1": 1.0, "final": 1.0} + assert model.forward_inference_batched(x).exit_taken == -1 + + +# ---------------- C. calibration edge cases + determinism ---------------- + + +def test_calibration_is_deterministic(): + torch.manual_seed(5) + x = torch.randn(24, 3, 32, 32) + y = torch.randint(0, 10, (24,)) + + results = [] + for _ in range(2): + torch.manual_seed(7) + model = _build() + loader = DataLoader(TensorDataset(x, y), batch_size=8) + r = calibrate_thresholds(model, loader, target_accuracy_drop=0.05, device="cpu") + results.append((r.thresholds, r.enabled_exits, r.final_accuracy, r.exit_distribution)) + assert results[0] == results[1] + + +def test_calibration_survives_nan_emitting_head(): + """A head whose logits are NaN (e.g. numerically broken fine-tune) must + not crash the threshold search or get enabled: NaN comparisons are False, + so it never fires and stays disabled.""" + model = _build() + with torch.no_grad(): + model.exit_heads["e0"].classifier[-1].weight.fill_(float("nan")) + x = torch.randn(16, 3, 32, 32) + y = torch.randint(0, 10, (16,)) + loader = DataLoader(TensorDataset(x, y), batch_size=4) + result = calibrate_thresholds(model, loader, target_accuracy_drop=0.05, device="cpu") + assert result.enabled_exits[0] is False + + +def test_temperature_fit_rejects_nan_logits_cleanly(): + from earlyon.core.temperature import fit_temperature_full + + logits = torch.tensor([[float("nan"), 1.0], [0.0, 1.0]]) + with pytest.raises(ValueError, match="NaN or Inf"): + fit_temperature_full(logits, torch.tensor([0, 1])) + + +def test_calibration_single_class_loader_works(): + """All-one-class calibration data is degenerate but legal; it must not + crash and the routed model must still work.""" + model = _build() + x = torch.randn(12, 3, 32, 32) + y = torch.zeros(12, dtype=torch.long) + loader = DataLoader(TensorDataset(x, y), batch_size=4) + result = calibrate_thresholds(model, loader, target_accuracy_drop=0.05, device="cpu") + assert len(result.thresholds) == 2 + assert 0.0 <= result.final_accuracy <= 1.0 + + +# ---------------- I. staged runtime genuinely skips computation ---------------- + + +class _CountingModule(nn.Module): + def __init__(self, inner: nn.Module) -> None: + super().__init__() + self.inner = inner + self.calls = 0 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.calls += 1 + return self.inner(x) + + +def test_staged_runtime_skips_later_stages_when_exit_fires(): + """The whole point of staged deployment: when exit 0 fires, later stages' + modules are never executed (the eager wrapper claims this via exception + unwinding; the staged runtime must deliver it structurally).""" + torch.manual_seed(0) + counted_late = _CountingModule(nn.Conv2d(8, 16, 3, padding=1)) + backbone = nn.Sequential( + nn.Conv2d(3, 8, 3, padding=1), # 0: exit 0 attaches here + counted_late, # 1: must not run on early exit + nn.AdaptiveAvgPool2d(1), # 2 + nn.Flatten(), # 3 + nn.Linear(16, 10), # 4 + ) + wrapper = custom_ee( + backbone, + exit_layers=["0"], + num_classes=10, + input_shape=(1, 3, 16, 16), + confidence_thresholds=[0.0], # always fire at exit 0 + ) + staged = staged_model(wrapper) # build-time probe runs the model + staged.eval() + counted_late.calls = 0 + + result = staged.infer(torch.randn(1, 3, 16, 16)) + assert result.exit_taken == 0 + assert counted_late.calls == 0, "a fired exit must prevent later stages from running" + + # sanity: with the exit disabled the late stage does run + wrapper.config.enabled_exits = [False] + staged.infer(torch.randn(1, 3, 16, 16)) + assert counted_late.calls == 1 + + +# ---------------- G. dict-output custom backbones ---------------- + + +class _DictBlock(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv2d(3, 8, 3, padding=1) + + def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]: + f = self.conv(x) + return {"features": f, "aux": f.sum()} + + +class _DictOutputBackbone(nn.Module): + def __init__(self, num_classes: int = 10) -> None: + super().__init__() + self.block = _DictBlock() + self.head = nn.Linear(8, num_classes) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + f = self.block(x)["features"] + return self.head(f.mean(dim=(2, 3))) + + +def test_custom_ee_dict_output_via_feature_extractor(): + model = custom_ee( + _DictOutputBackbone(), + exit_layers=["block"], + num_classes=10, + input_shape=(1, 3, 16, 16), + feature_extractors={"block": lambda out: out["features"]}, + ) + model.eval() + model.config.confidence_thresholds = [0.0] + result = model(torch.randn(1, 3, 16, 16), mode="inference") + assert result.exit_taken == 0 + assert result.prediction.shape == (1, 10) + + +# ---------------- F. genuine prior-version checkpoint migration ---------------- + + +def test_real_v1_fixture_checkpoint_migrates(): + """tests/fixtures/v1_cifar_resnet20.pth was written by the ACTUAL + pre-hardening save_wrapper (git main), not a hand-built dict: no + format_version, scalar temperature 1.6, and the v1 disabled sentinel + (confidence threshold 1.0) at exit 1. Loading must migrate all of it + deterministically, and re-saving must produce a v2 file that round-trips. + """ + import warnings as _warnings + from pathlib import Path + + from earlyon.utils import load_wrapper, save_wrapper + + fixture = Path(__file__).parent / "fixtures" / "v1_cifar_resnet20.pth" + with _warnings.catch_warnings(record=True) as captured: + _warnings.simplefilter("always") + model = load_wrapper(fixture) + + assert model.config.backbone == "cifar_resnet20" + assert model.config.confidence_thresholds == [0.7, 1.0, 0.85] + # sentinel at exit 1 became an explicit disable; the others stay enabled + assert model.config.enabled_exits == [True, False, True] + # scalar temperature broadcast to every head + assert model.config.temperatures == {"e0": 1.6, "e1": 1.6, "e2": 1.6, "final": 1.6} + assert any("enabled_exits" in str(w.message) for w in captured) + + # migrated model routes without error + model.eval() + result = model(torch.randn(1, 3, 32, 32), mode="inference") + assert result.prediction.shape == (1, 10) + + # re-save writes v2 and round-trips silently (no migration warnings) + import tempfile + + with tempfile.TemporaryDirectory() as d: + p = Path(d) / "v2.pth" + save_wrapper(model, p) + payload = torch.load(p, map_location="cpu", weights_only=True) + assert payload["format_version"] == 2 + with _warnings.catch_warnings(record=True) as captured2: + _warnings.simplefilter("always") + reloaded = load_wrapper(p) + assert reloaded.config.enabled_exits == [True, False, True] + assert reloaded.config.temperatures == model.config.temperatures + migration_warnings = [w for w in captured2 if "migrated" in str(w.message)] + assert not migration_warnings diff --git a/tests/test_staged.py b/tests/test_staged.py new file mode 100644 index 0000000..4132c10 --- /dev/null +++ b/tests/test_staged.py @@ -0,0 +1,110 @@ +"""Staged-deployment contract: staged execution must match eager routing.""" + +import pytest +import torch +import torch.nn as nn + +from earlyon.models import custom_ee +from earlyon.staged import StagedModel, staged_model +from tests.fixtures.tiny_models import TinyBackbone + + +def _sequential_backbone(num_classes: int = 10) -> nn.Sequential: + torch.manual_seed(0) + return nn.Sequential( + nn.Conv2d(3, 8, 3, padding=1), + nn.ReLU(), + nn.Conv2d(8, 16, 3, padding=1), + nn.ReLU(), + nn.AdaptiveAvgPool2d(1), + nn.Flatten(), + nn.Linear(16, num_classes), + ) + + +def _wrapped(**kwargs): + backbone = _sequential_backbone() + return custom_ee( + backbone, + exit_layers=["0", "2"], + num_classes=10, + input_shape=(1, 3, 16, 16), + **kwargs, + ) + + +def test_staged_matches_eager_across_configs(): + torch.manual_seed(1) + wrapper = _wrapped() + staged = staged_model(wrapper) + wrapper.eval() + staged.eval() + + configs = [ + dict(thr=[0.99, 0.99], enabled=[True, True], temps=(1.0, 1.0, 1.0)), + dict(thr=[0.0, 0.0], enabled=[True, True], temps=(1.0, 1.0, 1.0)), # always exit 0 + dict(thr=[0.0, 0.0], enabled=[False, True], temps=(1.0, 1.0, 1.0)), # exit 0 off + dict(thr=[0.0, 0.0], enabled=[False, False], temps=(1.0, 1.0, 1.0)), # all off + dict(thr=[0.5, 0.5], enabled=[True, True], temps=(3.0, 0.5, 1.2)), # per-head temps + ] + for cfg in configs: + wrapper.config.confidence_thresholds = list(cfg["thr"]) + wrapper.config.enabled_exits = list(cfg["enabled"]) + t0, t1, tf = cfg["temps"] + wrapper.config.temperatures = {"e0": t0, "e1": t1, "final": tf} + for i in range(5): + x = torch.randn(1, 3, 16, 16, generator=torch.Generator().manual_seed(i)) + eager = wrapper(x, mode="inference") + staged_r = staged.infer(x) + assert staged_r.exit_taken == eager.exit_taken, cfg + assert torch.allclose(staged_r.prediction, eager.prediction, atol=1e-6), cfg + assert staged_r.estimated_backbone_flops_fraction == pytest.approx( + eager.estimated_backbone_flops_fraction + ) + + +def test_staged_matches_eager_entropy_policy(): + wrapper = _wrapped(routing_policy="entropy") + wrapper.config.entropy_thresholds = [1.5, 1.5] + staged = staged_model(wrapper) + wrapper.eval() + staged.eval() + for i in range(5): + x = torch.randn(1, 3, 16, 16, generator=torch.Generator().manual_seed(100 + i)) + eager = wrapper(x, mode="inference") + staged_r = staged.infer(x) + assert staged_r.exit_taken == eager.exit_taken + assert torch.allclose(staged_r.prediction, eager.prediction, atol=1e-6) + + +def test_staged_stage_specs_partition_the_backbone(): + wrapper = _wrapped() + staged = staged_model(wrapper) + all_modules = [m for spec in staged.specs for m in spec.modules] + assert all_modules == ["0", "1", "2", "3", "4", "5", "6"] # full cover, in order + assert staged.specs[0].exit_name == "e0" + assert staged.specs[-1].exit_name == "final" + + +def test_staged_rejects_non_sequential_backbone(): + model = custom_ee( + TinyBackbone(num_classes=10), + exit_layers=["stage1"], + num_classes=10, + input_shape=(1, 3, 32, 32), + ) + with pytest.raises(ValueError, match="nn.Sequential"): + staged_model(model) + + +def test_staged_rejects_batched_input(): + staged = staged_model(_wrapped()) + with pytest.raises(ValueError, match="batch-1"): + staged.infer(torch.randn(2, 3, 16, 16)) + + +def test_staged_model_validates_stage_count(): + wrapper = _wrapped() + good = staged_model(wrapper) + with pytest.raises(ValueError, match="stage"): + StagedModel(list(good.stages)[:-1], good.specs[:-1], wrapper.config, wrapper._flops_at) diff --git a/tests/test_temperature.py b/tests/test_temperature.py index f833992..3c7ee88 100644 --- a/tests/test_temperature.py +++ b/tests/test_temperature.py @@ -69,3 +69,42 @@ def test_nonfinite_logits_raise_instead_of_poisoning_temperature(): inf_logits[0, 0] = float("inf") with pytest.raises(ValueError, match="NaN or Inf"): fit_temperature(inf_logits, targets) + + +def test_fit_temperature_full_reports_status(): + """The status-carrying fit must return a finite positive temperature and a + coherent (converged, fallback) pair on well-behaved logits.""" + import torch + + from earlyon.core.temperature import fit_temperature_full + + torch.manual_seed(0) + logits = torch.randn(200, 10) * 3.0 + targets = torch.randint(0, 10, (200,)) + result = fit_temperature_full(logits, targets) + assert result.temperature > 0.0 + import math + + assert math.isfinite(result.temperature) + assert isinstance(result.converged, bool) + assert isinstance(result.fallback, bool) + assert not result.fallback # nothing degenerate here + + +def test_fit_temperature_full_never_returns_invalid_temperature(): + """Even on tiny/degenerate batches the fit must return a value that + EarlyExitConfig.validate() accepts (finite, > 0) — falling back to 1.0 + with fallback=True rather than emitting a poison value.""" + import math + + import torch + + from earlyon.core.temperature import fit_temperature_full + + torch.manual_seed(1) + # single-sample, extreme logits: classic LBFGS divergence bait + logits = torch.tensor([[1e4, -1e4, 0.0]]) + targets = torch.tensor([1]) + result = fit_temperature_full(logits, targets, max_outer_steps=3) + assert math.isfinite(result.temperature) + assert result.temperature > 0.0 diff --git a/tests/test_thresholds.py b/tests/test_thresholds.py index 090325c..df3a3dc 100644 --- a/tests/test_thresholds.py +++ b/tests/test_thresholds.py @@ -63,10 +63,10 @@ def test_calibration_thresholds_are_in_grid(): def test_calibration_with_fit_temperature_writes_to_config(): - """When fit_temperature=True, the fitted scalar must land on config.temperature - BEFORE the threshold grid search runs.""" + """When fit_temperature=True, per-head fitted temperatures must land on + config.temperatures BEFORE the threshold grid search runs.""" model = _build() - assert model.config.temperature == 1.0 + assert model.config.temperatures == {"e0": 1.0, "e1": 1.0, "final": 1.0} x = torch.randn(32, 3, 32, 32) y = torch.randint(0, 10, (32,)) val_loader = DataLoader(TensorDataset(x, y), batch_size=8) @@ -86,10 +86,16 @@ def test_calibration_with_fit_temperature_writes_to_config(): # no calibration signal, so we assert the contract, not a specific value). import math - assert result.fitted_temperature is not None - assert math.isfinite(result.fitted_temperature) - assert result.fitted_temperature > 0.0 - assert model.config.temperature == result.fitted_temperature + assert result.temperatures is not None + assert set(result.temperatures) == {"e0", "e1", "final"} + for temp in result.temperatures.values(): + assert math.isfinite(temp) and temp > 0.0 + assert model.config.temperatures == result.temperatures + # deprecated scalar alias still reports the final head's fit + assert result.fitted_temperature == result.temperatures["final"] + # per-head fit status is reported + assert result.temperature_fits is not None + assert set(result.temperature_fits) == {"e0", "e1", "final"} # thresholds still calibrated normally assert len(result.thresholds) == 2 @@ -119,21 +125,21 @@ def test_calibration_fit_temperature_warns_without_separate_loader(): def test_calibration_fit_temperature_runs_before_threshold_search(): - """The order matters: fit T, then sweep grids using logits/T. We verify by - asserting model.config.temperature is already non-1.0 by the time the search - begins — done indirectly by patching the grid evaluator.""" + """The order matters: fit per-head temperatures, then sweep grids using + logits/T. We verify by spying on the grid evaluator and asserting every + evaluate call already saw the fitted temperatures.""" model = _build() x = torch.randn(24, 3, 32, 32) y = torch.randint(0, 10, (24,)) val_loader = DataLoader(TensorDataset(x, y), batch_size=8) - observed_temps: list[float] = [] + observed_temps: list[dict[str, float]] = [] from earlyon.core import thresholds as thresholds_mod original_eval = thresholds_mod._EvalCache.evaluate def spy_eval(self): - observed_temps.append(self._model.config.temperature) + observed_temps.append(dict(self._model.config.temperatures)) return original_eval(self) thresholds_mod._EvalCache.evaluate = spy_eval @@ -150,11 +156,12 @@ def spy_eval(self): thresholds_mod._EvalCache.evaluate = original_eval # every evaluate call (baseline + grid sweep + final) must have seen the - # already-fitted temperature, proving the fit ran before the search. + # already-fitted per-head temperatures, proving the fit ran before the search. assert observed_temps, "expected at least one evaluate call" + assert result.temperatures is not None assert all( - t == result.fitted_temperature for t in observed_temps - ), f"expected all temps == fitted {result.fitted_temperature}; got {observed_temps}" + t == result.temperatures for t in observed_temps + ), f"expected all temps == fitted {result.temperatures}; got {observed_temps}" def _build_entropy(): @@ -183,6 +190,11 @@ def test_entropy_calibration_updates_entropy_thresholds_not_confidence(): import math model = _build_entropy() + # sharpen e0 so its entropy actually drops below the grid and it can fire; + # an untrained (near-uniform) head never fires and correctly stays disabled + with torch.no_grad(): + model.exit_heads["e0"].classifier[-1].weight.mul_(30.0) + model.exit_heads["e0"].classifier[-1].bias.mul_(30.0) x = torch.randn(24, 3, 32, 32) y = torch.randint(0, 10, (24,)) loader = DataLoader(TensorDataset(x, y), batch_size=4) @@ -195,8 +207,9 @@ def test_entropy_calibration_updates_entropy_thresholds_not_confidence(): assert model.config.entropy_thresholds == result.thresholds # confidence_thresholds must be untouched by an entropy calibration assert model.config.confidence_thresholds == [0.8, 0.8] - # they must have actually moved off the conservative 0.0 seed - assert any(t > 0.0 for t in result.thresholds) + # the sharpened exit must be enabled with a real grid threshold + assert result.enabled_exits[0] is True + assert result.thresholds[0] > 0.0 h_max = math.log(10) assert all(0.0 <= t <= h_max + 1e-9 for t in result.thresholds) @@ -357,12 +370,11 @@ def test_budget_calibration_trivial_budget_keeps_exits_disabled(): assert result.budget_met is True assert result.thresholds == [1.0, 1.0] - # NOT asserting avg_computation_used == 1.0: the routing hook fires on - # `confidence >= threshold`, so a float32-saturated head (confidence 1.0) - # still fires at the clamped disabled value 1.0 — the library's documented - # convention, shared with calibrate_thresholds. The budget contract is - # only comp <= target. - assert result.avg_computation_used <= 1.0 + assert result.enabled_exits == [False, False] + # explicit enablement: a disabled exit can no longer fire even with a + # float32-saturated softmax (confidence exactly 1.0), so the average + # compute is exactly the full backbone. + assert result.avg_computation_used == 1.0 def test_budget_calibration_validates_target(): @@ -405,9 +417,9 @@ def test_budget_calibration_entropy_policy_writes_entropy_field(): def test_calibration_without_fit_temperature_leaves_config_unchanged(): - """The default path (fit_temperature=False) must not touch config.temperature.""" + """The default path (fit_temperature=False) must not touch config.temperatures.""" model = _build() - model.config.temperature = 1.5 # user-provided + model.config.temperatures = {"e0": 1.5, "e1": 1.5, "final": 1.5} # user-provided x = torch.randn(16, 3, 32, 32) y = torch.randint(0, 10, (16,)) val_loader = DataLoader(TensorDataset(x, y), batch_size=4) @@ -419,7 +431,7 @@ def test_calibration_without_fit_temperature_leaves_config_unchanged(): device="cpu", ) - assert model.config.temperature == 1.5 + assert model.config.temperatures == {"e0": 1.5, "e1": 1.5, "final": 1.5} def test_calibration_rejects_empty_grid(): @@ -462,19 +474,126 @@ def test_eval_cache_matches_real_router_exactly(): cache = _EvalCache(model, loader, "cpu") settings = [ - ([2.0, 2.0], [-1.0, -1.0], 1.0), # non-firing search seeds - ([1.0, 1.0], [0.0, 0.0], 1.0), # disabled values (saturated head fires) - ([0.8, 0.5], [0.5, 1.2], 1.0), # mid-grid - ([0.5, 0.5], [2.0, 2.0], 2.5), # aggressive + temperature + # (conf_thr, ent_thr, temps, enabled) + ([1.0, 1.0], [0.0, 0.0], (1.0, 1.0, 1.0), [False, False]), # all disabled + ([1.0, 1.0], [0.0, 0.0], (1.0, 1.0, 1.0), [True, True]), # knife-edge thresholds + ([0.8, 0.5], [0.5, 1.2], (1.0, 1.0, 1.0), [True, True]), # mid-grid + ([0.5, 0.5], [2.0, 2.0], (2.5, 0.7, 1.3), [True, True]), # per-head temps + ([0.5, 0.5], [2.0, 2.0], (2.5, 0.7, 1.3), [False, True]), # first exit off ] - for conf_thr, ent_thr, temp in settings: + for conf_thr, ent_thr, temps, enabled in settings: model.config.confidence_thresholds = list(conf_thr) model.config.entropy_thresholds = list(ent_thr) - model.config.temperature = temp + model.config.temperatures = {"e0": temps[0], "e1": temps[1], "final": temps[2]} + model.config.enabled_exits = list(enabled) real = _evaluate_with_router(model, loader, "cpu") sim = cache.evaluate() # accuracy must match exactly (integer counts); avg compute may # differ by float summation order, so allow one ulp of slack assert sim[0] == real[0] and abs(sim[1] - real[1]) < 1e-12, ( - f"policy={policy} conf={conf_thr} ent={ent_thr} T={temp}: " f"sim={sim} real={real}" + f"policy={policy} conf={conf_thr} ent={ent_thr} T={temps} " + f"enabled={enabled}: sim={sim} real={real}" ) + + +# ---------------- staged pipeline (v0.3) ---------------- + + +def test_calibration_rejects_empty_loader(): + """An empty calibration loader must fail fast, not silently produce a + policy chosen from zero samples (or hang).""" + import pytest + + model = _build() + empty = DataLoader(TensorDataset(torch.empty(0, 3, 32, 32), torch.empty(0, dtype=torch.long))) + with pytest.raises(ValueError, match="no batches"): + calibrate_thresholds(model, empty, device="cpu") + + +def test_calibration_result_carries_rich_metadata(): + model = _build() + x = torch.randn(16, 3, 32, 32) + y = torch.randint(0, 10, (16,)) + loader = DataLoader(TensorDataset(x, y), batch_size=4) + result = calibrate_thresholds(model, loader, target_accuracy_drop=0.05, device="cpu") + assert result.num_samples == 16 + assert result.objective == "accuracy_budget" + assert result.method == "greedy-coordinate-grid" + assert result.schema_version == 2 + assert len(result.enabled_exits) == 2 + assert abs(sum(result.exit_distribution.values()) - 1.0) < 1e-9 + assert abs(result.accuracy_delta - (result.baseline_accuracy - result.final_accuracy)) < 1e-12 + assert result.target_accuracy_drop == 0.05 + # enablement written back to the model config + assert model.config.enabled_exits == result.enabled_exits + + +def test_calibration_disables_exit_with_no_passing_threshold(): + """A head that is pure noise (uniform logits never reach the grid) must end + up explicitly disabled, not parked on a sentinel threshold.""" + model = _build() + _flatten_head(model, "e0") + _flatten_head(model, "e1") + x = torch.randn(16, 3, 32, 32) + model.eval() + with torch.no_grad(): + y = model.backbone(x).argmax(dim=-1) # full-network labels: baseline acc 1.0 + loader = DataLoader(TensorDataset(x, y), batch_size=4) + result = calibrate_thresholds(model, loader, target_accuracy_drop=0.0, device="cpu") + assert result.enabled_exits == [False, False] + assert model.config.enabled_exits == [False, False] + assert result.exit_distribution == {"final": 1.0} + + +def test_fit_head_temperatures_can_differ_per_head(): + """Heads with different miscalibration must receive different fitted + temperatures — the P0 defect was one temperature fit from the final + classifier reused everywhere.""" + from earlyon.core.thresholds import collect_head_logits, fit_head_temperatures + + torch.manual_seed(3) + model = _build() + # make e0 wildly overconfident (large weight scale); leave e1/final alone + with torch.no_grad(): + model.exit_heads["e0"].classifier[-1].weight.mul_(30.0) + model.exit_heads["e0"].classifier[-1].bias.mul_(30.0) + x = torch.randn(64, 3, 32, 32) + y = torch.randint(0, 10, (64,)) + loader = DataLoader(TensorDataset(x, y), batch_size=16) + + cache = collect_head_logits(model, loader, "cpu") + fits = fit_head_temperatures(cache) + assert set(fits) == {"e0", "e1", "final"} + t_e0 = fits["e0"].temperature + t_final = fits["final"].temperature + # the overconfident head needs a much larger temperature than the final head + assert t_e0 > t_final * 2, (t_e0, t_final) + + +def test_collect_head_logits_shapes_and_order(): + from earlyon.core.thresholds import collect_head_logits + + model = _build() + x = torch.randn(12, 3, 32, 32) + y = torch.randint(0, 10, (12,)) + loader = DataLoader(TensorDataset(x, y), batch_size=4) + cache = collect_head_logits(model, loader, "cpu") + assert cache.head_names == ["e0", "e1", "final"] + assert [t.shape for t in cache.logits] == [torch.Size([12, 10])] * 3 + assert cache.num_samples == 12 + assert cache.flops[-1].item() == 1.0 + # cumulative fractions are monotone non-decreasing + flops = cache.flops.tolist() + assert all(a <= b for a, b in zip(flops, flops[1:])) + + +def test_collect_head_logits_restores_training_mode(): + from earlyon.core.thresholds import collect_head_logits + + model = _build() + model.train() + x = torch.randn(4, 3, 32, 32) + y = torch.randint(0, 10, (4,)) + loader = DataLoader(TensorDataset(x, y), batch_size=4) + collect_head_logits(model, loader, "cpu") + assert model.training is True diff --git a/tests/test_types.py b/tests/test_types.py index 3324156..a83a81d 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -82,3 +82,152 @@ def test_config_rejects_num_classes_below_two(): for bad in (1, 0, -3): with pytest.raises(ValueError, match="num_classes"): EarlyExitConfig(backbone="tiny", num_classes=bad, exit_points=exits) + + +# ---------------- centralized validation (v0.3) ---------------- + + +def test_config_migrates_legacy_scalar_temperature(): + """The legacy scalar `temperature` is broadcast deterministically onto + every head (exits + final) at construction.""" + cfg = EarlyExitConfig( + backbone="tiny", num_classes=10, exit_points=_two_exits(), temperature=2.0 + ) + assert cfg.temperatures == {"e0": 2.0, "e1": 2.0, "final": 2.0} + + +def test_config_rejects_scalar_and_per_head_temperatures_together(): + with pytest.raises(ValueError, match="temperature"): + EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=_two_exits(), + temperature=2.0, + temperatures={"e0": 1.0, "e1": 1.0, "final": 1.0}, + ) + + +def test_config_accepts_per_head_temperatures(): + cfg = EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=_two_exits(), + temperatures={"e0": 1.5, "e1": 0.8, "final": 1.1}, + ) + assert cfg.temperature_for("e1") == 0.8 + assert cfg.temperature_for("final") == 1.1 + assert cfg.head_names == ["e0", "e1", "final"] + + +@pytest.mark.parametrize("bad", [0.0, -1.0, float("nan"), float("inf")]) +def test_config_rejects_invalid_temperatures(bad): + with pytest.raises(ValueError, match="temperatures"): + EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=_two_exits(), + temperatures={"e0": bad, "e1": 1.0, "final": 1.0}, + ) + + +def test_config_rejects_wrong_temperature_keys(): + with pytest.raises(ValueError, match="temperatures"): + EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=_two_exits(), + temperatures={"e0": 1.0, "final": 1.0}, # e1 missing + ) + + +def test_config_rejects_duplicate_exit_names(): + exits = [ExitPoint("e0", "stage1", 16), ExitPoint("e0", "stage2", 32)] + with pytest.raises(ValueError, match="unique"): + EarlyExitConfig(backbone="tiny", num_classes=10, exit_points=exits) + + +def test_config_rejects_duplicate_layer_names(): + exits = [ExitPoint("e0", "stage1", 16), ExitPoint("e1", "stage1", 32)] + with pytest.raises(ValueError, match="unique"): + EarlyExitConfig(backbone="tiny", num_classes=10, exit_points=exits) + + +def test_config_rejects_empty_exit_or_layer_name(): + with pytest.raises(ValueError, match="non-empty"): + EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=[ExitPoint("", "stage1", 16)], + ) + with pytest.raises(ValueError, match="non-empty"): + EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=[ExitPoint("e0", "", 16)], + ) + + +def test_config_rejects_reserved_final_exit_name(): + with pytest.raises(ValueError, match="reserved"): + EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=[ExitPoint("final", "stage1", 16)], + ) + + +@pytest.mark.parametrize("bad", [-0.1, 1.5, float("nan"), float("inf")]) +def test_config_rejects_out_of_range_confidence_thresholds(bad): + with pytest.raises(ValueError, match="confidence_thresholds"): + EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=_two_exits(), + confidence_thresholds=[bad, 0.8], + ) + + +@pytest.mark.parametrize("bad", [-0.1, float("nan"), float("inf")]) +def test_config_rejects_invalid_entropy_thresholds(bad): + with pytest.raises(ValueError, match="entropy_thresholds"): + EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=_two_exits(), + entropy_thresholds=[bad, 0.5], + ) + + +def test_config_rejects_negative_or_nonfinite_loss_weights(): + for bad_weights in ([-0.1, 0.6, 0.5], [float("nan"), 0.5, 0.5], [0.0, 0.0, 0.0]): + with pytest.raises(ValueError, match="loss_weights"): + EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=_two_exits(), + loss_weights=bad_weights, + ) + + +def test_config_rejects_wrong_enabled_exits_length(): + with pytest.raises(ValueError, match="enabled_exits"): + EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=_two_exits(), + enabled_exits=[True], + ) + + +def test_config_defaults_enable_every_exit(): + cfg = EarlyExitConfig(backbone="tiny", num_classes=10, exit_points=_two_exits()) + assert cfg.enabled_exits == [True, True] + + +def test_validate_catches_post_hoc_mutation(): + """validate() is re-callable after mutation — the wrapper and checkpoint + loader use it as their trust boundary.""" + cfg = EarlyExitConfig(backbone="tiny", num_classes=10, exit_points=_two_exits()) + cfg.temperatures["e0"] = -3.0 + with pytest.raises(ValueError, match="temperatures"): + cfg.validate() diff --git a/tests/test_utils.py b/tests/test_utils.py index 4f6382b..55bd453 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -120,3 +120,200 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: save_wrapper(model, tmp_path / "m.pth") messages = [str(w.message) for w in captured if issubclass(w.category, UserWarning)] assert any("load_wrapper" in m for m in messages), messages + + +# ---------------- checkpoint format v2 ---------------- + + +def test_checkpoint_v2_payload_shape(tmp_path): + """v2 files carry format_version, library version, and the full routing + config including exit points — the documented public contract.""" + import torch + + from earlyon import __version__ + from earlyon.utils import build_model, save_wrapper + + model = build_model("resnet18", num_classes=10, pretrained=False) + path = tmp_path / "v2.pth" + save_wrapper(model, path) + payload = torch.load(path, map_location="cpu", weights_only=True) + assert payload["format_version"] == 2 + assert payload["earlyon_version"] == __version__ + cfg = payload["config"] + assert cfg["enabled_exits"] == [True, True] + assert set(cfg["temperatures"]) == {"e0", "e1", "final"} + assert [p["layer_name"] for p in cfg["exit_points"]] == ["layer2", "layer3"] + + +def test_v1_checkpoint_migrates_scalar_temperature_and_sentinels(tmp_path): + """Unversioned (v1) checkpoints must load: the scalar temperature is + broadcast per head, and the legacy 'disabled' sentinel (confidence 1.0 on + the active policy) becomes an explicit enabled_exits=False with a warning.""" + import warnings as _warnings + + import torch + + from earlyon.utils import build_model, load_wrapper + + model = build_model("resnet18", num_classes=10, pretrained=False) + path = tmp_path / "v1.pth" + payload = { + "state_dict": model.state_dict(), + "config": { + "backbone": "resnet18", + "num_classes": 10, + "routing_policy": "confidence", + "confidence_thresholds": [0.7, 1.0], # 1.0 was the v1 disabled sentinel + "entropy_thresholds": [0.5, 0.5], + "loss_weights": [0.2, 0.3, 0.5], + "temperature": 1.7, + }, + } + torch.save(payload, path) + + with _warnings.catch_warnings(record=True) as captured: + _warnings.simplefilter("always") + loaded = load_wrapper(path) + assert loaded.config.temperatures == {"e0": 1.7, "e1": 1.7, "final": 1.7} + assert loaded.config.enabled_exits == [True, False] + messages = [str(w.message) for w in captured if issubclass(w.category, UserWarning)] + assert any("enabled_exits" in m for m in messages), messages + + +def test_v1_checkpoint_entropy_sentinel_migrates(tmp_path): + """For an entropy-routed v1 checkpoint the disabled sentinel is 0.0.""" + import torch + + from earlyon.utils import build_model, load_wrapper + + model = build_model("resnet18", num_classes=10, pretrained=False) + path = tmp_path / "v1e.pth" + payload = { + "state_dict": model.state_dict(), + "config": { + "backbone": "resnet18", + "num_classes": 10, + "routing_policy": "entropy", + "confidence_thresholds": [0.85, 0.80], + "entropy_thresholds": [0.4, 0.0], # 0.0 was the v1 entropy disabled sentinel + "loss_weights": [0.2, 0.3, 0.5], + "temperature": 1.0, + }, + } + torch.save(payload, path) + loaded = load_wrapper(path) + assert loaded.config.routing_policy == "entropy" + assert loaded.config.enabled_exits == [True, False] + + +def test_v1_checkpoint_invalid_temperature_falls_back(tmp_path): + """A corrupt v1 temperature (NaN/negative) must not poison the load; it + falls back to 1.0 with a warning, matching the v1 runtime guard.""" + import warnings as _warnings + + import torch + + from earlyon.utils import build_model, load_wrapper + + model = build_model("resnet18", num_classes=10, pretrained=False) + path = tmp_path / "v1bad.pth" + payload = { + "state_dict": model.state_dict(), + "config": { + "backbone": "resnet18", + "num_classes": 10, + "confidence_thresholds": [0.7, 0.8], + "loss_weights": [0.2, 0.3, 0.5], + "temperature": float("nan"), + }, + } + torch.save(payload, path) + with _warnings.catch_warnings(record=True) as captured: + _warnings.simplefilter("always") + loaded = load_wrapper(path) + assert loaded.config.temperatures == {"e0": 1.0, "e1": 1.0, "final": 1.0} + assert any("temperature" in str(w.message) for w in captured) + + +def test_checkpoint_from_newer_format_rejected(tmp_path): + import pytest + import torch + + from earlyon.utils import build_model, load_wrapper, save_wrapper + + model = build_model("resnet18", num_classes=10, pretrained=False) + path = tmp_path / "future.pth" + save_wrapper(model, path) + payload = torch.load(path, map_location="cpu", weights_only=True) + payload["format_version"] = 99 + torch.save(payload, path) + with pytest.raises(ValueError, match="format_version"): + load_wrapper(path) + + +def test_malformed_checkpoint_rejected(tmp_path): + import pytest + import torch + + from earlyon.utils import load_wrapper + + path = tmp_path / "junk.pth" + torch.save({"weights": {}}, path) + with pytest.raises(ValueError, match="not an earlyon file"): + load_wrapper(path) + + +def test_checkpoint_exit_point_mismatch_rejected(tmp_path): + """A v2 checkpoint whose recorded exit points disagree with the rebuilt + factory's placement must fail loudly instead of loading mis-wired routing.""" + import pytest + import torch + + from earlyon.utils import build_model, load_wrapper, save_wrapper + + model = build_model("resnet18", num_classes=10, pretrained=False) + path = tmp_path / "moved.pth" + save_wrapper(model, path) + payload = torch.load(path, map_location="cpu", weights_only=True) + payload["config"]["exit_points"][0]["layer_name"] = "layer1" # tampered placement + torch.save(payload, path) + with pytest.raises(ValueError, match="exit points"): + load_wrapper(path) + + +def test_custom_checkpoint_loads_via_factory(tmp_path): + """custom_ee models round-trip when the caller supplies the factory.""" + import torch + + from earlyon.models import custom_ee + from earlyon.utils import load_wrapper, save_wrapper + from tests.fixtures.tiny_models import TinyBackbone + + torch.manual_seed(0) + + def factory(): + return custom_ee( + TinyBackbone(num_classes=10), + exit_layers=["stage1", "stage2"], + num_classes=10, + input_shape=(1, 3, 32, 32), + ) + + model = factory() + model.config.confidence_thresholds = [0.42, 0.9] + model.config.enabled_exits = [True, False] + with torch.no_grad(): + next(iter(model.exit_heads["e0"].parameters())).fill_(0.5) + + path = tmp_path / "custom.pth" + import warnings as _warnings + + with _warnings.catch_warnings(): + _warnings.simplefilter("ignore") # the save-time custom_ee warning + save_wrapper(model, path) + + loaded = load_wrapper(path, factory=factory) + assert loaded.config.confidence_thresholds == [0.42, 0.9] + assert loaded.config.enabled_exits == [True, False] + p = next(iter(loaded.exit_heads["e0"].parameters())) + assert torch.allclose(p, torch.full_like(p, 0.5)) diff --git a/tests/test_wrappers.py b/tests/test_wrappers.py index 14ebc9d..1de2701 100644 --- a/tests/test_wrappers.py +++ b/tests/test_wrappers.py @@ -91,20 +91,128 @@ def test_batched_inference_prediction_is_grad_free_without_external_no_grad(): @pytest.mark.parametrize("bad_temp", [float("nan"), float("inf"), float("-inf")]) def test_nonfinite_temperature_does_not_poison_softmax(bad_temp): - """A NaN/Inf config.temperature (e.g. from a degenerate fit) must not turn - every softmax into NaN: _safe_temperature falls back to 1.0 so routing stays - well-defined and confidence is finite. Inf is the consequential case — the - old `max(inf, 1e-6)` left T=inf, making every softmax uniform.""" + """Config validation rejects non-finite temperatures at the trust + boundaries, but a direct post-hoc mutation of config.temperatures bypasses + it; the runtime guard must still fall back to 1.0 so routing stays + well-defined and confidence is finite.""" import math wrapper, _ = _build(thresholds=(1.0, 1.0, 1.0)) wrapper.eval() - wrapper.config.temperature = bad_temp + for head in wrapper.config.temperatures: + wrapper.config.temperatures[head] = bad_temp # bypasses validate() result = wrapper(torch.randn(1, 3, 32, 32), mode="inference") assert math.isfinite(result.confidence) assert torch.isfinite(result.prediction).all() +def test_disabled_exit_cannot_fire_even_at_saturated_confidence(): + """The P0 enablement invariant: an exit with enabled_exits[i]=False must + never fire, even when its softmax confidence is exactly 1.0 (float32 + saturation) and its threshold is 1.0 — the numerical edge case the old + 'threshold sentinel means disabled' convention got wrong.""" + wrapper, _ = _build(thresholds=(1.0, 1.0, 1.0)) + # saturate exit e0: huge constant logit -> softmax max exactly 1.0 + with torch.no_grad(): + wrapper.exit_heads["e0"].classifier[-1].bias.zero_() + wrapper.exit_heads["e0"].classifier[-1].bias[0] = 1e3 + wrapper.eval() + x = torch.randn(1, 3, 32, 32) + + # sanity: with the exit enabled and threshold 1.0, saturation DOES fire + wrapper.config.enabled_exits = [True, False, False] + assert wrapper(x, mode="inference").exit_taken == 0 + + # disabled: must fall through to the final classifier + wrapper.config.enabled_exits = [False, False, False] + result = wrapper(x, mode="inference") + assert result.exit_taken == -1 + assert result.estimated_backbone_flops_fraction == 1.0 + + # batched path honors the same invariant + batched = wrapper.forward_inference_batched(torch.randn(4, 3, 32, 32)) + assert batched.exit_taken == -1 + + +def test_disabled_exit_cannot_fire_at_entropy_exactly_zero(): + """Same invariant for the entropy policy: a saturated head has entropy + exactly 0.0, which meets any threshold >= 0; disabling the exit must win.""" + backbone = TinyBackbone(num_classes=10) + exits = [ + ExitPoint("e0", "stage1", STAGE_CHANNELS["stage1"]), + ExitPoint("e1", "stage2", STAGE_CHANNELS["stage2"]), + ] + heads = {ep.name: EarlyExitHead(ep.in_channels, 10) for ep in exits} + cfg = EarlyExitConfig( + backbone="tiny", + num_classes=10, + exit_points=exits, + routing_policy="entropy", + entropy_thresholds=[0.0, 0.0], + ) + wrapper = EarlyExitWrapper(backbone, heads, lambda t: t, cfg, input_shape=(1, 3, 32, 32)) + with torch.no_grad(): + wrapper.exit_heads["e0"].classifier[-1].bias.zero_() + wrapper.exit_heads["e0"].classifier[-1].bias[0] = 1e3 + wrapper.eval() + x = torch.randn(1, 3, 32, 32) + + wrapper.config.enabled_exits = [True, False] + assert wrapper(x, mode="inference").exit_taken == 0 # entropy 0.0 <= 0.0 fires + + wrapper.config.enabled_exits = [False, False] + assert wrapper(x, mode="inference").exit_taken == -1 + + +def test_disabled_exit_still_produces_training_logits(): + """Enablement is a routing concept: training mode must return logits for + every head, disabled or not, so a disabled head keeps learning.""" + wrapper, _ = _build() + wrapper.config.enabled_exits = [False, True, False] + outputs = wrapper(torch.randn(2, 3, 32, 32), mode="training") + assert len(outputs) == 4 # all 3 exits + final, regardless of enablement + + +def test_per_head_temperature_is_applied_at_the_right_exit(): + """Each exit must use its own temperature. A very high temperature flattens + exit e0's softmax below the threshold (no fire); resetting only e0's + temperature to 1.0 makes the same input fire at e0 again.""" + wrapper, _ = _build(thresholds=(0.9, 0.9, 0.9)) + with torch.no_grad(): + wrapper.exit_heads["e0"].classifier[-1].bias.zero_() + wrapper.exit_heads["e0"].classifier[-1].bias[0] = 20.0 # confident but not saturated + # make later heads and final produce flat logits so nothing else fires + wrapper.exit_heads["e1"].classifier[-1].weight.zero_() + wrapper.exit_heads["e1"].classifier[-1].bias.zero_() + wrapper.exit_heads["e2"].classifier[-1].weight.zero_() + wrapper.exit_heads["e2"].classifier[-1].bias.zero_() + wrapper.eval() + x = torch.randn(1, 3, 32, 32) + + wrapper.config.temperatures = {"e0": 1.0, "e1": 1.0, "e2": 1.0, "final": 1.0} + assert wrapper(x, mode="inference").exit_taken == 0 + + # flatten ONLY e0 via its per-head temperature; other heads unchanged + wrapper.config.temperatures = {"e0": 100.0, "e1": 1.0, "e2": 1.0, "final": 1.0} + result = wrapper(x, mode="inference") + assert result.exit_taken == -1, "e0's own temperature must govern e0's routing" + + +def test_wrapper_rejects_mismatched_exit_heads(): + """Heads must exactly match the configured exit points.""" + backbone = TinyBackbone(num_classes=10) + exits = [ExitPoint("e0", "stage1", STAGE_CHANNELS["stage1"])] + cfg = EarlyExitConfig(backbone="tiny", num_classes=10, exit_points=exits) + with pytest.raises(ValueError, match="exit_heads"): + EarlyExitWrapper(backbone, {}, lambda t: t, cfg, input_shape=(1, 3, 32, 32)) + heads = { + "e0": EarlyExitHead(STAGE_CHANNELS["stage1"], 10), + "extra": EarlyExitHead(STAGE_CHANNELS["stage2"], 10), + } + with pytest.raises(ValueError, match="exit_heads"): + EarlyExitWrapper(backbone, heads, lambda t: t, cfg, input_shape=(1, 3, 32, 32)) + + def test_safe_temperature_clamps_to_positive_finite(): """Unit-level guarantees of the temperature guard.""" from earlyon.core.wrappers import _safe_temperature @@ -112,8 +220,10 @@ def test_safe_temperature_clamps_to_positive_finite(): assert _safe_temperature(float("nan")) == 1.0 assert _safe_temperature(float("inf")) == 1.0 assert _safe_temperature(float("-inf")) == 1.0 - assert _safe_temperature(0.0) == 1e-6 # non-positive clamps up - assert _safe_temperature(-5.0) == 1e-6 + # non-positive values fall back to the no-op temperature 1.0 — NOT a tiny + # positive clamp, which would produce an artificially razor-sharp softmax + assert _safe_temperature(0.0) == 1.0 + assert _safe_temperature(-5.0) == 1.0 assert _safe_temperature(2.5) == 2.5 # finite positive passes through