From 3e69bbc7d3453dc10ad3ea77e5ed26a4665efe7c Mon Sep 17 00:00:00 2001
From: sohams25 <86193836+sohams25@users.noreply.github.com>
Date: Wed, 22 Jul 2026 02:45:13 +0530
Subject: [PATCH 01/16] feat: per-head temperatures, explicit exit enablement,
staged calibration, checkpoint v2
- EarlyExitConfig: per-head temperatures mapping (exits + final), explicit
enabled_exits booleans, centralized validate() re-checked at wrapper
construction and checkpoint load; legacy scalar temperature broadcast
deterministically
- routing: disabled exits never fire (even at confidence 1.0 / entropy 0.0);
each head's softmax uses its own temperature; invalid runtime temperatures
fall back to 1.0 instead of a razor-sharp 1e-6 clamp
- calibration: staged pipeline (collect logits once, fit per-head
temperatures with convergence status, greedy grid search on cache,
evaluate); exits that never fire stay disabled; empty loaders fail fast;
CalibrationResult carries exit distribution, sample count, objective,
method and schema version
- checkpoints: format_version 2 with exit-point metadata and cross-check;
deterministic v1 migration (scalar temperature broadcast, sentinel
thresholds -> enabled_exits=False) with warnings; custom_ee loads via a
user-supplied factory
- InferenceResult/BatchedInferenceResult: computation_used renamed to
estimated_backbone_flops_fraction with a read alias
---
earlyon/core/temperature.py | 50 +++-
earlyon/core/thresholds.py | 548 +++++++++++++++++++++++-------------
earlyon/core/types.py | 160 ++++++++++-
earlyon/core/wrappers.py | 53 ++--
earlyon/utils.py | 181 ++++++++++--
tests/test_cli.py | 9 +-
tests/test_integration.py | 4 +-
tests/test_temperature.py | 39 +++
tests/test_thresholds.py | 185 +++++++++---
tests/test_types.py | 149 ++++++++++
tests/test_utils.py | 197 +++++++++++++
tests/test_wrappers.py | 124 +++++++-
12 files changed, 1395 insertions(+), 304 deletions(-)
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..2f98761 100644
--- a/earlyon/core/wrappers.py
+++ b/earlyon/core/wrappers.py
@@ -16,24 +16,26 @@
import torch.nn as nn
from earlyon.core.flops import per_layer_flops
-from earlyon.core.types import EarlyExitConfig, InferenceResult
+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:
@@ -97,6 +99,16 @@ def __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
# ordered list of (exit_idx, exit_name, layer_name)
@@ -195,18 +207,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 +255,20 @@ 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
+
head = self.exit_heads[exit_name]
logits = head(output)
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
@@ -323,19 +342,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/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/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_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_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
From 5f03a98910ebba73641d53f2e950a40f9f095c06 Mon Sep 17 00:00:00 2001
From: sohams25 <86193836+sohams25@users.noreply.github.com>
Date: Wed, 22 Jul 2026 02:49:56 +0530
Subject: [PATCH 02/16] feat: fair benchmark runner with identical boundaries
and inputs
- benchmark_models: one measurement core; every compared model (backbone,
early-exit wrapper, static baseline, quantised variant) sees the exact
same preloaded sample sequence, warmup, eval/inference_mode setup, device
placement and per-iteration synchronization
- explicit boundary labels (model-only vs end-to-end incl. H2D copy);
loader-vs-noise input source labelled; accuracy reported alongside speed
for labelled loaders; batch-1 validated as the primary mode
- BenchmarkResult carries dtype, input shape, warmup, sync policy, and
avg_estimated_flops_fraction (computation_used kept as read alias)
- CLI benchmark + scripts use the fair runner; docs/benchmarks.json
restructured with pre-v0.3 records quarantined under legacy_v0_2 and
marked not comparable
---
docs/benchmarks.json | 183 +++++++-------
earlyon/benchmarking/__init__.py | 4 +
earlyon/benchmarking/throughput.py | 373 +++++++++++++++++++----------
earlyon/cli.py | 26 +-
scripts/re_evaluate.py | 103 ++++----
scripts/run_benchmarks.py | 51 ++--
tests/test_benchmarks.py | 123 ++++++++++
7 files changed, 573 insertions(+), 290 deletions(-)
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/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/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..b58a899 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(
{
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..6aca7a7 100644
--- a/scripts/run_benchmarks.py
+++ b/scripts/run_benchmarks.py
@@ -20,11 +20,13 @@
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
-from earlyon.benchmarking import benchmark_backbone, benchmark_wrapper, evaluate
+from earlyon.benchmarking import benchmark_models, 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
+METHODOLOGY = "v0.3-fair-runner"
+
OUT_DIR = ROOT / "docs"
OUT_DIR.mkdir(exist_ok=True)
OUT_JSON = OUT_DIR / "benchmarks.json"
@@ -102,36 +104,46 @@ 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 +152,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 +166,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/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)
From 5724afe4103873e281bfb9d428c18ec3663412a4 Mon Sep 17 00:00:00 2001
From: sohams25 <86193836+sohams25@users.noreply.github.com>
Date: Wed, 22 Jul 2026 02:53:26 +0530
Subject: [PATCH 03/16] =?UTF-8?q?feat:=20honest=20compute=20accounting=20?=
=?UTF-8?q?=E2=80=94=20lazy=20FLOPs=20estimate=20with=20provenance?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
- estimate_layer_flops returns FlopsEstimate (fractions + method + reliable
+ notes + excludes_exit_heads); per_layer_flops kept as thin wrapper
- reused/multi-call leaf modules are detected via a counting probe and
produce a warned, low-confidence uniform fallback instead of a precise but
wrong cumulative fraction
- EarlyExitWrapper computes the estimate lazily on first inference (cached),
so constructing large models (ViT) is fast; probe runs on the backbone's
current device and is materialised before routing flags flip
- estimator explicitly excludes exit-head FLOPs and routing overhead;
documented on FlopsEstimate and InferenceResult
---
earlyon/core/flops.py | 149 ++++++++++++++++++++++++++++++++++-----
earlyon/core/wrappers.py | 41 +++++++++--
tests/test_flops.py | 91 ++++++++++++++++++++++++
3 files changed, 258 insertions(+), 23 deletions(-)
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/wrappers.py b/earlyon/core/wrappers.py
index 2f98761..7d1caa8 100644
--- a/earlyon/core/wrappers.py
+++ b/earlyon/core/wrappers.py
@@ -15,7 +15,7 @@
import torch
import torch.nn as nn
-from earlyon.core.flops import per_layer_flops
+from earlyon.core.flops import FlopsEstimate, estimate_layer_flops
from earlyon.core.types import FINAL_HEAD, EarlyExitConfig, InferenceResult
if TYPE_CHECKING:
@@ -114,9 +114,12 @@ def __init__(
# 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
@@ -136,6 +139,34 @@ 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]
+ # device-neutral probe: the analysis input is created on CPU, so
+ # run it against the backbone wherever it currently lives.
+ device = next(self.backbone.parameters(), torch.empty(0)).device
+ self._flops_estimate = estimate_layer_flops(
+ self.backbone,
+ layer_names,
+ self._input_shape,
+ device=str(device),
+ )
+ 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]:
@@ -184,6 +215,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
@@ -328,6 +360,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:
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
From c73254410d6b94ece79becf86725a70568c1e7ae Mon Sep 17 00:00:00 2001
From: sohams25 <86193836+sohams25@users.noreply.github.com>
Date: Wed, 22 Jul 2026 02:55:58 +0530
Subject: [PATCH 04/16] fix: restartable tegrastats monitor, honest
power/energy semantics
- TegrastatsMonitor.start() clears the stale stop event and sample buffer
(a restarted monitor previously collected nothing), no-ops on double
start, degrades to unavailable on spawn failure, and stop() closes the
child's stdout pipe
- missing telemetry fields parse to None instead of invented zeros;
JetsonRun documents its power as instantaneous
- integrate_energy(): pure trapezoidal integration over the timed window;
EnergySummary distinguishes instantaneous power, window-average power and
integrated energy, and reports energy-per-inference only when a real
integral and inference count exist
- profiler exposes profile_with_energy; CLI profile reports the energy
block and nulls for missing data; parser covers Orin-style rails
---
earlyon/benchmarking/jetson_profiler.py | 217 ++++++++++++++++++++----
earlyon/cli.py | 24 ++-
tests/test_jetson_profiler.py | 168 +++++++++++++++---
3 files changed, 349 insertions(+), 60 deletions(-)
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/cli.py b/earlyon/cli.py
index b58a899..5729b24 100644
--- a/earlyon/cli.py
+++ b/earlyon/cli.py
@@ -303,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,
@@ -317,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/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
From 7b84818aad910f5e0d26b71737b3dd6833f2a868 Mon Sep 17 00:00:00 2001
From: sohams25 <86193836+sohams25@users.noreply.github.com>
Date: Wed, 22 Jul 2026 03:00:08 +0530
Subject: [PATCH 05/16] feat: safer, more capable custom_ee
- example_args/example_kwargs for the inspection forward (single-tensor
path retained); dry-run tensors created on the backbone's own device (or
an explicit device) instead of blindly probing a CUDA model with CPU input
- per-layer feature_extractors convert tuple/dict/odd-rank outputs into the
Tensor the exit head consumes, applied in the dry run AND in the routing
hook (wrapper feature_adapters)
- the dry run now validates each exit layer executes exactly once and in
the listed forward order; reused, missing or out-of-order exits fail with
actionable errors instead of building a mis-routing wrapper
- freshly built exit heads placed on the backbone's device; wrapper FLOPs
estimate degrades to a warned low-confidence uniform fallback when the
probe cannot run the backbone
---
earlyon/core/wrappers.py | 49 ++++++++--
earlyon/models/custom.py | 189 ++++++++++++++++++++++++++++++---------
tests/test_custom_ee.py | 153 +++++++++++++++++++++++++++++++
3 files changed, 339 insertions(+), 52 deletions(-)
diff --git a/earlyon/core/wrappers.py b/earlyon/core/wrappers.py
index 7d1caa8..e0f422f 100644
--- a/earlyon/core/wrappers.py
+++ b/earlyon/core/wrappers.py
@@ -10,6 +10,7 @@
import math
import threading
+import warnings
from typing import TYPE_CHECKING, Any, Callable, Iterator, Mapping, Optional
import torch
@@ -94,6 +95,7 @@ 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
@@ -110,6 +112,17 @@ def __init__(
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)]
@@ -151,15 +164,31 @@ def flops_estimate(self) -> FlopsEstimate:
"""
if self._flops_estimate is None:
layer_names = [ep.layer_name for ep in self.config.exit_points]
- # device-neutral probe: the analysis input is created on CPU, so
- # run it against the backbone wherever it currently lives.
+ # the probe input follows the backbone wherever it currently lives
device = next(self.backbone.parameters(), torch.empty(0)).device
- self._flops_estimate = estimate_layer_flops(
- self.backbone,
- layer_names,
- self._input_shape,
- device=str(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
@@ -294,8 +323,10 @@ def hook(module: nn.Module, inputs: Any, output: torch.Tensor) -> torch.Tensor:
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
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/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)
From 43e82b1a1cd119b976e1637d1cf7346b532f6175 Mon Sep 17 00:00:00 2001
From: sohams25 <86193836+sohams25@users.noreply.github.com>
Date: Wed, 22 Jul 2026 03:02:58 +0530
Subject: [PATCH 06/16] feat: staged deployment contract with verified
reference splitter
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
- earlyon.staged: StageSpec/Stage/StagedModel protocol — ordered stages
emitting (continuation_features, exit_logits), routing applied between
stages so later stages genuinely never run
- staged_model() splits Sequential backbones (custom_ee shape), refuses
anything it cannot prove, and self-checks staged==eager on a probe input;
tests pin equivalence across policies, per-head temperatures, enablement
- docs/STAGED_DEPLOYMENT.md: protocol, scope, per-stage ONNX/TensorRT
procedure for later Jetson validation (no fabricated numbers)
- export_all_exits_to_onnx alias makes the static multi-output export's
compute-everything semantics explicit in the API name
---
docs/STAGED_DEPLOYMENT.md | 80 ++++++++++++++
earlyon/onnx.py | 7 ++
earlyon/staged.py | 219 ++++++++++++++++++++++++++++++++++++++
tests/test_staged.py | 110 +++++++++++++++++++
4 files changed, 416 insertions(+)
create mode 100644 docs/STAGED_DEPLOYMENT.md
create mode 100644 earlyon/staged.py
create mode 100644 tests/test_staged.py
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/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..8da7091
--- /dev/null
+++ b/earlyon/staged.py
@@ -0,0 +1,219 @@
+"""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
+
+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
+ features = self.stages[idx].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/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)
From a75058628a42c38157e9d59a63e42477c70f2982 Mon Sep 17 00:00:00 2001
From: sohams25 <86193836+sohams25@users.noreply.github.com>
Date: Wed, 22 Jul 2026 03:03:35 +0530
Subject: [PATCH 07/16] fix: cast ModuleList access in staged runtime for
strict mypy
---
earlyon/staged.py | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/earlyon/staged.py b/earlyon/staged.py
index 8da7091..6970fec 100644
--- a/earlyon/staged.py
+++ b/earlyon/staged.py
@@ -29,6 +29,7 @@
from __future__ import annotations
from dataclasses import dataclass
+from typing import cast
import torch
import torch.nn as nn
@@ -113,7 +114,8 @@ def infer(self, x: torch.Tensor) -> InferenceResult:
ep = cfg.exit_points[idx]
if not cfg.enabled_exits[idx]:
# skip the head entirely; still run the trunk to continue
- features = self.stages[idx].trunk(features)
+ stage = cast(Stage, self.stages[idx])
+ features = stage.trunk(features)
continue
features, logits = self.stages[idx](features)
temp = _safe_temperature(cfg.temperatures[ep.name])
From db3832e1c19409cf7463d42fe23ca1200cb05ee4 Mon Sep 17 00:00:00 2001
From: sohams25 <86193836+sohams25@users.noreply.github.com>
Date: Wed, 22 Jul 2026 03:10:43 +0530
Subject: [PATCH 08/16] docs: honest repositioning, contracts, release draft,
overnight report
- README: deployment-oriented framing, batch-1 primary use case, legacy
labels on pre-v0.3 numbers, routing-overhead and estimator limitations,
static-baseline guidance, per-head temperature + enablement description
- DESIGN_DECISIONS: per-head temperature section now matches reality; new
sections on explicit enablement and estimated-FLOPs naming
- docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md: split discipline, pipeline
stages, estimator limits, benchmark fairness rules
- examples: test-label hygiene note in 01; 02 rewritten on the fair runner
and energy-aware profiler; pyproject description de-hyped
- CHANGELOG unreleased, RELEASE_NOTES_DRAFT.md (0.3.0, not tagged),
OVERNIGHT_REPORT.md
---
CHANGELOG.md | 71 ++++++
OVERNIGHT_REPORT.md | 238 +++++++++++++++++++++
README.md | 94 +++++---
RELEASE_NOTES_DRAFT.md | 66 ++++++
docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md | 93 ++++++++
docs/DESIGN_DECISIONS.md | 68 ++++--
examples/01_train_resnet50_cifar10.py | 12 +-
examples/02_jetson_deployment.py | 45 ++--
pyproject.toml | 2 +-
9 files changed, 627 insertions(+), 62 deletions(-)
create mode 100644 OVERNIGHT_REPORT.md
create mode 100644 RELEASE_NOTES_DRAFT.md
create mode 100644 docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9cb062d..81a93c3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,7 +4,78 @@ format follows [keep a changelog](https://keepachangelog.com/en/1.1.0/).
## unreleased
+### 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
+- **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/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/README.md b/README.md
index 2725820..18b1771 100644
--- a/README.md
+++ b/README.md
@@ -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.
+
@@ -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%
@@ -238,19 +257,31 @@ 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.
## Contributing
@@ -291,7 +322,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..894472d
--- /dev/null
+++ b/RELEASE_NOTES_DRAFT.md
@@ -0,0 +1,66 @@
+# earlyon 0.3.0 — draft release notes (NOT published, NOT tagged)
+
+Suggested version: **0.3.0** (routing/calibration semantics changed;
+checkpoints migrate automatically, but calibrated behavior can differ —
+disabled exits are now strictly disabled).
+
+## Headline
+
+earlyon 0.3 is a correctness release. It makes the calibration methodology
+defensible (per-head temperatures, explicit exit enablement, 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").
+
+## Breaking / behavioral changes (all with migration paths)
+
+- One temperature **per head** (`config.temperatures`), fitted per head when
+ `fit_temperature=True`. The scalar `temperature` field is a constructor
+ convenience only; mutate `temperatures` after construction.
+- Disabled exits are explicit booleans (`config.enabled_exits`). Threshold
+ sentinels are gone; a saturated softmax can no longer fire a "disabled"
+ exit. v1 checkpoints migrate their sentinels with a warning.
+- Calibration leaves never-firing exits **disabled** instead of enabled at
+ an idle threshold; `CalibrationResult` gained enablement, exit
+ distribution, sample count, objective and method fields.
+- `computation_used` → `estimated_backbone_flops_fraction` (read alias
+ kept). It is an estimate excluding exit heads and routing overhead.
+- `BenchmarkResult.avg_computation_used` → `avg_estimated_flops_fraction`
+ (read alias kept); results now carry boundary/sync/input-source metadata.
+- Checkpoints are `format_version: 2`; v1 loads migrate deterministically.
+ `custom_ee` models reload via `load_wrapper(path, factory=...)`.
+- Invalid config now raises from centralized `EarlyExitConfig.validate()`
+ (previously e.g. a negative temperature was clamped to `1e-6`).
+
+## New
+
+- `benchmark_models`: fair multi-model comparison (backbone / wrapper /
+ static baseline / quantised variant) on identical samples, with accuracy
+ reported for labelled loaders and explicit model-only vs end-to-end
+ boundaries.
+- `earlyon.staged`: staged-deployment protocol + verified reference splitter
+ for Sequential backbones; staged execution matches eager routing and
+ genuinely skips later stages. `docs/STAGED_DEPLOYMENT.md` documents the
+ contract and the Jetson/TensorRT procedure.
+- Jetson: restartable `TegrastatsMonitor`, `None` for missing telemetry,
+ `profile_with_energy` with real integrated energy (`EnergySummary`).
+- `custom_ee`: example args/kwargs, device inference, per-layer feature
+ extractors, and dry-run validation (exactly-once, in-order exits).
+- Lazy `FlopsEstimate` with provenance; reused-module detection.
+- `docs/CALIBRATION_AND_BENCHMARK_CONTRACT.md`.
+
+## Release checklist (when actually releasing)
+
+- [ ] bump `__version__` and `pyproject.toml` to 0.3.0; move `unreleased` in
+ CHANGELOG under 0.3.0 with a date
+- [ ] regenerate `docs/benchmarks.json` `runs` with
+ `scripts/run_benchmarks.py` on real hardware; update the README table
+ from those records and drop the legacy labels for the new numbers
+- [ ] run the Jetson procedure in `docs/STAGED_DEPLOYMENT.md` §Jetson if a
+ device is available (optional for release, required before any Jetson
+ claim)
+- [ ] full gate: `pytest -m "not gpu"`, `pytest -m gpu` (CUDA host),
+ `mypy earlyon`, `ruff check earlyon tests`, `black --check`,
+ `isort --check-only`
+- [ ] `python -m build` + twine check; follow `docs/PYPI_RELEASE_CHECKLIST.md`
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/examples/01_train_resnet50_cifar10.py b/examples/01_train_resnet50_cifar10.py
index e1b1152..9308718 100644
--- a/examples/01_train_resnet50_cifar10.py
+++ b/examples/01_train_resnet50_cifar10.py
@@ -45,14 +45,18 @@ def main():
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,
)
- 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}")
+ 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}")
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..9c4d3af 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,28 +31,39 @@ 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("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: {speedup:.2f}x")
- print(f" avg compute used: {wrap_r.avg_computation_used:.2%}")
+ 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)
+ 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),
- "power_median_mw": statistics.median(r.power_mw for r in runs),
- "temp_median_c": statistics.median(r.temp_c 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))
diff --git a/pyproject.toml b/pyproject.toml
index 8cc8e47..1b89487 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
[project]
name = "earlyon"
version = "0.2.0"
-description = "Production-ready early exit for PyTorch CV models."
+description = "Deployment-oriented early-exit toolkit for PyTorch CV models."
readme = "README.md"
requires-python = ">=3.10"
license = { file = "LICENSE" }
From 5d821aa36382db7e362146565ccc070d2ddfb83e Mon Sep 17 00:00:00 2001
From: sohams25 <86193836+sohams25@users.noreply.github.com>
Date: Wed, 22 Jul 2026 14:48:44 +0530
Subject: [PATCH 09/16] chore: bring scripts and examples up to the lint/format
gate
ruff --fix (unused imports), noqa for sys.path-dependent imports, rename
ambiguous lambda vars, black/isort over scripts/ and examples/ which CI's
earlyon+tests scope never checked
---
examples/01_train_resnet50_cifar10.py | 26 ++++++++----
examples/02_jetson_deployment.py | 49 ++++++++++++++---------
scripts/run_benchmarks.py | 57 +++++++++++++++++----------
3 files changed, 86 insertions(+), 46 deletions(-)
diff --git a/examples/01_train_resnet50_cifar10.py b/examples/01_train_resnet50_cifar10.py
index 9308718..5e026fd 100644
--- a/examples/01_train_resnet50_cifar10.py
+++ b/examples/01_train_resnet50_cifar10.py
@@ -36,13 +36,20 @@ 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
@@ -51,12 +58,17 @@ def main():
# 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"val baseline acc: {result.baseline_accuracy:.4f} "
+ f"val routed acc: {result.final_accuracy:.4f} "
+ f"est. FLOPs fraction: {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 9c4d3af..9f7d6b7 100644
--- a/examples/02_jetson_deployment.py
+++ b/examples/02_jetson_deployment.py
@@ -34,17 +34,25 @@ def main():
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,
+ 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" 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}")
@@ -55,16 +63,21 @@ def main():
)
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))
+ 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/scripts/run_benchmarks.py b/scripts/run_benchmarks.py
index 6aca7a7..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,10 @@
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
-from earlyon.benchmarking import benchmark_models, 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"
@@ -33,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},
}
@@ -52,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,
),
)
@@ -71,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,
),
)
@@ -84,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(
@@ -109,12 +118,18 @@ def run_one(backbone: str) -> dict:
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,
+ loader=test_loader,
+ 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,
+ input_shape=(1, 3, 224, 224),
+ device=device,
+ num_warmup=50,
+ num_runs=300,
)
save_wrapper(model, OUT_DIR / f"{backbone}_cifar10.pth")
From de8ce1c950f82805a1cfdfc96c2555575a315ed2 Mon Sep 17 00:00:00 2001
From: sohams25 <86193836+sohams25@users.noreply.github.com>
Date: Wed, 22 Jul 2026 14:51:28 +0530
Subject: [PATCH 10/16] =?UTF-8?q?test:=20independent=20release=20audit=20?=
=?UTF-8?q?=E2=80=94=20adversarial=20verification=20of=20hardening=20claim?=
=?UTF-8?q?s?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
- decisive counterfactual: a single global temperature misroutes where
per-head temperatures do not (junk overconfident head vs final classifier)
- batched routing pinned to per-head temperatures (previously untested path)
- calibration determinism, NaN-emitting head safely disabled, clean NaN
rejection in temperature fit, single-class calibration data
- staged runtime proven to skip later stages structurally (counting module)
- custom_ee dict-output feature extractor
- genuine v1 checkpoint fixture written by the actual pre-hardening
save_wrapper (git main) + full migration/round-trip test
---
tests/test_release_audit.py | 305 ++++++++++++++++++++++++++++++++++++
1 file changed, 305 insertions(+)
create mode 100644 tests/test_release_audit.py
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
From 1a0d90fbe55418209b1deeb951603492b051f28e Mon Sep 17 00:00:00 2001
From: sohams25 <86193836+sohams25@users.noreply.github.com>
Date: Wed, 22 Jul 2026 14:55:21 +0530
Subject: [PATCH 11/16] =?UTF-8?q?feat:=20packaging=20validation=20?=
=?UTF-8?q?=E2=80=94=20sdist=20pruning,=20wheel=20smoke=20script,=20v1=20f?=
=?UTF-8?q?ixture,=20evidence=20runner?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
- MANIFEST.in prunes tests/ from the sdist (it shipped a partial test tree)
- scripts/smoke_test.py: post-install smoke for an installed wheel, run from
outside the repo (import/version, model build, training+inference forwards,
deterministic calibration, v2 round-trip, v1 migration, staged==eager,
benchmark smoke)
- tests/fixtures/v1_cifar_resnet20.pth: genuine pre-v0.3 checkpoint written
by the historical save_wrapper, force-added past *.pth ignore with a
provenance README (the committed migration test needs it on fresh clones)
- scripts/evidence_run.py: bounded seeded CUDA evidence experiment (CIFAR-10,
explicit disjoint train/temperature/calibration/test splits, epoch caps,
wall-clock budget, fair 3-model benchmark incl. static MobileNetV2
baseline, JSON output; negative results preserved)
---
MANIFEST.in | 1 +
scripts/evidence_run.py | 285 +++++++++++++++++++++++++++
scripts/smoke_test.py | 155 +++++++++++++++
tests/fixtures/README.md | 11 ++
tests/fixtures/v1_cifar_resnet20.pth | Bin 0 -> 1182235 bytes
5 files changed, 452 insertions(+)
create mode 100644 MANIFEST.in
create mode 100644 scripts/evidence_run.py
create mode 100644 scripts/smoke_test.py
create mode 100644 tests/fixtures/README.md
create mode 100644 tests/fixtures/v1_cifar_resnet20.pth
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/scripts/evidence_run.py b/scripts/evidence_run.py
new file mode 100644
index 0000000..c1de90a
--- /dev/null
+++ b/scripts/evidence_run.py
@@ -0,0 +1,285 @@
+"""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),
+ "temperature": mk(eval_train, temp_idx, batch_size, False),
+ "calibration": mk(eval_train, calib_idx, batch_size, False),
+ "test_batched": DataLoader(test, batch_size=batch_size, 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 ----
+ baseline = build_model("mobilenetv2", num_classes=10, pretrained=True)
+ t0 = time.time()
+ stage1_train_backbone(
+ baseline.backbone, data["train"], epochs=args.epochs_baseline, lr=0.002, device=device
+ )
+ log["stages"]["baseline_seconds"] = round(time.time() - t0, 1)
+ 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/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 0000000000000000000000000000000000000000..b1e155498bfefcf27d10e2f50461cf58b173d811
GIT binary patch
literal 1182235
zcmdRWbzBu)*EU#)fru>zDk>OA*O|4!!oWNzV%><8D2Qh>*o6TYh%FYj*xk(Dc6WDo
zch|RO=71dT!~NdR`+R?VpZ-|;n#J1d+Oy`GeP;H|p`CyJ0wyMf3!D7yUqzF`CjFg)
zL!%->XfO?r>l-dQJ2nUl5knmM^@;XtRKTPQ{u`NZn9js5s0enC6GO!C;IODrQD+(&
z8XFxQ9x6t~_KhpjmWGAX@USNMnLX${?0{o1?}vD%oN608WI_1({-uqpomJer!Yg#p}#lZ!i^gMcVtD*!f%f=a>JBx}tXe*8b**
z@mIyh^^(U|T+~@;btQsIx?1dxfRj8NB;G`?r%iP~m)LCI-f6KADQvRl6
zb*0r*Wkg+Bt*%^9$=s=WMuo)b${Qqd)>TjwRn#Z4mJ?M{6IB*S}3qwS&s%m9T#Z)!7;3bJf|a`ReHN
z)s^$rQ}fjqbq-oxgCLvS`TEAk1osRPLwkkC1&b7qd3c!4QO&DAZmnYbMrI#2C*`oY
z>70=Y4~Yv>+1slWuApXjzw>B{sbzYcGKKRPv
z7j@ot{*$u~zD6YxL>*{#B&d>6k{l}VF`)3!`Kq;PsIN^Uxi*c}oJ~YsQ?1S~sA8Tu
z6*)9B$mpqSu4ZhZ&)8DV=&xpMCF)vhb!~#G=b6z!7i|qPd+FM#ncM3#Yvs%VYUU22
zuA^2L7*y?l%q$C}Q+DEf81Lk$3sN(8)@SY_XYQ(I?k4KGYjuAG+2@(rm{fxG3{F7`
zgOjcYGN3R*kReMLJw;upRu}dU!f?)-UcwCr?5vB>D|L)lvsUH`pHGnsY$7*i__}FJf-36phzR$Fr}-mznXG@KIK3;U6-h4
zPSR&imNSo2GmjT_6STUCd5Xi?!I(Hw^bC4&c<3e}1BxRR8M4GNS=3F@>ZTeMhm+x~
zG0kwmp1SFJam>JcIm9tDhd5@b<4+THv$eW8M#bUekYn1IYslfHo2M3KzP>2wa#0qj
zNf(N`MOxirqvBALD&kmTkkZ*vw^U8JOrLVOoH9dAxkA*f)aq6l6^EMAKpd+LQ#dkqKds@ta*$9~M0LmUTkh~uC-
z{zIbfuvT{@PjO^(95v){(;ZWba$H}O6LL{bs!30Yy3<F
z;<#&=*;98<&3s><`GK7Ip_=)TsC%r{J;_rX#%GSFdIr5XymZfy0mbng8M4IjLe#y~
z>RuTYhqK|#@!DX(E{?i4dU3qPd^yDNE{8bYtKON|9pNxtFkKv!o&d-J%PP#8@
zQNHSn@=Y$vcQxq`QTJ1;`(;!dYEot9_-&ZdnVR4=vyzpXB2p96dEUv=CaOicxuTPAAk1V+|=?XkpnWUqDlHM0(012xZEm*s7Gp5Q(%nS}nCeYg9>w
zl&By>j!(DTyRb-!XsY|OOM20`&qc-UI
zqd!^=H`P6@j!2|6&}ykI+B~(G&0Q0t47ojME!n4(wKgJsS?v(YWwpmvC98u-q;=71
zX+5-gY7|S0H%YWU#u!q0QU}?mP&GiLr*cFnQ#oO)QaK|MsS8>ybw!(}O#jS2^2A2n
zFxHULi@M7`h0+6&p3)PcOzDNKO6iSAq#CqZDxl3%rhh19nF59yFuFQYBKs6ZA4Gaa
zUxYGaLu^&XMugh$HVwsCLrOOqCi@i1
za722_2!t|aB(^GLFGM1ZLaU{{(HfCyZc16EeK6FJ(Va%iK7}y`k)E+HLYXlZTa~dN
zB9ZFQYAHo)M5bAc#%9wv?4oDZ%hZF4vQLp|JR)Syl4*b36zKr8S~?KzKgcwDMjM22
zhP&oT2g^Q`sjWzdAVdaPrZ&pL5gm#iyVf?Q`t<`k47b%?9gaw(BhYH;NVIwCG@CpD
zqYTNt=qTBz)O9o>eO<>OliG!?CuPDY!jPXEkp%5*vfV+|>t=v3LKP)huq#tkc;TYRKqJ=g2;VaV{b~<2-~i<9uvY#&kp?U4T|g7oyEmr^ct#
zMc75ptm@Rwg)WwTicXgxLgp-;F2zleE<>xO%h3ka_}kKmd``smilt&`yogX;B&8V`
zY`A}}bcO6gvE~!$N`y$G&Z(yI<;E)X$X5cmA#bk6O?7wIAQI_Xv|73j?f*1{tVK||
z9)k^s;6^vdKBdGP5$Q|3387r#&Dg4(TM&tKD_SkxhBl9bn!Ve2izi6!N-l$rNq
zt1=%zB+`Ruwe%3$AiIA)#7^0ZGqR2kV-G`OPkKc5Da1z+>4}dal!=dHs}i3;B+`>;
zwe%F)AjkiS*x0mt8hh#opqH~3JtO-RIiE#@1IUu|IouTKd9+%30j<$#*Ko;&Uc@+q
zJ#=@Zmt>zh?OsNR3^}IVE9l8N?Ow%gb&szh66tlcT6zPm(P>v*F3_JwH!;eP+={Df>5S<
zimgiZ43S8mqt((EXpK(0DwTn1Ut+8wr3-x}`xMI8i1d_i5XzKqu~jMGArk3(v|9QB
zt0MbejDMogd>2cg>v^kbSC73nD~@96Bw8
zo}4;0!)RKF;zOEJs<+_%@Rwc7UB+`;-
zwbTl2o;uZE9+biuLn=>NTJ|YaWf1A9$|97h%3-Tgl}Cg(wrI7qBHBE4`scG+xjeAO
zSVKxLT1oaPl$8$C=j8Zde|Qd`-lFxEt*
zXRL)#W~_~^%4mm3r1ogFv<})lb!z$r;rw+I&(gtX?
z)Df*wof;kLS~?ubO{fmTaB(HhmMy8fY`
zPrWe8klclO%RZ&98btcK3JB%80=6m{K_pTiv|8$m)~HTZGG#t(h%ts#uC$TtQ>Ypv
z(o;1-C{s1XR;BVoB+_PRwX``}qdHZo4Cd1o7;8xBMqA20h0-69p0X7}nX)yuDrFl)
zB5jLSOWUC}s#8NsS*PtW)R57gYGt3o7=TF6*a4x;*b!TmF%XeRJE7H5`N^44ovMt+
z=F`sDMbE6)sR!*M`xKpaMTE>*I_-v=BJGY=OaDUq4?4}B(tcGz1|s
zj4`0{bfo=cpF*WWq^F_?WvV!ARVopYNaN9JX@9hN>h#a#
zrp%`UFxHULi4K%~3gsX~ddk5FWy&Ggs+2u%2E;L#8DLNg82${2VIvzJg
zIsvVgPDE=&r!KiyuTn71aMxVvB-w{LZH_m?2$8|g-#2SnU^04g>U0WjtGhZCkw~YZ
z)zax`jp$US&#`{BQ+5a}sbAe1RrVyjZFLWFm*Xti_=S|d8mO)2YiEruF0deU{WPhnh-NYA(dq0G1u
zTa|GWB9U%JtEF4e8qsMMqp|sPD|XQ{>vifyx5++5r`r)BbCynb;HF4-qSew}X#YW{
z*;863#u@CImm`&AAL?{6-ufa$23e=R>hcvwk6mlud|As^yK!6H)jfzrx)-gM?n9fW
zPP57PW0WDe6FngNl)4^7q_68CgmPUEW2=%KK_t?nXtneh+B|itpHGitj3JdXJt6xP
zs*{NHRHqQiRHw03sm>r0=~=W|dJb)#I{h=bDf8)hj5VZmp%-MILU|F9p7IhxnesBW
zD&-YKBE5=MORu5LQ>T9@Wu0EfP(wyndPDXpj5iVK8E+w!8E<2>uQX(z|H2^d8zg
zb!vP*y^me=%&Jbk+~@<@r|9${B4p0e=_A|}>0`87`UI_!wIlhQ$X->+UcP#Y!G`6l-C;)kTOj>ZGcySG_=wd@Z0~zIut9>h8WmB+}Ptwe$_z|7i%>%U5qP*l-9Q
z^quTeO8g#?zQi98$|e4Yt;+cckw`zI)zU9$^Ejy4yN!3P@|D0>3^io*q~Bzp!uTDL
zp794lneivKD&sFiBK?h4OHJ@?#yk#c4n~8C))c!KGJDZ{vQJ^ok4Vp40HMrW5L=bG
z5F(M9q1DpDXpJmi860B6xfnv3xj42evjrlNmO!he
zmS~NvU*%#pHt&|ip863uDer+gIXY1**{KM+6e1i!mY_@HCO+{+i%)#f8l8A^Ent^*r5GiuZycN-tbLO?iZFQe3Arfh2v|3sPt!a0D2ed}#UX{y0x(zTk7pEI_l${Ev6CypQGeViu1zVNV6_H5Y&}ykW
zTBGNwA*Vd~dSGZSN_XlhI~7VVM0!eZgfgWDTa{8kBvL@Dr39_f6IG=&rdA*9qNmPM
zs|WR!or+o;B0}mcwKl>{eDI4FAN->I2eoE@Qa8o8Ts!AU{bZ-A)@BHiB8OU=qbH|Y
zTi~|3vn>&c)E}*uwnCeyTC>?(V^l77FWN?SDy40UNMG7^2<6hY$5v(2A`)o;S}pB>
zHcz$cFBCdrj3JkkBMp?D3RfpYdM_9{!9fVd(2ctErR&_x{KerCS
zs9fxBbg1l9N;?dZzO=&;%B3BFt;#kMkw_EJYUwDnM%AjaDRb*+jLF62PRGbjg=;J#
zJy#+^nJWofl`9#MNXMbo((!1Gs#WDOm|G`cY%Wd@I#G5ioGFO(oRbjBoT=EVoRblW
zbP8H6or>0|S`9g6wNAs(T$G-4y6jXaXCTs3&O|6v&caruOhY8n*=V(N4qBsXRVj_l
zt#h%9o;pjdUUZ)9RMa{j5mINVH61tcDKJ`m3XJw2)S5lJEy6g%opW}ii)E*()+Gp$
zB8OU+q9>?#VpJ}6C%Q^@Dy3bGNMG7D2<6hQ#a3lo
zhe)LB(Q4@iw0Wu(%PU{XZp4^eT+Vcp>{PfmBhqtiK`3)=#a89ohDfB_(Q4@qw0Wxa
z&pfBhtvfL`7pDu|B|8<)OhkH438Bo%uvIxZB9ZP!tEGF;=Bd^{oU&T?VrVW(SGrGj
zDwO*X=_wB&lqnBlt5P09B+|oZwe$$uJk@G^Zas=!^wfH_I=j(hvQts(aYRU+rPdR;
ziO+h`;${CyMYIo!a-0@Fl#
z9+A3N7Z8c`B3doIgf_@B7m*^5%lMmIG#>Pd>@@c;g15tn6qd2{8bY~l*RfR@ZXgor
zO|)8i3$1df!Usl)!M(ym!r~ml@rypuA#rh05mDjv&kGT^F*FyYC%q#(@!TtbH^c~)
zW94s|Nbli>{4PWyy^l!ceSk=$57BDrBeX$PanJsNSJ~0W7?_LIi$0N^ia?$sQb@2;U;iK7{XZLp_A=5vjZ%
z5Q+37S}pyA_MdmuMLpp^V_+^;C;CNpDo6AyBK;BlhEU$o@7St@KM;xZCt5B2h4ybd
ziZ@h7j_7X;%|(f?C*zkDjap}8nMsiW*vkEj#E>?7)o8|sd_AX0f<5#gOHTD)^b`_DV7KcXHOn2Xhm
zddg1ah9I?l&I!8ql)E>!sjEjg%Z;;J0@RTP>Gx|>X$;PzY^9eI}8V8@I~CL?~c;eN?@Cu
zf#v?5FYDP5ZL5U0%NcI)h32fh^0OS;UJ2G_2jf`t&h1PCl<*GO;W*~Jay!$GN^oFy
zFy=P?9n-9(LE1^_5tQ8{kCHgi&PrgHoPoxMtzQkJU6t@|Im3+&TUq;{-Id_KvV-#|
ziJIG)1}ot`vcvNziJaS+hA6>3vxD;}iBtBnG!0dHgk|^0qa@BWTnUWG8EC8|`kPo9
zsf72+8E&j3$}J0xQi6MD2j@`|HMcYEql8Cihv!ieIkz*7QG)ws2j@`|=RYNhReJQx
z?t$t5ettPqof1ZKhW-8g!Z%oHoDwMJ4E%fLtNIXO7{{2He87@pk&)BpXO|mpn%gXIc
zmnh*&v%`&6PR;F1mnp%^vxALRP7cmKk291WE3$iF`oEvYIeTDrRw{i~{Re%V=xU|U
zn(RI}GGnE8qHC4Fb=iS9AY-M^I*->Y;Tv*>8!Nqf9&c2FH)RLsQTnX&c(W3|B|AKi
z(yQn3RwZ~_c5oi0&v71aS92W1zkKC1XqF0rmYdM2*SJElFlGl~s8##k>R}#mDxj
zz0%`Djvn!GJ%bA2>4l%!5mAx!qZ055t)1cR`)70o$v5nVe-n&;k^S;CBR_UgDi)4!
zyP(s*h!Rd?@lSc-U#RLA8yyuoh<=m9bij5A7
zqd)N1-RV!XcFlYqMJWH8`Mkk+{MXFqxr_2ITv_%G?<>;SeuHv$|D|;Qt#*IcXJh8p
zj%Azq
z9@8(JhKTVroaV!>&3wM9zmAW0Xyl^Y{OikpUyhmhfQs3B!;?%#l790{C71=iitMAvy
z{@)ME-Blid&j6kLN|=fKllK6uP27wOpywblJXZb-bAFAS{ylG2#iFA`F*;lx!&taq
zK|TxQpS;IV-o(|&806Kl29AD>9RK41jELS2Gdt~slGaq@vmzEjl2tw+)5o%swP%yJEe}d#UV1X?luF!ZO9NhN
z=yT@dIa%o1Zvy*iRs&AA8V`3i`|xgu55kaWAU
z8WiADmpQPSK`ppl=Lr6!L%LLb?`pQ_NGxpW;>H34tCJxKZlGCEl1C<1;~&jW3M(U~
zvjcC5*bGlgIC{JuTkM+*>18g1_Rd)0@}N+*=JH`S!rp_AeDsl>-Lx&U
zRv}w>G0qN(4hU`vCxNipNx$fvVb{K1o>(!k)lnYT2pi07D4spf?ytYe95P=!_@U&=fY
zIy*9EmYT-y3`yq)eDJupYQ<{L{=_zv%McFtUB*5xKLFjkmIbd@n?rt5}?PsOQhMBRf;F~ZASvDWa#+Kt7QvLX?4M(8*
z)aGpUu*)pm@iA1hn+(tA1_;#)+LI1ZQLrU*wLo5N1gl+}ShM2QHC`QFv#UW**t2mp
z+0T$6aP~|puU2M0b8a0g)rxrwKI_Wx1}h!duPNPGvlmm@5&IiL^<9>Hd|VLk9xP_o
ze!QPu>%2+$z1LJY^KC8L+QeBn@}(E+KFx$Qf(Oi{;Vigj<1hX4d$7xWUk_Gs+*t@n
ztpaq*TnM)~EF5c?k3UN~
zX;r=i_GZRaXmjnDaP4PiX5*d_q*1d4>}}ipQ2bq}RCNW1Qp-Cq@pT0lZc>@=aa<==
zwja#qRkk3NmbYUaZX`<^x*cY7TDF0ObORf(Cz8b;>jyznKC-jvBuEIo4Er8DWOco3
z^V0*0@$(z!u-W6@!<5cHS;<1hdBZCKY@TjQ=GkW-S;RnBG9%GX>Nfo>`>Ses?(pC`
zbMrhXT#abO?%aI}U%p+@92zqW0;wzimhXtvW$$rK?MxGX2X(f;!9AFINJ9#38wxcK
zUW2t4%fN}f6G5}H73|(00Q+a^SdDWP$&roASv7GCdyrHS-b|drQkJ$B;sUm0&NzMv
zTrSRorNtM*o2$k7vxavh-!Ap|njK-@14ng*B5XeUUOSBQF&R+HyArwOzYkXI-om!M
zwjkCfHw4WRU@c20@V_EP3n3Nz^7IlxyqZ{CczmWCxj#LaMYrn52VT*TBi+mKa?q3&
zTV9-(eSR|YMaR9F+_^My_WsOPogWGdyvhnoOL+0Ke*3`sS7)|p&RgLaTPq!{dzWcL
zlVOSVHeq^#J0u)fof$dx2#n9zN;*50Ltyd2@!TWxIlI!W
zFn8@d1+F_hWxQEC*xjkD)UD%tc6Cw%{_8{qc4YR}U7K9Zd16VN7mgoefu5-l?Ol?r
zUssQ>DsM{Kmn%v9A2s0{8sC9dGY-O`<32*mpox4-=o?mKWD1Y7=|z6U&&Ts_4%-zr
zMDj{>G
z|L5a({cpzq|A2nn{+sdRL068uS>dSL)x+iIt@6)5>A*R!^U&R|k;8v@9wxr6FGRa*
zVQvL4a(4DC*m>Mqc-mnawDai9$_Ay8<3YQ@-p8MBH7N?3`&ts<63th1Hz$eWCm2w5
zpK$MKLGo>FB{Cw?k5$REgY>{6r23a&k`YjeN1C-F6Y5kYZ7jNzeSNj;xmz0B|6w?L
zAeJB%*5`+_X_a6>i|Rc8hAnWS)H`Onq*3Ov2}9W8k(tnAc01CixDU6Qzfd##)h#x8
z$0(5EF0wJL!yqPlJFBs54ym!YI7vD?HIrPj=2dq3u=7<2J2q>IbalKAs`yP|jjz-u
z4=*1g9)os5%+UsXU_l%9&ZGnj$d}F{LXNY*0`@%c=NlF^r6Qj^XM*tK<~->FI3
zHj22Ip41!mJ=d&T(V8b-Ud?8#Y|M+?@!^B=mm+nB`AKEV*W{O{lqNB4F0%XZ+GDQS
zc`3YeHL^N>FRZJOpFJCNfK@*Ze5=JJ>DT)O(86;p^s(;=J8Y}-H0w#Qs%HbC&FkZk
z6xa*89J#`}ejLx-PvmDMUiE>M6K4y<_N-@h&6}|qTb$U^Yi>|}k|UWj=PDz?u6*|h
z9n@Wa9eP}zC)N05%FnefK`#5hfvL+T3l?_=vLBJ=aD1UNv96T>56c7s`#qd>dg;#0
zYSkq>>2zVbsSj(j)14hVs^$CY5Z)>vg{>aFFLP=g4^q9SEuVWJ9F7(*&hJ#;4-Gco
zgMh+IpmJpkX~KR9VvltPP4S{otVeUwxo<6Uq47dUu3DKKu#Sa3Pkl(F>s4WmXwUDb
zzK8eU--9OZvy?Ao8oavp6OIgDB0Xt5Thh3GX2Bb~?20DjTL%hIuot|)ew(c_n*vvM+y>|8waE9hASu1cBvxApfTQ2ba@(}YQn9~Iv-gb(
z@xm?rVBg2F{AE{bJ~*if^Kr4{5rMuuBzi47)#j0=grkmqT|Su9zA*`Uv`>J(`%bXr
z%k~gAb&}MvdK-vey-3*8C>jo}D<)Zg{>|zol_I?>r@^=f{)~rudKa?|f}4Fhkc;gd
zNkZ69wtR9j>s#)GkYV>1d^DZG8+b2;9_x#e9e$Ic;l7z5c=li}+PP%mmLB9i9SQx*
z{-ZQ4{Z4m4Xl1W
zK{#$w7%Ii2!3*0G#H8pAVTyevHgjk*ay{e&n-X7%l&YI7RsC@SR^K)e0)Cl5%ynzx
zu(lf%n|Mm{e>z%fp5VamMy!K?8^w903E`yoB4-{py#b%yYX-k*;RQqT3H;%q5O{F(
zkW{|KG${wywPzuVt*Lst}VdLYNrsX2x?&65?145Cw4l2J~rrceez+)IM#aX
z8}@jk9pAIV2JStw<<2!4vJXeB_`z3$VBX4fNQiajW0y{YJ(I?>Ptj{2VeA`NGwUns
zzr4KU>FdR}HChEh`3jN~Et`{SC6n1C_v>Ktr2~09evfAHjKNTGoh2922J>rpj}tvG
zO6WZOU8cF$F=_vfLVRp|f4FL0lsm_pYwnb}BlyN8kQbzqFl@jyR&~*B>0agw@MxZ&
z>^eVKxRk#bd7g2Bb$hvqt!QG#?Jib>m4jBXXRwL~;R3V*)jwVeiu4d2o4`V+oMe%D516h>!A}I6u
z0t?x2lvUdA2u=1zNPRz*;$gnkNw3TTya5s}oEu{1ciU
z+b@09O~KFa?riFY!92oe1;l=6KxZ(|K5s(6w%?1&hUMryHqW!|8`HiLYP1{&p^E1$O;tuxkxjV6I
zoy3}ZCNlrGJETK<%y`Ij$IJOlhP+h!=9GlZiiiX@qt#b
z*-{iP^q9MWO?
zQ%M--nZ(vfo_y`T6+&IB+N^TD&&)HkBBN8+vs2A?z=E_|WW2|HD6-v_yRNEDoLCA(
zZ5#{!$7=%ijeh#53$vEU$PpBlfb`h2?ktX{zPs^&jH;>~$75zdreLbBZQr
zoQ=@Cy$P&b8UStDcOlanjAX+$?a8!0eFg?3q;RpS1@W9#hqV6d5ED;NlROe0LXSG<
z*oC(l(#69?Sx>XM0`Hm%!7u=(bZ}sEMK{tp+FhvoD}^-n^pf^%S^@*7_k=2Mf>@n2
zfvhmQ%~sx@$FKYMWPQJE1Lszetio$Y_^~Bec-FOz(6~~9&}UH)SUoAos^lNRV`eXw
zc3GJ6h?#3x%|jyW5A@>2+m~hT&f|Eyt%do$b4w)`lk&W}mm3+C))qchUI@J$2zfcC
zFrUA>0&B5%0yFWO0=pk{~pDImt$aMYQ|enixhsAp93jwiKIi}rEu6&VkIl~}H_+0OS{
zz%RyEIyPg!M88$z6tBD5_Tuc*(tfbiKZU`~p_$za^@mj_N3(7)6Bdm*0%L2HQ3ZbK$ukyp?gOh{|0=s@vlup-q{7ny^PyEqD~OEtgaXq`lAlXIK}E+3
zWbMdSWPFz-7@j%6yPa<>p{$)3U%kVU&vJRn8oYQQG`%SDv>y%mk|&O=vyCTT=@qT9
zsAs}&yX5EXf3zVMGs}?t_?f-TEP-zce9U|%TC$7z%kf>ck7!;l=?y)?`@)xuG7vxQ
zE^AhF477<1WWQG@z^5lG;n!^wD00e*N6sn^kJrD3ZT0aQB=)9oc*r&=-Fq}^baoH?
zifGQHWj$Ho(z#v}dN%;OKEugF+rG^2P(yh3dNz#MSVX#$J_XA4^Wm+lhO@4Z>yhns
zLfPInA0TDk4L1C|3yb)87y6CqN1n}J$|gNr4Fe(<3w!!>f$wFb%Bxl&^!w;2{pd+R79c@=>0UYp>!=8bfqX*qr>
zqbl@%Hv}rJ9L!#QT?BpDW{H2<%&s4p#ggxShFfWs;b;BB@ZcKdYrbz`ySwgXLEZ&<
zg4J5EKd?~RShxMG4luP&^3&}yk08%-LHo(Ipp
zuSgRk`Z51+iy{5jT$qy}vTv{bQi^Pmh?Z%Z4#GNwE$Hu#WmJ7cx_oMiJFEdn2*jZo&(eZk}XUbdmm!AzlLL0-c0GP^bSy;Qx$
zQQ=6j2u-z?$?X1xlDv4oQDF7u0)x3un%ORWZq9(|f
zds6dlUuW{}K}Ql&uMx3$9hce9p%01C)nHQWCl*=VmL)7PAt%3D34^1z!YQ2o~6@J9WAV8}cnl81Zv3Pbq&^G8!n59{*xX2oDDR+11Fan(U(Uo%p5T}M`6(OcH5mOXcmpCm2s_egNbd=B?Mbz}GE
zxJv=Ar%RQB%dvB2o{$k<7RIDifFjo0p~bZ)Fk$2hm=g~W`CSqQSQdw;!;8Z5AYCp2
z`8nA~XjsPs+;&;R%pI=Wk3Rrhz5vq$)ktgG4Z`|5qu8vZs&Hz2F&29IX9r`ig;X4i-2dzQnOMNim8+qzQHqW+qs3BGJdfoClI
zbQl{|a3S+;xd40}XG!OS@Y%~_AF?g5mN209E7;ine&!u}H|Fm>8J11N>)2Y?g~M@G
zS+d6)_9d(kUox#Zytnrxz3YWZ{yyK~&g1dI;tw@g@bWa)>3%rrf7XJuOACQg3FBb0
z#t#;}z-#u8B1w#$1eQJKvzJcyG7F6|kp>+}B89vA@?pMKeCDb$(8c~Da{_NB98eErD
znU$O{kOh9JEZG&eKNEp^ta@wImsHtjpDpd4L+yS12VFA
z94y=CDkQ{Tfa4nqkrk)SSjJ|5sN1gx^pCg?m!0RbdTDorYh??=NArrjMNm=ZaJ4R9
zK5UaFU&8?2LsN)27o0DRd)J*L*FFqOYrBva(>>VM>r;e}_lon8K2M>+wBbUWb~Kw)
zbvj&a)s!DB{sFq32@#IOd}LvL3H;UX2s_~?k;zl;2}eVZz_gU{PLesY9OKTtF7NFJ0zZ3?(^e<1erv$?
zt`iygBOZ!Mt)cvhb;6sHF}(7ieVWil_u1-hDVfK9J`k=iI0oO}++yLHX}qD!7MT2X
zJzM`Mhz-&%f(BlDVc`eLYt$?z-Gf^EM`<_mbw?9^tL!W`dPsgU^S}@&e<>L*@2UpX
zuk~V|FMMawgPQOQy^659Ek?4Ro#Mc%RZk&xdMF%rFDp%bkw$u%7bWlG*ReLG9O2=L
zU^ecW6(4Ci19n-J;p;XYlTJ@N2xD|6Jha7P>C;vlvglkDwqkr)VaJUQ#G}JzX1QoT
zOSx{!DSrkhK8=FELJE@J-OfoHX76P#xoGzEowz0C-31
z@%xJk^S6$}!17HYGGkagyU=q8D|7LaFuis`l9+Lxl@7CIi(RH`yut>t-)$GMLl4{W
zzH8SBUQ=hXACCLktgoIV^H6*K=*bn<;n|tYF47XR>*YgsZiWqUYc@yY>TyommVZ67
zUjKj%P&sfcXUme=WtVBp!4B6y`yYS{i>|>dyCN_wb|wohRf)~G_MUy(
zwH*$WG2FV$BVP+y^?6b)n*5d-J?(PQNO*?xP
za=s=!tLUpSty~i_XODxlSF43__2OaU%1Tn;$PB38FkSfGbS!kQr4}krq=ircA
z5;U9T&DU%{F1#P-OB$315Uvl0dws5C=9?NvewFzK-8b$OE{q%^?O1F=t`siK6RTL0
zh^X5_yKo26vcVDdAZShzekr{Qt9WuE)6}e(xvuwkY1zA*P;t;eZjoTldJLN3nY5@oA6PD38ozii
z6CM?lF4z1^DqZw7tm*BJt)hxU}6NWi{2aCI<
zc%V&V{_AUkuy)8;sI%=ZtA3;(D{a;tmZ#lhv&MC2E!V$-M-|4Aq{sPq=*I)#>m-uP
za|`n05!SGD=sMUi>5vfbbxoSNH-^PMFejcqV_2)EwTSbR(!Bb>Vo-I>4oy|__N>R3
ztI(|>m*RR16b{zwAuVk`fG1474;!97f#YTI+Ti3NY31i1aGUPZjCGv^lO9G3I=nwD
zQse>jY*&xp_}NnG(_$K1)T-c4^PT48L_}Fl@ViFDZhP}x0r%5*qe;>5!X=E*<{#On
z{$per+mc`N6Jg0j8$Rb+5V<(nk`-h3CC|deS)FDQyZ_@l
zd;EMb8|ak^*3;uz#j$OmOTMF0?aiw*PfWZ4ReHB1GpoO0uTNN$`$8QkUF0@I@4Uo{
z_t?ZHFIdXXnw4f3yWt#@(h1te9|P}E_GDf2TWsh4DdhBpV4i%#o3C`3%I9tv!=J_-
zhOLS1_|UUc*_%5%*`tnAVc#@D+%F9gbRqT0xAk+u&bm4&{8HerO#Abb^?mvFCU;q(
zwddKy<2Cq=s&Ay-Ar_F_dz!SdSp~NIbq|R18pz^Ob#SQia|mmhkKA6`iWP|6525qx
zGq(o#UJ)6|$F|MSE1xVxOdl8Km(O?P#ZD~1d;9II+`%$@<+c49&j?w`$q(u(k#Cl?E^oA?Vg4;3ag4z3j%pLz}VU)+#_Z>2HTye7FleKb4s
z+7wP!I>}D^29a|2QlV)-TVZolW%Bmk1OUA5DOC9|lsI@p*b-QuKmGBDRmFF6goy`O
zhng1T;PZ4g+HoT*R>W3vc{71s=`@Fp54x?{`q_&2d^ieP-7N~w6VAZ&@fV=fj!Hak
z&qK}YNe)n}P*WHxw1%2HI`MLo3y`%HKVp4Ozdd6L?p-BIMgP0~{~hWoJ&0
z1nbi?;ZSBtqAB!TvQAx3s#xul%9XeZX|FZX&(u*;JIfS!ot~fgRA@#9j(jJrFOiRg
z%uXehzT1$RKaa7?`}acEq2JkrPL{;8j+U2Yx4`x4T4_bCrR>6?Q*65JNS3reL8y5$
z3NqpbvP*@NV2I$%Dm6H-XF959IqoA_YoHOt^lt)CzY)YFyjm7o@PBB9FQKL80dX9u_H@`1FLQ;Tq@5*OZfAK;Ver7M*`fC-l
znR8QkUv(y{I<+ESZQh=5X}Oiz4lK@3q;w=@N}gxlGn{xrz*)9ne*(?Z>m0@BO^^s1zdmI*N4Mvl
zJ`I7UgJ*(Q!8xqYkuu(v?>Y;wA2%lN-(6>|jx%V{#TBl8up;d*bR)5&1Hfb3HRikW
zx6nx12!%VWkOmfO0}rdLllGMxMIJ<#gGI#>`Npp0AYxW0;cPuy7&dB<6y4thEJK%S
zmW0e;M^0DgQ`~^ewyX%P3JriZBagsC_nRQ3tb%qA2f^)ATVcqaI(*onr&7ZWbzor+
zA08d~jdc^>vA)~pvuYM4U@RZNHsu#s!Fq+rf*~&jpKe8^r}q1p)w3FW6P?3OWOy@v
z{WsjbV$aeiPJ+Ns-Ynqlcs6X>P-s2Jm85TL%{9JYVZTS5i
zT;qpn2DV8inu$M{(0mk>n-oQCW`BcsHu;Fh`CIIiWJXA-u9A(g9QNJ+$Oa61D)cW`
zPm1eOoKK1OAq|h=n)Utd?3e}cves?z*;pm1-j(v?MX(L2cBwO~WHu(VPdyX1b-#~v
z+p{ouT5M)ROB5wv8@7;+CEM}emmIj4_X%O^)h$B1gLoW+8j_y9N^6X9u+wOue(_BNg;)J>62xNo$D_K
zO~H
z9IWT%ZC}ES_#Ep3Q5H@4yn{M$n$nN`oKu7>wx1xaI1<5Ul}v{hhi4li7ItuZI-7817>jtBkNjM?T)0~pFC@d;L!-V8;L7t!Z1!tG
zXgp*gL_{X@u^X;yHg6~bgPUw)b4+T&(Fb+hASYFKnpWJ?Ok<9NV=Zgt*oyPpXbQBfJV-%UYEx&znUDunrCF*r;li+1k4t
zItM(2>cuMZd^=jPKF4c8ay4MZ+d1-lY4_OKOSnEVw=kLc`wQEDc|R*Y$(~%VxLsle
zkFr5SOYn>N3JY_pneodR)_gi{+NUGu`Qe)MD3
zDfKw(`^<;T+j|+}QY%2>BV6k*wON>We;&K?#+@~(z8ltyL*cL3Bdq6!nVF++MM^uY
z>{yx3KAE9QwY=r7)llHnK4$vSoV(q&Cl4zefQ4(0vDV!;Nt>htLYF8KRalA=RNx|FBNJHISPe)4Q6BNl;ErV
zX9<(E>B9ThBYCl%JF`NRX4R56q<2qyzN_~t>EY<<L@GMF_jf~>McxpZNuiS&R{2N*^r;3BO#(!2E@)^3)jvZ5_)?C!CdRH
z&}rLe_N+!HsrR&WINB+bS=Sv4Mf^-jpYy3Oz4{hrI?kQfEjX98>^Yp%!BbgA>Dq#4
zuTreiXHV`>VLkh_I)HSaUY?iuoylz0TJRzVT0^HT`vj0OnB(m|nGagHlKXWlLq5L@
zW?5l`;Mv|04y?KWgYg+%Ag-bRI_w6|ihB~BMM?NS)SdY^mQmR6&GS4K5{Z&b87j|x
zZ9=I;k_v?elm=3iG?H1SGDV2WkdTt;xvwoWsWfP$SxQBkY2@7R`2*f_&U?-}>zrSo
zwVt*1bFY2x=i1lz^Zj1e+Fvo68;^rS%2D|GO#<4k#lq9_5UL{Q#?-xcgSkSZj1d1G
zvnY2B-MZ;M`waZ?czrIq4X($}9YS=>)0~Fcs)LWB29?v%BYSu6XZs^c;MjyLPM`I`
z%Jy9lG4d1#=~uX;R>TffWs(UAYBq+X4TkjeFi(6Xh-y1P_S+cpX?G>F%M*x9xg=3M
zoCBMMr}E{@K7s4mxmc!0&_QH1bM=xEDQs9prmIe1qE}rRzNSQRM<5C$^w-P!~Af
zyo0U_B48wYJCS>5g3Egyu_30NHN83+^1k-qMXfzFpz}D+x=e7O)(2-xOk^&E*P)9-
z2m5Ss5Ne)igXm4G@ppwO)(@>G1`E$HK5MqeSE<@{<9R%tk4C$b{30&tY!t2=8!!`L1-+Z
z!YV7hfj9h8DD;?0yHazR-@%)pY^V+=bvx7DMTNv=i#kzvRLAJ@vy8^1HV|g5VOv`k
z4&Q47Z@YK=ySm4j`DL}3=XRSdd)a_5Q@L~k~pURjQd?hNv~JEstzPnKlEx^x>=o&_`V=MKBu{{rfCoyWBcf&k|gfOTjC
z9k=g0wuv9csQ#5iE7=czwj^TK+!Uzlv}Tx|mBdiz9eb72^D`$5qpY+r95H_jVqYjL
zCwCO>RrW*DiDA&6QH~95g78aqh-q^)z=5}Wm;`=3Zvo>?onLNdRs`O~k11c79J5uN
zr>O_M0ta!coH*UJd;?yY=?x!VMlpI%67ac`F5PlX1xkj!$bg;@Ra(~zdR0MSwb%!5
z)`y|NI8hScxF3dsbqN2O4Ar=q1#h<5lAtAVy!wtNGNth)Zu4ouv&RBJx22YSC3y_1
zD}VaVA@>}MGfI=DX<&R)oX!*0O*m>xwWG-u-PU@5ZN$eul%`kSw(>_|l-kFsk$
z%HYh+mtf+b0m-LTsn+5EYU;iT8_!~zv@&sfItw(^KQRR&t3YnyTb`Ch5FILgj_1Fr
z!#wGeSh*;hy>j{@yjy+^8H2kh`L=~MtLTIH`;^_huNBX`+LK^Car(N!k@fRw0%iSs
zY~+^TtX!%w5&t%yB%X|99m)iW4h=n5ske@0BAGejDp80u*OuC16%k4>0yAsWnDI()VIepXX
zC+~**Z%ltvjoOQZsn8V%{`twLV5wgYkjPrdTJ{l|9*?KaM^3Z8lJdB&ZwV6|*@yY@
z#?;)}1DB+&CXM2~Y}w!)s3zxG(b(hQXDLYbj5P5FUgg4JxiuKNMT+Xwt)t(EUN;Bo
zIib>mh0wOF79w)2X~=dn{aXS*LFTff8S({soiiSe}lmK%6_T*LLz*^sev7sD+5jz0`8
zU~1TNSZ|cbIQ$(?4~MD1r6FUY=cxxC_BZ)KT~Akzvu=QZ43
zc?7Gu{`be|r!di1pG|e~!Wo~E>D-iF_EG3rCh%n(=G|4qxlZzUqAC;R6}0JfC3o0<
zb0+;Yp@*$rasn0q9z|JAGi(k}MyGj2%-p;2%;yuI(0ZF7J#S@;TU6^I(V-gkIHuHw
z(+YR@H$ub&71$L}!SL!H;G^HlaLt1AO<1XtlPiwF9`(QMa}7Nb@b5dEr%I%HmIWr>
zCCp5ddypxx64NdyaNN5Ix#)e9-MMlNJLh!)Gl#jt+r6PV-jEFd_G1;dVHh@ImKx(zy{#oq+n{p8W3?J}CDPJ&UxDtMTF
zl2^O%19PEb2n!c@Gp}FPgVj?#2A9^dEjWqyIn0opth)mdM=Z&)+4bnubq`9WCPQhs
z5GjA22f=6R@e=)lbuA}A@U$NY6@7rF?`{yUu7-2V%TOjh5rW?=B&{D0LQ}RYy`Sug
z4A-X-4Y8s(hq7!OFLYw?qGy=r=|p6|&ahG6q=^n4>R1>y8IBE&rxpzYWY*yxe7AQE
z^LXVYT%~j#i#<}|v3Sn{0Xlc9^
zgNG>W)nw?m!%plPD=9p(M4F!K9D#v-1fu=x+2mIh%!X&*utLXxr)Y4Vw=BVcZjG%2
zkvk`teGbPUeb@mun>B&D!Wr!94TGe`dr|MuT^#VG_*GGt_5YlMIvZwC9j+HHHJ@dE
zE?z`yVvoVw*|X_&K?^7^e}Fs_Z-{Svfz{8)p;l=USQUSPkG?gisO^aEQ;eznkzlf%
zzlj<MsU;
zvxUtorZYeHbi#{Z33zxg8Zy6nLJ2PxZiFsJ|6k*Yb-n_b@mP+k6ep6LLr0slX1igp
z(rpa6FbR!bTarZ4G_tj>n9ftyz@9CKFw0SsIiI)#{IYrQQX!k=%^qe%d|!jm@C4!~
zSB5j+Yw=#Mlwb!Hl^74DGpIRvANzg52nbu1!1oGaqPISuthP14J=&&Z!QyxD*F%;R
zG|5u4p~HAUriibvwT`BFP9iz2Rg4`gO<#t|aeeDJ=FuWqs7XDJtA0j+_rNjwSNj8G
zs;;E*-L3G_O`55`8HktQ02Smkn2xOpI28L9K8q}-&7m!@ZBWYXoDv
zY8-6x=D?*K1;`CoR4OC1C*IR?#&
zf>f@tkBvXnhU+FSLfgf0bo+W)lDS=+M73Ely6?&v_PGS8^sL4HZ9+`U*BCUi3#ApS
zx_If^PBgog0rB5kite_0nAPD$mvUJM7Y>Ny4g)oK8!bbE`eSKPX$bz5m*R!))FYb(
zq`>`F3RDbHFz*j-%7Xw0|3x=-oAY|4bmCXG1jLIkd|2!_Nd@sQP#g4wI9pmV1!
zv2F;#2jhRDnBft&`_LPe*Fk`~!}V3X@!4Lu(XK1v~dT@d^&-7yyJL$$8410Y0(?|XR#h}r}AtLd?xT3
zymA;)?7~pj>78ij8_D2`Hh#_4o2YxzmEMxopwE>2u)$x8x%i?4$F-e=nVDf2UmA_6
zb;oI;QX|{^Mv)FAYmhDTM^UJM7Q5LfgSO8(Oi#@*!K)Q%$h=l0w(q&RK;4m+f1ZdU
zHRdGmpd$2$6{A9EAy}F1L#@Aqc+c
zDcAs?)^4$2#!{r|&RZ||31Q0GjRVg?ZL;UC9&x|d1+TdGVd1Ja@tSu6K@+Hd({GOTfRZ&
zrSa5bs~t_Zk|zsBM;Q6p3K-13hP3dd^c+11(dR{p>iDU|(Qi3wg$FZI!GhHF+)UcF
zItxsL-$0?g1~JP{V9L+E#)7$f*%q-O7|60F!oHqh_}30(@;jL$!inHo><9YKR7s#k
z9c%4m2i{+H;^6>q=H{q6Nm9|IbNX(fe8)C6@8l=kklx97zA9$#4|2NYKSS6Om4|Cr
zoQ34^pWv$MJ&NfK&@y=#GH@IktQNwZvLo<%P9A1n_rs5(D(v-L-m^2Oj5ROZQia>1
zUa+?&JAlc~Q~3F`A-(=60YsNmW@isWva5%2;kq5rzEz0Se~U${s(Hk}?+;V_cYtlZ
z;0(Vc|1h(+|3sy~E{syQ22M$uMb|tM28l!i%==ov#Qz8+>)*yPMI6&xa9|Udx38er
zKTadsXG7`l=6}$jSjkLYHk&RnUpTxFv
z8l>xz4P~worZPAR9X4J@_DUb~R(K_7pXp{gpK=_C^gZ~odN(dmO=qW8e}u$NSN5|?
zA*QeWiN#x`>Bp`w?9CUkIR1hl@udm$@V{EzqW=@~IIpDMPGzF9Y7(d2b-=p_Mc<9QclpCZF)Nz4Nt;Sv-Ny!|?1dT6E$A8MDQ2#$
zV&7=#k~!^TSi4h?rp;F;y{tGL`!b*03q~lcJk9xM^oiBEaeTx1mKZX2hZUW44%azv
zLhB9-rlx&9(Y|bo2UOj`-1Z?m?eAo|=2Hi|b=o9a&>acmw9>J3yD&-ikfPh)Mv?9t
z%{amBG_#~+Inf%^#;6}^HYAwS$DLB>E44ECm?cM>L?cOU>RJ*y-IJ_dr%SCW_&}n9
zagzQX)T!pYuVzoMxAQ!`{^THWYnP!5KPQ09-xTOI-2{>O-8{FAsk|9RZ%C_R1Js<|
zkDKbmDE=BpCQTQ^3A3A6e~Uck;H9!v+S|Y(dxG3e?|W
z5;-2FhkZ(QXj+(pVqce0t+&%@N~tK>cXkGR9pE(nfkQa>P!Og&h!Tw%!lYxZCE5Nr
znbBT*0j5-zgHrG>Cby8=-_?}2rQjVq{h0$9cqp;j|sW7Ad_CbmV}S@E0a=)gv)z&
z;uGamniMLGUS~fw`z0@jX`NH)+rA5cLb@cqC<_l1SisM;dTji9nn~WMMC-VFR%>=J
z@z{P2V^aOtSKVAMkffn}{xUdweF5wYaiv3gvbaXY*m!9_^n^N~+l~n&bN@tYAKnK?`!uMB
zu^9gD)21)n-?4S8|Dp7#C{;{o#oy(3v7n`jnUWWX{=5j3|NWDBT|Ub5t`KDoyiGGzoM@)9W9_hP%X2eol=!X@{GsHVR*30pnH8e0E`8+%?se&{^nb*%vEzRJLr&zHb7
zoexrno0(~XA^a)#9x@7-1<)Rb*)5+;I1kBra2tP_ms!0HuT7bYj*9!>;FXJvc47~T
zOBfRs2P@>A+yv=@AJBau7lw0!u*+l>$9TTPzd-^-xZaChvL=USlx?6juMc2-i4+yo
z*Tk5XQ}n)PH+~bZX6|PUqG{kQ6g#pSb4tQ+Tk1DVSBZu0Lq+)gdI0nb&LJDjPvVmd
z9WuXC9c(7Yg3J@{OuU!@H}-0vxs)~L3l!gLQ03W^>BB`2nWn}zTx}SKp1tGfu$dCHT8yI`=>&8=
zx|#l5K9PR62qanVqSS72I2`z5!{|$FX8RS@NWfM#^5J7BjX55Nb|I_aK~f#ah;3yI
z|J{c~ts2HAqm#8R>*eL<>XSL1CUA0&2^{*>0W}*9$-@gVklFqj)O=l_tK9_aR;(h?
zxwQ(0j(DJ$@iZ_O
zAA>h(e3p+-z{kxN1!WB2;%Er{U1>tJzHf$W4=u>I;j^sV+yGe9>jVphv`O5|70f)+
zk5*w4WbvI@=#Xen1n0cujrc^Ou|N+K=5v^Nc)yQbyL<(;vF!xorIV>sz9q``y=M1K
zoeNf9gTZFA0^PJq5;AjHeB~QM%UrE-R*NIilHNzV4}3t6xAz(44F{MlQDs=Ot&pjy
zV!%x6Fxj@Z2YEkA@$O$$>{lJRac$-Nxha0sjN`(%(+&`k(jz%3GFZ7x0*WTqv+~&z^o6xJFM5MG$qUz@OGm`W>!vKI%U??M_a(C%
z40WlV%x6Zfq?Z+)>54r&>Ax>x1#L0|x|z`^D>Q`TT9da^k$T&Um01QgV9A_+1raHyVvcSUk2=Kh;E!TBRQHeG_OBMVD#YrSF!tKYEAkN_}!oF@Wc=wb&u{(zAckn<$^BLO7Yf{0LspP{C
zHFEaeAbQ+g#2QT=u(!>R(Y&qA9%Rgomi{Vsw5^s~p4BL%0;B*81D6I+_{ftjq&iIIxtnNANru0$ahM}AgpJDh87
zMc--SbX%RytpN=+5iKg=0UcAPt{Iwu=GoEgE
z)`|rllFY?(vCukhF&xjJ@B
z8u6U=9QKjf06e>94!!4`$xic1Y^tb6*%%dw{OO6$T590lOmSwTNCrD3xd^l^YH=Vb
zj6QB`MFGW?B=#Ex!9@o2`o?2$TRx`w8UHpc)|m%`dD_Gyx`&k(??Q+19UPlO&};ug
zRJJuE*FQv~pLh?d9iGD7m%_;eTN7S>Zzt9~v>_=;3yE9#4iE}=BLa(U$^M&uL_6~m
z=FYvyc_7^RD?P^3b?c1jQLfKww=IM}K|CI^+bB$yn?(aRy@uwli|Ec%A}%Y;kp5Xk
zoc@E8o8Qh?U!_d%KOYAM-$t?YZ6;AHxQ|u8FYxOlHNbG|GIYA}8?Sa0frEbvMg_-0
zO34rYiq}r;=o)!wY^-9;7cIfCQen&vRHw$?wa{D>2YNG-al7PZ*!)$A%9XgiviRqSeVvk)G
zpzw~R<5fCf%V|rRA1i>lMS&2lmyd-{qA(-Ij68_D#MfTGgqbS83!-kCfmCP}$o1x-
zU7r-|-mFVoB!p;(Ul5$=H0LV=>`e1>&FzpR>Bfm}in3ljGIN7v|zu=NEIGcFU)ZkEfep!+_%@Kef
zg*|Nl{7y&_&0yw*i2XA
zfr&l`__a57vk_d@!tUH%aEI?ex7nHD6n$M%<)MyG;wlz1Jyo*+$KAVqWx?d;L)#3%U2aQ3iX8{a8tA_9fXELUVjN-Pz9MlZiZS87{!3B!lKgsbSbN8>p5o
zW0yE0iY=F-FLa~XJK2Sp&)rM@6p9gp&*9|3%Y~@xuMY2)g^@S=Q}ExCxwunkFPS&U
zL)26tJ&z8f*S2Kv&|%?t-g=^LnFMLs8}X5w0qKdKN^gX6-nVzAuxig^2n@`j;Av+!T(3mk+V8E_!gehi}*vZEG^)=}fk%dm;(FMd5R&
zHEs^*fjb2yz}uHZSI5cOl${yEw`P%W!f-X~w=9ugBP@Ih
zi2pJF_n+}kTmapFnSc7f#sB?h>*oUK{>$tC-+=$Pul~>TfB)I~xd6KVuh-87*Zr5*
z|Np}N;Q-kG^7{Wz@_!u2`oGBkeUmyt#&uR;Nv$X2dB_WD1fJmYNgcRg6@xEly=Uc1
zo`KEver6zUI%_y?fc5xP$$C%y$$U2QVXk-=!$`zoC{Vk{T)Pnj_onu-lW%gRWLIMI
zt%;VTV4Dir`!yMT_EfPS_>$y7UIN_kmWMr?*Pz#MC>=~mf$1M5$+T(009<`S%1&;=c#Vz|+iYOKiRi55#a9;FN1XbmwXRTnrWE${)6n5Kawf`J+8$KdbG77QBXa(F8b
z(*<#@ctUL_gz{wIH`?XQHUbi4{KRp7cPJ7U3Bx_JB^@~!$RFMmTz)w3(HHX0@h_IiiJ1wQk`VMGL+Lh
zUI{{~UmiR*+y@O{Wy}Hxl594AV1k3b6eoQf6m_J
ziH}lV$=06?yQUHHCbx4rItdVeB?MijMnm1lK^*wGAD`b^35|6_uy#%wucbzn-`pER
zBT97fD;;KvEsIHRqa``FZ9nuM5hr`P%OSc~ooIIl!gEVcDpW8FAGD2vVI{|6alS)`
ztS_uafFudI?SMK-DOjC-10%FP;PJ2y5Y%l(u_Vq1IbMNW3%UWbQiqsoaZck)sp1(W
zT*LSj0qXN8g?j1>FlxT}EW70mY6VQ8w~o4y4T_dv+dUg^!bEbUFb%vGrV}GSAZ;`d
z%07rSuV#XnrR%nn!FEBq$Jq-M_Qk+NwcBtCvdEH>eQ5B{k@)?LhTA91n?=9PAmu+T
z>5Z*s)byY$KA63PBrwjrnuuifg}fy_FQ`q#>W)Eex(*tej3-vShhV5YhEoE)Y0uqC
zma+FHYZFf6i|Qk|apoweg_@HE$-C)-Sq<=F-VExNAdXqR@`U+oM{W126ZfRaKomo$
zsmgW!beklUC|`rt*7+#8o)1&6WnsLSDOt<$V;wvHLB-7oUi*pC=Jcj+Fxfa8zF$$N
zrS@(R)tbT&5X-TYEl<|+Aahu*0(c;{e3)r$ptiEK^S|)+MnLqrvry%O(@Ul
zGb8h_2+FTcW?EmAkQD-7;M~z%=+15i@7+(K=$138a?I)aNf+UhA(z)b#fk|TxQ@De
zKeN@-3t;iCsr13(FBlmXh5=V%!C@=|)>Yhx)NRgmwtp(*ixgmYkRiGjU0|L?N8s$;
zKj2epKFmnl3zsyb*gIM0+22)3yrAcjB*FV5?(1@ZVfhibW|K^*v=}RwQHehVl-a`y
z#E~_}JZt-cWbZo#F#*X?yWtZ&UZ+HlOO|2QUIil1Q49HLTd7=A
zHE4B;qS~gPIQNSGNW!ODoKgYG$!eYan
z@V%7tpGgZ)1A&iVXVi)6(Q9!=|2@b}?1CX}5fTX3P~Kb@T$l0L+O8icsQM4-fh$m4
z?TbGz5W4)^9MpVKh<2V2F<$>V)AU>u-3t35cc}t_`%;t|#M2@ND7%teO<@}t$-1
z#YAvffZ3+ucvk8+^M0Bgm)CNGdHkvzlg}E_H)$1=Np)qT{Kt_E-Rqh4rmrx?MvyNe
z_Z%XMKe9n0HC|Ps@ZcE(aJ2$10HdP5I%vlHrWZL0O+aWxy
z!_A%uoJ4;=HAWMK09=3Y2vE*rcz$aXzInNk1YMgzo+gEXOwVRq9oL6vWW8W-S3axN
zq)2;gQ*CrZ(wNipBS31I5GAQ!IBl>3`5ezbR9g;b9koD4c_AHcO9O|M`gBsp26FRr
zC7ylj#eDf90&-_vNa1lID*Apgomi?sJImhU)tQOR6p=@~hozHX$yhn_YQrn^+8II0
zq~iEy4G+Ls!V7T!U%ZySpLxzsL(Pc+Amq`7n-&SO@ydPd>g$a#VY?17w3<%OHEzIX
zoc1}}u86gcO`%Q0d9Wqt4DNA^BbS}csIQn4#qDO;FB*ofbJ8J7KL!U^_u-pr9Z*`x
z${2c
ze1bc#%DafAmyV*-iOW#9cM|QH8i-w1XBc9?iOzi5h7ps^f&Hh3ujQ5TZ0SolI-AAc
zW;$dxS1sQvQ3iQ6X}B4@4Oii9Xs|)
zjTlKQy9ldQ1F`(|0Wz1nM17t2OYf
z3+z3kLuf;U$YJwG5GQ<~P%g{?hzaeBs7+>>o
z2k7tBqYvNv^D=TKVZog^vRzP~=*x`Y;65j+f5-ywE2qM6Wh49GyA0XcQVh~V3gn=$
zEnb%zPvva(V&ZKrXmT>8dtOMRlfE}&bjF9;Ow5G5K8Ab{X+rralzAeejIQ)K6ZWqb
z(%0FuJq=ZG;Jg!gX;BIV@1=->&nF1`vK%H5L{`&)t}1q)&^eoqff-zmW;`?v9t5-Sbx;{N2cIjjjK{Vs@SLjy%7oR4
z{c366&&nTg_~Qr3xCvHiaMjI{1QTsbC
zgn5{QiMI36BI^VsR&Y7dk)mK8PH~0GNmlZHEF=u~VR79jv>Q%gIw~n>O+O0fu1b>=
zuM}y>eq%O%)D8V*^&nujB2E}Pz&Cr>&(0XV1f^>=P`b;M?2x?-XXiA5Pwxtnacu+H
zXlD%j#`;)+MFU`*9YMYasFB{M*+jCnn%y-zpRT-a%zXQ`gv-_Gu~B<*8m#5&(IVg$
z+aIe+d}hkihKz+c&r1PhCyWCni8HwRRy*cyT!SJ<-om<1=H$-0DU7g24xTs>$9BHV
zVU@PWg7*O_v%OOSE5cY}egGF@$c4Nqsr
z(v|5esbxXTWtfWLr0!A24H1=M=%Orz*_F1Y54p
zq(bJ6nA6Iy>U4ft4f|)}7CLX2Dh^#!hdpCXWOTD1t|-$cq4PiEJD&kI@MsDilV5^6
ztPpScZEx<*>|}+x{DN1NqcHPR0=~5{AhN+?c(r&leVwrr?u)qLX6aq9HKGK&jK>qZ
zj5sh?tLIx@9nS>?_MyJrR*26&M6aGd$hQtWhTEp!1>4V`*=FBpvbbUl)gGLH`_d!0
zJ7@v&#?E7_yBe{6rU~=Smcr&32h6#wPk!AQ!i8g7=+5(N@OQ#l>>V-3JprdM|Bn`I
zSg;Y&R+}=%0v@1eq9DDjH;Wc6{=iBVvnZ5Ujq@&B;AbwQ)#_9bSz2a9<097ay2tpi
zu1tnb?a7An3&wQG3~u&{YBG9`%!i=eH`uVL>mmD{Jv|+!%39lw1F4cE8&{331Y}P`
zK3@^;HcC>==Ko2aL`yBPxl`!sD#6FhJ1QmO2lJ?V=)JV$Ft1BAN?XfT|I|!_d
z)B`L%I0)xP+tAlVgqy)p%N8#)$C;fI>2FjfQ++(q{a8BLu>C0dq=%A?{1|+^Bo)>>
zXpwI^524}Jc4SkV7`4(}IQVi4qy29K)7kG&1Co46!^&6CF-e1X=AFXD+Y*_XTZdqt
zg)r@yv}+GIn?(lrV8rN=XYgz
z_ZG*Il9o@nEjEfLn|O;|V=hBww|kMN{u79qWgru#-JCSz}-XjM*T4EPYz>Ue409cIgV-Fx6r+-1Z)R?u_Dd`
ztXx_v9(=QvYKXN#&POTIAhnHpzupbJ)^@nGU`fwcx6&PnlGo>nQ1DL2MJx
zq3VeQtbT9{UhKA`37jsW@b?^(sT+j59A2=7buGB}g*)Ek8?vL~eo&iq3iVXQz|ZS5
zieGiY@mqz-l7Dllt-%B`nJ-Vr`pQuGxfz-C=Pld&u7=&m{Z^ZQT@ZAFn=Mei7Rpa<
zBu2kOP~3GKPP24GU7x?uI+TjQZ81>Dv3k%_EmfXzZlsCOxgn_sn^zsletNNw$B
zM~|$7?uctl&w4E~`?(RaTb469!P3Nbdkxk%_|j~4CqG^pA>@8LGp$~el@QtiwTYW)
z)`_Dq^FtN8X
zCNa0)U__S!{#Ceuo73co=~g#-bGkTv7RA*c<9jh;K@O`MB1!$;G-BXG6$oC*`6#%)PIn=WAN@$att)+-nqydv0mE@0)|x8tSjQ}BbK
zK2f!6fv~+{kSz5BcFO;P-NFn3BN3ETdVvtFhGtiLVe3C754+xp#+-0BzwmD`g^R^bT8
z?zMx+D?fU9+Fjh@Ii5s^bNO5MP6J~RP1RdnhzZ>R-^5m7!-f^qQg189NQIN|9e?>T
z2C49H%>&k7xfC@Vk3v;e7Ky7dq<-!v_$_N!QR78dkox@soq+{J@U}AkGuui3u#NDx
zViS{W@(F^!X^;i%FLv*{E=J^*3)#5$7p8xyLg5hsXc3U5GG&&qkoN{=WgdeQ#~Cj3
zClN(8)tS@GGW2yE$XcB54b!qeN8+8s6tJ-QO)H
zeiq}2!5lucak^Hk-8fc#WIG
z7S;j{hdnUueGj;L7J;8o5Na&lNoQoP1eIT!)W7E@D8rhai;l5F=xs{=x0e!!YIMeEc*s3J+?^@kC;-z@j69q-2Q+)lDiVMwedccIDQSoU0PHt6oJ;qptL;*1bQs@SRo+pM)9Z0>tpak(A8aecyy
ziBWPr>?d84
zJ$nP3;(RTR2~s!$!I;;YhW`XUvvYUnB8KLoYj-O7;ado6pIeYCIw$c=v=((mBJnxQXmF|^-
z8{eH&uA2|4M(=Tav;l08H6vvwUU0rHE=whMJgHQ9gx|S-if@Ywh=g~-yosmbZ}nrS
zd(g(Liq0nQXJwMv922@cuaJbi=CKvaEm6nl4BqLPPRqS)x%z$!
zqcLY4zL!)dzVU+Og?cbZ2XT3@=nj{)PeQKbFZ|#f3t603WTapjxCuB@!-TKw87{lq
zgU%uknm3ahjS^VmYfM@`I59r+EU1OM0%>hOjsx9O>F{4}PsiA!`lCLS=Zi3F_IXi%
zFD>#D_fj>%I&kMD!pHR{kXGrA4Q>xn(qNeRcekGNzBC|z=Mei$M}tnAU1Cy3gjjdM8^fBDMT?$Lzcy%We3pZ73$xg7RAGl2E+t-L=Y
z1MKaS2O;iwIKGy9imRp-;M_Ai*t>e%Vv*-6vraO5@Sg5zzdzITj~DkLC@
zI_xSrd&J=|_R|+F5|!=(fm&-w8Z(1fpBV@22Vr_5$dTE6)CD?>rK$Ms6p&ti28I;M
zplMDXudVbo#5w;#RpUeq&AEqC868l$-vTrOw}D{dczSw7jFIlvWMo`yn1JVjq{T6m
zw;HP9${}%j@0u%J!@XOk)>rV*xl|%D72HkKIZH^1+(E7j*aI5&m
z3@TgOES8*3Qg~bzvVI{MezOs~+=f`)wf1nyT7i_#RpZC4cji1z8BpRn88fUT@Oh~o
z%3IHXi8}W1vG^7fFXqkZ#?KHtS7ZIO6dD_o#(sOF2V?g>FfJ1gL4}h%-L`riJ8NJM
zeR}dDoQrJ28wz|jq*{p9t1fM}#QQKBFA2Y{iPPq{SMgcn4ywf#!{xm)xVc3P-mOjG
zGKP$C;OhfstI=eXJ`Lph9WioqFbm5qr*VTuKA`T)gD|aDkA4ih0?T_V!Ovp_4Lals
zb3Sl)-p`Zq*^#S=n~q?f0mtVR9zlgQg2eRNd+1+!2sRaS{l#U3n*GoxP3qD#Mj#06
zN-FrdVak}&z6ib5=aU&NhM=Ah2;S6;?P`f4Q{sEj-D(juPXB{VQ8v^eaw6zj&!p<_
zGoXHrHZ?1of&xig&qSA-S@|puMRiRfSY$UHynF`)b_>!KQ93kko)NpTxqwEQZD%$a
zyuq8VtMG1bDY&b~60ZpEe{SqB+&#qIV^&=RS=S%TLF2i!$ip7NdnKK_A(U2+#elcf
zI=K5llrnSYpxRt>I1wsGt9J6~LEWz~?uabAqMe7EkIp6Y@&vHLxd$J(tbiTj2b)X{
z??PC#HXRjv##F4zg$K9Fu~u1=D$l)!j%uy&=#K=sT-VEvIsb(({mP>GFN{G#C@;9XC|WHCkBmqSR}$Q*LGb~%bG8#5ng`O|>|
zHT?No@-XICG;Z}Cg&;R+s+xWs8%M%frfC4RXK~M3@&09jc64T*g(I@`?S-I%Kx5ZR8kI|pMjodHhQ?aa@u#aa<8DAChLpcF<
zckYIrX@0QEazCA}^$4fFx1!$yW)W@&!x1kh%y^fBS00)YkDgc}b=r|siA@8;R|iPl
z0*=pA@TboXHlymyo3QRqJu|kLAs2QmqvNXuXkNjx|BUT6bqHw6!$+bV%io5flP)t{qiS>U28s@zLopIKA$JeF1z*}ypB`W
zCfXIBTTDgwzz=+hNCA3a)-$M^FG#{I@_Dz`MZxSzF2p$@n0-2y2_8D)4nd0Dd*LL940Di1c_7x~6UdV^qC|
zMnz`9ixuZsEfm5ZzXR#{u(ecr!Ew+Wola6`8PPn!aOUHFVLEZzc4n>o4#;XW=bycu
z$F6uEMPlxSvn~tI(M!Y0_^aj^`Lf1{yv~i|cs*rkI`2WkSN+0wF;mDKAzvJDZGsxt
zD3YgK3R7;IkXMflZ}UW6fl)5-N0zV
zdU)ueKyx--;PP_{WN?=inWDFh1`O9j#f*0JiD|&dd>+mz^
z5%xb-C!5dOV)VL|L`3HSzH^Btg9U4tfrGNN>}Up_?EJtStkx#i9ymg6_60Q9JBgk2
z*p_DI>yu|n`9#ESJldWAiN4mAApX#qb52m4vVyzTifge_m!yd9hi{CGa~{U7}eEFpup5xtb^#F5t3tX9al8gSPbA{@bYF!H||1Ga~q74NNb1$krx|
z^4$9?$bHWLCQ@^o{p+hoZuuyX3iG+N<&7OX>b0C0JC||3v~2ukTnHui&!flE+c@ED
zGtOkkQ;oz9?8xHscXSFpZg&cf4<9DJXEXV(gIO3*lt3IMmy@;Je4{oggYWk>9vc?)
z;?aqj_+@n^ZhdtU=U8UolDNav)9?a6{_-`ZPj)ufEB1!9EUpFLSP`23Q=1t6PUXK)
zb|N+pjEH?5WlX1O64uOs7}nL`lgkctsL+HOpG4f67fy>zH{wbgeKNL!%a8~v#
z?6cmpkb9&E_If^pl-G-SQG3)_wT9F9PFkB5jTPXaOcW}b>hkVXi@>HiW&C=Rtt943
z0CCwlje4x8XUL^I=o(rDmprYR!kw!0LZ2A?9yba8G=xE0rY0>aH-)H=>$Nc0fOtCmgf{%AH!!eoQ;wMHM+p{r2ajs>V~FfJF!GK>GEEnM?~G-)UP!9$Nb2xi(0hBH
z^p1C;VzX4pK_ekL79mOhi$7bPm^tM?@55>Y{rOlXQd^H
z{rr>fR&l%kD`4`<6n3c^a`%hegl)1DUiPv>!RR$a)t4}@S1M4yyk0n84D^8338*c;
zjdluKal=6_J9xw%6BI1L`~7N~u%{m^y2J6_srhiK!I1`SO2zK+qV!el7Brieh+#=W
z^ni>VcBIaPrs|EP&d?0)Y7B`)%4~chPho5CUShwVvTH03kZ7JAm-o=7ZY!qZf|$Md
z&U*#k6jy+bfh_)sW^=MQWDQXc6`;{&#>~A?H6qiS2TSlKG~L<)65j1lkd;V8Bp!nL
z%2@EbdW>o6p3n8Dd6UN{9^=~e+?k~H3y;l^ppJFIR9Rq$hxRr{N^(FhLnQ
zxAyW5AKFD<|6*W5`dJ)a@&WoUbH0R1Zl_U2an_C_?AQ1k_={%=yi+dhk)v;MWqv4|
z*RmX1*pAgVL!ZR#L=!kelE8!etF*v2ZQ&(TIE2dsk8{UY+$Hw@eX(y#&wqPk6Q8g*HG}^
zG*Wr9{N~5Jal~Z%OL%nHlza?aLVj-F#z9)HY~P-7bnE#9`twC6CVzW~ZEqU5uS1x*
zI|l{+UPSRh&KG!gI_I0d0h+7y(Qs`U6ozX+Kw;C1buQ&ApxvOw5?%${J|i@R!xQWloH)lwgk9R#+#S
z3d7E(zJn_tF5}!%^~~b+EnxA`2lvfXq9Nik
zTyLW+3bGDxF=-)f3a}*BKa9zD0crSWsKxx#Gk_A;T6lX|6wJRj!bH<>JW#1buh%CE!RJkL8*_1dNtncKaz&9yYqeqr*lh}evx&Ls9
zi6fO6>}AV+EYR5ICG7uTNqE0*f}Ynk{t;hw`o>=fe7e8mjZ4zRYlAn8<=HxE0@Qt^NPfK$nX9l4$M$Z+r*%S%c6J88=!pzte$kgK
z=yj!^+6HkpUxrAToaD_?-c4^!;qE!l%9G*4_ZWN5G5>q#P4<-4L}Ifm2mAh|fW_uT
zR4zsV7Wa>*Irp#Q8o?M)*(gmGU);bJcFE&Z9~FA-Z!w^?`J
zBRKVQE$({rjkoV<4ahmIrBVO%cs7+kpgwvltX%7gW!JQ^{G=BXB54fKi(MF9>6bX;
zt|qcILL~q73~)$N0qt4Wa6^w2R?ZzyBsd@FSvxKA>996gy?PQxRqByh=`LJfjvalb
zvVazZ`k;&7Y*@Va0Ql*dz>UBEV8Js14jh?D$kJ;tCu}+OE;>x+37@c<+Q)T0a-H2r
zGvceZ&;~rC`3UxPJV)*AX^j2+8?0tq5q?(iWR(wuqU%&W`tRdh2$bNumbLHmr&cS`
zPU8!Ni^7oz(;cV{X>fAv2k)h2HY+~%8<%n!yh%m&Oo~T6VbC9ai%nr10d~rw|uDl>f56yRFg8I5a-DDy?G%u2AKNLy5(td+RT{b$-
zE{DAi&G&_(i6aK-W-oWAKun&M1S6GxW3nvl1=5{*RYe^xM)BW7HpvRTTO_Vdoo^@-3dm^;&6eAI9>SSCr&><
znU3C(r@{IvP`mg9bd-H!?0qI+=HZjrylXwiig&_Gho9JyAcwj^gJ{3O8fN5|^0p=J
zpsU_AVWU+q4SKzbKOjAsTpg)lpP!h_Rx`%9Rq!=>pQG&CW*KVs^D56FAQE0FaGa`e
zJFYV=z{93zSQVScfLkx(<{o<#Hy%Nu98dHV5g^w=9^bEbrkU4l$?~Ek*4}O%F@F%u
z3-z;rhs$H=QSC%JX6#09&HIVpCa%J9kK$o!!fduXf47y^Un@ZFtlrQTgFhpT@LH)I
zd7++-w-oissYpZSnsE@7SaBbeY}$cxnaZcpPjKhmckI?{u}tfs$9QOR7L}YS!^@k_
z^|tkKd8rL6$-71kEZB1lN2?WR(!nk4!-{(JzM=;+f6k;~Vk|85X@IuTQPxd$J~;05
z#J1&!=)cK7As~unOq^3Oq8;jf^z=85RoD208kIWiHiZ_pmk1VAw&PuO
zA#$`ln?C-$kJ!9wV2&&l#;8jJaB6S_+>4vp^L25!(Jc=mxeoYj6|Tc8k@P
zRo28zFMrF3Y(2_uEm8!#1}9K%O{W+X%qnUf29v33#HPf7bdL2x_|(_9qB$HKz>+MH
z5hYRjmULE=Fm?JC4!Ym(L$uB)s#m7slxv&GjrpeJ?+r;5_I?LSqgP>*`Z4A}ZVE=K
zuc8hy)mSJH$EJpFYbL`qPFPI~DR^a`2Au!oEA9^<{L)A_vR9dnNB!(NIKT!)af>o)NL^if%1;UPZ
z(_m}5C|tf|%V^r(2mhBv7+5iiQk%`Fa_|{?@kD`@H!L%d^{4GkB}P#`vrvAe$k1f^w3={h-bMdTLtXg+0@&Abb`nh}hP-AeK`
zxemvqBuVLr2z>)q*mFCY{<4$7-o(YQgzM-nYM?memJj!%2XIr!Z1Ud95mf~&=_lV3
zI2Rg)SO4Aw{}Ux}HGL3xT{6^yb11kMa&EfwvD7^wnh6m)i`z%Xk+m@gnUN-Oa7&P-
z-us2oQvWMvq;oxvSN_7=wKLeLAPMr-k?UpAUX2HOP02p*V02DD%H}C+(924Ru+B=D
znD5*{_wBN#_Q%9w(<)`+F@bX;8W(}lQ_f%U{S!L%+Y`IWW;iPLgx@7v0P9xOupg58
zz^`>~zLcNC!<%nn_%CE4qC!C=cr^|P-$jwQcJ>K(e=^BZAGCeaabfUeI{QZ*
z^L)W%qM){eiL5vdM`}AD+Ea}NXG=rrw;kZjcW1uUoVd;eHOz*pXbb>xe?fBVSrObD=@=GpUMiqhmh1)
z=%A}gE3;OCan3t#r>}rUcmA;IH@ab1R~x=1^T?GUGprW)U={j)0zYT#L7LXw!dfN?
z6QQ?Z^!}l7V9=WeIXerm%kMSQJ-z^&&nUBIo62BAg$(?7Q-Uw$-7zlpJUsd6PaEF2
zqM*}fM)|-j5)@@mFK@pM2ODHKFU?FW(YS;%%w#-k$i1d9hq3Wr1=@y5!!pe|c=Ps2
zP#ao-Hv<2m$KpojQFIYRsBUMEXWYi&+wzRa>=zi4DnvCEGVnn2dpyRhBtgm(NT^x}
z4e?OJUhSLk@^v{+aibVm^r}*y@u$$}?I;sxa{)r$Nst@X5lo$-AZTpKfeW6em>%VB
zsM@{;CwVFm^@qmPZ9+1u?C}tEQog{Ob1&J1NviN*`YuLlU@Fmkuoo0(eu7OW9iVEu
zF$v`|_m?zZvod?WGs}i8Nq<-_D*hagH?o4j>G=a3-LeK=4yB-O%T)3=-ws1|`GEH9
z6w;pJ&p+635(MxxZ!qK|BRqZu?7EywUd*`&J!-;KqWlGJ&0N8_eA@-5X4dnU4SvVW
z>`00Rs{rh)@u`U>k&f`CD!RfPv$hr{PcUI_I;&E>uZggl%>uWSJGjj80NyDTgg&QD
zFmyE<7u~6Vx%Ur~A(2zKX|odhNHB~jRd0Zt{VFuY@fax2oDA!JTT&YyL(^;>abrOn
z{Ag966AHPmksWKWZO39d{N*NQt`Y#>$&0AO`XYAHCpYT2>mt^iTn-nk=J9iHnNky#
z$MEM|3v}|s=!}DII6hdMhObb;eLtqL%iiz9hjCxA*|GziC1=yMb5!YaZFM~UYdpj}
zF{jSn&oEfj19JA1LHsW{@XU2wmADy<0FhfeZ8_`;i~
zgWk|v_}P90n{0AeeJ6eX!gy<*L$?JgjuoPX<3h6YdnEMO7-H0Pce2tt3+%n$GY4z`
zK(E|t=4`n!;36q1kfFvgG}&T5;KXF|2p3hAUHDsP)k({LQRnA2b|d
z`HPOiLxFRQ(b#NyR7)9zrsy*##1!ZrElcWDkc3ihS6IdN$B?Zki=GQ^V)8v6n9d4h
z-f%kt`{|SD+|m{Z=;Zh*@{bcC*aSflIihGJw47?}qYqPjT^}5Pf9Rj_RD7^Ie@IRo-}%>?|@NYvs8P-A5Cd
zoX?^dtf)?UqvVLnYghE_GRG;oa(H-O7t=3%o1J~>ER;W60&R1z^6#~bqXoyMadp!V
zc0B(hPQ85(`VS1C)B4ls?f(FB-D6?t>m2&xhX&m_mJH7uju7`{qWD+;9Y!?i@w!92
zXhD}Tqhy~*EaLW~mSH@^<%psAmIr)hdnA0jq)MzhtZ-J4Jl+570=w{RGPsi;O!3{{
zJpPAEOz^By=7_rw+#y3u@J>UrrjlbEzP(_iL##+5clP#am`86V{brB)d67kb%b6z6
zL%?&)U<4!1A=i}4Opi6fHC~tK!o~-TUX?sq=u*i3E58r-@3*kwL$b`y*GbS~y@boy
zO<`THaNJj}D}GBDfL)v?MBSzeg%6sNZysB)c)A1GWGF+PZWhOS<>l~R=rA3yokG{m
z>tK$aRYTn`%h>rQ3u(!S1Lkd=4URiz&{tYlK|6Q}_WR`Hayvh6r)^DqW(bfP!vZ*Q
zyOdGT@xcBRZm*zI1)nS8V9B{j5SDTj)~qnY`h_*@IW#I^>y;LJB
z1m=M7Xb+pIDnNgK{Dhw}-9g)T2!i!|@YL{okdeHHo3||R
zE8yPY4p7sq$4VCw(wSgIJs1)CEoTW4y6=zXu3tDNYc?;vG8&Ea6_`#{Evl92L?#vr
zky$=(u!rvk8ku*Pt12B(x)k8FBgqC?(4jSvP@bM+yo*ipFB|G@i<<(~Roeco=M5
zL!W(p3>~}0A(UJCS0{65ZciwB$Ckp(g+rishU=UN{Ejt(EhK2bpXa*FlP2AsNpCyt
zCKVN5`3q65-69newWi4&94nB%<}_{2GrO$jn!Z4IL_#XW;{+?I!vz1@kYT{Vk0IZwK(8RQ?m
z!(LC@Ms!o7;P%i4xITRx`P*7uwf4sVZqwcb-r0x2_?RGhbnq#vzwKaOPckGyzf3Wd
zy@{4xzu|kTJe@J9Pv^5mEOgz*r%`KZub&uMIW~vh75<0@{Q`6}Rg0OQF@>0@Y^7mJ
zZ!rD768*XQ9(!r3EMBlT!C8MS$+s=}sB!Ev_g<;M^P}?>Yu%>;PU*F{ZPI
zXH$XP^~CY#TbS-pgI_AuVAJnDytMT;Ote|TNFH8F5A7OZYp=e=HK|8HXnFg>S@VLxW#p#n8psFjG%p2hJM
z`x8)v%Yj|aTR^%e-NnhTN5IANH!EB59{B!m*hMPM_`t`QhDW;*+I
zviLOT#+gGL{+{9H4UYKi=MbBHNSg>9Wj&)0ZH&`z6K&HY2mqjh6UwT;#JJW^LpY8vkk;7N+`wS
zs_fF&Mf_dQcUN@{O{Px=Zm?@7$+2@-d0s{6Cs=OeNwm#U@ooHdeD}Hq57e&Y=J!Xq
za@Rn9-1=1f(0w0~T*1%didC|+Uu~?66doB;_-`Qhf
zuOP8+@P$R^uH%8B7PK6xWFH&MCu1I_w8dcuXol3HTtpl5WUVB8ag!n{V)9h$qAM*p
zU;+nkHgo>C73}`bgly@1#hhHU$L|X!QJ#j2gvm2Se~_Ce)$75oXVB
z1D^?#iA9h!{8HLNiv*jXjJuZF?tBR)H@85l>m`0&Y7?Xn36sJ*>se3VCYbP49DW$Q
zj)H7zO_vG
zoB`{ZKoFTdh2z>x;KKGiP{>TDPfe}K50oIqC4QtS={(#uAawJpa?V4M0Tl&ar1V1t
zYJ5{57%9*CM2k>qzui!MLWI23C}-|@R5LHR8UCUd3&`Cav#YV894*1+xL%9?VTlnPU)
zdj>bY4&&Xjk-@|hvq^01I=Xh)3)Xt6qqdzOPMVw#ITNHoLFYRzuQX*96-IeJ!h5m3
zC7LYow*%RxTAZ|T15UE%GkPZ;L9cHRZs`=H%-n2DJ)MO6S4J>FG@e;^e1x6k6M@=Z
z7jg9bD>!mZjvhz}B#mF4asCyq?`Y9IP%D#yvyrRjfv!qU;b@K5bN{JCaI{p~Nq{>F=Z
zG1F_%-y6&;zjG2C-*huCB}U-9;|F$GUn0HRx&i|t2bjJi@z}g+3oCD%jdgw8In-(#
z2|4}_yq>E^v23Du0xzo_=us$WB~w1G}GH
zVo&}}q@kQw?x6QAeC4r>w24a)S*t4;>e&b9KgH7e4aq3JREYoN#R>W^#ge4n5h7uq
zO>xe^ZRBm*4~{c8k(VXjXj^?6zihQ4%bpaom-o1_X|93r&y33+>{Z1{S1nNXz5o^U
z=iw&f*|hlBMRuFJ0Fg0UL5tJQLU^_gU0I_@$90@x=lSH~kg*ml?URN}uelwIfS@Il
z$Z@B&iC`Xn1El-sV_Ucxys%l1?hBryqjwL}C;o|*{U{EP2la^hyA`Zrx&!N4F33z1
z>jdrG1gx+cPc9sYV&9B)qNDO>eEP8whj$p@*F#$LZsI&FF}e+z{#&t1b`3HDt#HuF
z6XS}atQ7AYM#C5lD&&3{KFoU$qNh*6$DAt2>DGm)!7W@5jRD}ArqqUs{%)=Su(&5`!JhZg*W7>CTxcuY>W9jo0X8DSdi(#JN;T^#DeDVZW
zM3|ykaSmJkYz*Ec&85Ghb;;-qRdVBeD-}xC0^u1e$hbHMX2ch$(WN@Pe0VQPTh!u8
ze_;}b7IfyC1T=mtOt0qdW+pfv0`04xSugjEbXR}`eLQAI0}boR%%VR0e)9|FF7gE3
zf4A`PnK`(*tQL@W9NJqfs*VT&Ox
zlKcY=Tt?VV4*ecPbN*2=vNU-j^Y2X#ER|KFxp769**TX!+-J($XPbhBvlel)@o!M{
z+nkPtuB7UZtf@q9C!--$$#fXpf?0dckXj90T6H}ES5&DJC0$kc^KStpa-5rdMkITC
zn=LZzC_J#^dcu>-z%Ngksy4}!DIVNgU1-n-^;9D}zm~0b@B~
zA-I$T;Of|iFyI``Eb=k|waa|mZGD4%ojr|uPJW6pf1hFMpF%V<<}I91)4
zNW1I=iSuN6u>SOjq4ElJ(uSEt_18j1_@@Mj9(7=q_UMt00SPj1{Xs@8p$*rimN9Fu
zouJE3*W%skan!-=9L&5E!O~kXaEBJ*{dwbX{!UZ6czP(*yb>bDiM3Xd?Ly>b`f^y~
zdLHzh9U`%RXEzekb(pVpG>0`2
zSHxopnbB~iV16`0Ewv-mibE)?9$4r>-MO1Z~SQd6IL>j}ZdkO%N?
zS}DY9C!^%7Rya3s93G6Pn8{_;12oX>8`Q`i!;|c!^9T#%4RL`-5?%6fHo4BP!9M8(D2jQ?
zymb=5(jfslbhirwhA%K*O?q*#Y9<>#Clgr%anfSkz<5okb=LR&alYW9>*@^;Q3RwOx~1va6auQ%3e*!^o4uK^7bxt(2yVr
zFPG996;h?r+KE>G6lvSRMMSg5jcj?ei{!#X&QpI11eUyIz5T7Ro!ce$=T1ZO(e;>;
zqJwEg*FodP9=zE!kI{Yp6z}hnqwn@?A`_HHSx-iT8ZI`Wv$}H7&?gG}E%fL*_e!Xn
z{{+YOiBnO7F!Z>e#@=X(MO~ge8Hc}6Zq{)UI9G@Dt@{KS0Z9;ZVI@1!X)U}_JVUO!
z$C5n*_T>AgeUPriJy(xsL30lSI}feKk==2y*x(Mc&D4aqoC(3L6Fo3H_5^G>6^@O^
zelRCpW`SDeN>1rp4N$rf%DEYu;|Y0M@be6Oj;Vz_JwT0A36!$7rq5!&{onq-N!N;w
z{YU(t`G4Y{``>~8`%nCH|6OwK|K+|cHYM${;lBN^&;RfK<^Lc5pZS08pW}Y?f9vCa
zhX3Qdb^k~6fBy;p`u{xb|Mfin$Nj%EIga&zk^i%PI+wRiU7Lytrh%}XCWN@Tl2>8E
zbjDWBJ0od9yBf9NvV}D*7*r)IdQD+0ScC*j8Gs(SKvsLh8sf#xBc5xmLeK9E9X*!=
z?^#u%dn1w_kGsknZ&M5G@^j$5B^@u>m!XHrO6F9FI=8>!`gi{-Va2v8@Ci^L#}D1b
z^cfFf@r%W1z490|t$oF;j~_?GR7PN5PCS0{uj0pF$m
z4JfZ#O<#}RL%IGcOr252{xh1wC@LBgg+PC%Yw;!M-+LG@s0*-5XU4#qmpLeuFGF%x
zParxLP5jN_GNAfH2i!TI{#&EH_-D@;INm$W7jxPR3+@||gfsf|nTsfiWX4b-c_FP|
zDGxsT%;{&pI(%PXgiGe0z*$S?;AfrdaNz7zHa2M{omAyc)6YD?FKznt@yqF$+}i@j
zmfMi$)+T7C@4z2QFXm5mdW7albD?{HAYC28M(ix(`fwvKW8G2avYrJg4O&DSGYcVD
zXbFUtor41$M_;f11A{f(*~4zPL1nEMev32(ok{DdPdOpkQ?l5Z*N#D{FScE57-Sn74jy#x7?`L0P*Dn8t`-?JA
zyiA+qjT54)+eK*GH^Lfd*DwNoiD;ml1(kuL7$qnPJN;JUN@~tb`)baZYV*kBH<6q}
zrw-H3Sx`SM!_%Z_z3MjrxTO0OB_m_C8HHJ7n1
z`8fNx`9AN3=~XmaF$^PtM&z7c5#+Q?C+4?)z~D6*oB-dT;H^K6?bj#Eel@~=xevTt
z#|5aMt%?f=lkob)d2Duu8h_E5EpVXh0uH+Pg2=vC5dA)us!puuH%{#1y|LFJ3-ejB
zKed9*|2LW3TV>7bU#m&PpLt-!KF)u3M2ehN(WeJ$*3h@J1ZZflEEFht;_Egk(zWjp
zW^bs$KhI?7xmRV>iaVb@(2d3Qy}uaAeNDXZJV(>RSs0{W#O`w(`ea*Z){myE5|Rr
zhdNpnLxgXXb_R|f1ERZ5SME?lje3A>J(tl3MY0z#L!Ph
z%-;mAWgLbnhoqsVxPn(!WKF&wxeVC~2~7Fy3FJpZE_=5lfasKM$6eVP*z{sAJb2>B
z#5Sa};qRX_hZ|Qw?<^yX+9X6BE8bSAY|3WEj(RaVA!79Lx(jgbYbwmGRff)%J@^sk&RmIiA8Rw843+VBnIrsA?uX>wS8y+B1$VzKPRd(7
z$ef2~a71Ya&1!8Wn_myH{-eU=^kGR7YobMezHwourNv=^wJgW&sB=BXSK+ATYw#2j
zB1|>c3sJWQ7rV#epG^U*)P!@K%RY;27ZWDGR#?O8@Re-GB^|1%7Rx+}Y=xTKGiZBf
zE+&o-=Gek^yxE_OBF|*t^;SN7o0mf>yQJucx;&;XLxw7lQ1Wv}3vd0B>FA+87wb|^
z;(@?3(DHdEnWJhz<)S{K_{DSZtE(NJ6{oT{PnwVw_obkko6YqMq+wTUCiZw*avE4!>cRN`X9FeBdO~asnz?LqVetpkb3#*=p8b*mQKIwVU-uMpE_(z{Z+U=EcGkl0yR(U*tQHych-7~I
zsKOr04m|yQACzu4AZiPo;F-Q4IXK$@O>IS?