From 0a66c1131141d7c67d391f5d4180e436c7236bf8 Mon Sep 17 00:00:00 2001 From: Utkarsh Date: Wed, 22 Jul 2026 16:25:35 +0530 Subject: [PATCH] feat(inference): prefer ONNX Runtime with PyTorch fallback _load_model now serves an exported .onnx model when one is present and loadable, falling back to the PyTorch .pth checkpoint and finally to untrained demo weights. Each request logs and returns its backend and forward-pass latency (backend, inference_ms), and benchmark_inference() times both backends side by side. export_model.py gains --all to export every enabled analysis type to models/unet_.onnx using each type's num_classes from config.yaml, and validate_onnx now asserts the (N, n_classes, H, W) output signature. --- scripts/export_model.py | 177 ++++++++++++++++++------ src/climatevision/inference/pipeline.py | 106 ++++++++++++-- tests/test_onnx_inference.py | 170 +++++++++++++++++++++++ 3 files changed, 400 insertions(+), 53 deletions(-) create mode 100644 tests/test_onnx_inference.py diff --git a/scripts/export_model.py b/scripts/export_model.py index e855a12..9010e69 100644 --- a/scripts/export_model.py +++ b/scripts/export_model.py @@ -1,15 +1,19 @@ """ Export the trained ClimateVision model to ONNX for production serving. -Produces: - - /model.onnx — standard ONNX graph - - /model_quantized.onnx — INT8 quantized (CPU-optimised) - - /export_info.json — metadata (opset, input shape, benchmark) +Produces (per exported model ````): + - /.onnx — standard ONNX graph + - /_quantized.onnx — INT8 quantized (CPU-optimised) + - /_export_info.json — metadata (opset, input shape, benchmark) Usage: python scripts/export_model.py \\ --checkpoint models/20240101_120000/best_model.pth + # Export every enabled analysis type to models/unet_.onnx, + # reading each type's checkpoint and num_classes from config.yaml: + python scripts/export_model.py --all + # Override output path and input size: python scripts/export_model.py \\ --checkpoint models/best_model.pth \\ @@ -46,7 +50,7 @@ # Model loading # --------------------------------------------------------------------------- -def load_model(ckpt_path: Path) -> tuple[nn.Module, dict]: +def load_model(ckpt_path: Path, n_classes: int = 2) -> tuple[nn.Module, dict]: from climatevision.models.unet import get_model ckpt = torch.load(ckpt_path, map_location="cpu") @@ -62,14 +66,15 @@ def load_model(ckpt_path: Path) -> tuple[nn.Module, dict]: in_ch = val.shape[1] break - model = get_model(arch, n_channels=in_ch, n_classes=2) + model = get_model(arch, n_channels=in_ch, n_classes=n_classes) model.load_state_dict(state, strict=False) model.eval() logger.info( - "Loaded %s (in_channels=%d) from epoch %d val_iou=%.4f", + "Loaded %s (in_channels=%d, n_classes=%d) from epoch %d val_iou=%.4f", arch, in_ch, + n_classes, ckpt.get("epoch", 0), ckpt.get("val_iou", 0.0), ) @@ -112,8 +117,16 @@ def export_onnx( # ONNX validation # --------------------------------------------------------------------------- -def validate_onnx(onnx_path: Path, in_channels: int, image_size: int) -> float: - """Run a forward pass with onnxruntime and return inference latency (ms).""" +def validate_onnx( + onnx_path: Path, in_channels: int, image_size: int, n_classes: int = 2 +) -> float: + """ + Run a forward pass with onnxruntime and return inference latency (ms). + + Also verifies the output tensor signature is ``(N, n_classes, H, W)`` so a + shape regression (wrong channel count, missing dynamic axis) is caught at + export time rather than at serve time. + """ try: import onnxruntime as ort import numpy as np @@ -128,6 +141,15 @@ def validate_onnx(onnx_path: Path, in_channels: int, image_size: int) -> float: ) dummy = np.random.rand(1, in_channels, image_size, image_size).astype(np.float32) + # Verify output signature matches (N, n_classes, H, W) + out = sess.run(None, {"image": dummy})[0] + expected = (1, n_classes, image_size, image_size) + if tuple(out.shape) != expected: + raise ValueError( + f"ONNX output shape {tuple(out.shape)} != expected {expected} " + f"for {onnx_path}" + ) + # Warm-up for _ in range(3): sess.run(None, {"image": dummy}) @@ -186,58 +208,50 @@ def benchmark_pytorch(model: nn.Module, in_channels: int, image_size: int) -> fl # Main # --------------------------------------------------------------------------- -def main() -> None: - args = parse_args() - - ckpt_path = Path(args.checkpoint) - if not ckpt_path.exists(): - logger.error("Checkpoint not found: %s", ckpt_path) - sys.exit(1) - - model, cfg = load_model(ckpt_path) - - in_channels = cfg.get("model", {}).get("in_channels", 4) - image_size = args.image_size +def export_checkpoint( + ckpt_path: Path, + onnx_path: Path, + image_size: int, + opset: int, + no_quantize: bool, + n_classes: int = 2, +) -> dict: + """Export a single ``.pth`` checkpoint to ONNX and return export metadata.""" + model, cfg = load_model(ckpt_path, n_classes=n_classes) - # Determine output paths - run_dir = ckpt_path.parent - onnx_path = Path(args.out) if args.out else run_dir / "model.onnx" - quantized_path = onnx_path.parent / "model_quantized.onnx" + in_channels = getattr(model, "n_channels", cfg.get("model", {}).get("in_channels", 4)) + quantized_path = onnx_path.parent / f"{onnx_path.stem}_quantized.onnx" # PyTorch baseline latency pt_ms = benchmark_pytorch(model, in_channels, image_size) logger.info("PyTorch (CPU) baseline: %.1f ms", pt_ms) - # Export export_onnx( model=model, onnx_path=onnx_path, image_size=image_size, in_channels=in_channels, - opset=args.opset, + opset=opset, ) - # Validate - onnx_ms = validate_onnx(onnx_path, in_channels, image_size) + onnx_ms = validate_onnx(onnx_path, in_channels, image_size, n_classes) - # Quantize q_ms = -1.0 - if not args.no_quantize: + if not no_quantize: quantize_onnx(onnx_path, quantized_path) if quantized_path.exists(): - q_ms = validate_onnx(quantized_path, in_channels, image_size) + q_ms = validate_onnx(quantized_path, in_channels, image_size, n_classes) - # Export metadata ckpt = torch.load(ckpt_path, map_location="cpu") info = { "checkpoint": str(ckpt_path), "architecture": cfg.get("model", {}).get("architecture", "unknown"), "in_channels": in_channels, - "num_classes": 2, + "num_classes": n_classes, "image_size": image_size, - "onnx_opset": args.opset, + "onnx_opset": opset, "onnx_path": str(onnx_path), - "quantized_path": str(quantized_path) if not args.no_quantize else None, + "quantized_path": str(quantized_path) if not no_quantize else None, "val_iou": ckpt.get("val_iou", None), "val_f1": ckpt.get("val_f1", None), "epoch": ckpt.get("epoch", None), @@ -247,17 +261,17 @@ def main() -> None: "onnx_int8_cpu": round(q_ms, 2) if q_ms > 0 else None, }, } - info_path = onnx_path.parent / "export_info.json" + info_path = onnx_path.parent / f"{onnx_path.stem}_export_info.json" with open(info_path, "w") as f: json.dump(info, f, indent=2) logger.info("Export metadata saved to %s", info_path) # Summary print("\n" + "=" * 55) - print(" Export Summary") + print(f" Export Summary — {onnx_path.name}") print("=" * 55) print(f" ONNX model : {onnx_path}") - if not args.no_quantize and quantized_path.exists(): + if not no_quantize and quantized_path.exists(): print(f" INT8 model : {quantized_path}") print(f" Val IoU : {info['val_iou']:.4f}" if info["val_iou"] else " Val IoU : N/A") print(f" PyTorch (CPU): {pt_ms:.1f} ms") @@ -266,6 +280,87 @@ def main() -> None: if q_ms > 0: print(f" INT8 (CPU) : {q_ms:.1f} ms ({pt_ms / q_ms:.1f}× speedup)") print("=" * 55) + return info + + +def _discover_analysis_checkpoints() -> list[tuple[str, Path, int]]: + """ + Resolve (analysis_type, checkpoint_path, num_classes) for every enabled + analysis type that declares neural ``weights`` in config.yaml. + """ + from climatevision.data.band_mapping import ( + get_model_config, + list_enabled_analysis_types, + ) + + found: list[tuple[str, Path, int]] = [] + for analysis_type in list_enabled_analysis_types(): + model_cfg = get_model_config(analysis_type) + weights = model_cfg.get("weights") + if not weights: # e.g. SAR ensemble — no neural checkpoint to export + continue + ckpt = PROJECT_ROOT / weights + n_classes = model_cfg.get("num_classes", 2) + found.append((analysis_type, ckpt, n_classes)) + return found + + +def _export_all(args: argparse.Namespace) -> None: + """Export every enabled analysis type's checkpoint to ``models/unet_.onnx``.""" + targets = _discover_analysis_checkpoints() + if not targets: + logger.error("No analysis types with neural weights found in config.yaml") + sys.exit(1) + + exported, skipped = 0, 0 + for analysis_type, ckpt_path, n_classes in targets: + if not ckpt_path.exists(): + logger.warning("Skipping %s — checkpoint not found: %s", analysis_type, ckpt_path) + skipped += 1 + continue + onnx_path = PROJECT_ROOT / "models" / f"unet_{analysis_type}.onnx" + logger.info("Exporting %s → %s", analysis_type, onnx_path) + export_checkpoint( + ckpt_path=ckpt_path, + onnx_path=onnx_path, + image_size=args.image_size, + opset=args.opset, + no_quantize=args.no_quantize, + n_classes=n_classes, + ) + exported += 1 + + logger.info("Batch export complete — %d exported, %d skipped", exported, skipped) + if exported == 0: + sys.exit(1) + + +def main() -> None: + args = parse_args() + + if args.all: + _export_all(args) + return + + if not args.checkpoint: + logger.error("Provide --checkpoint or --all") + sys.exit(1) + + ckpt_path = Path(args.checkpoint) + if not ckpt_path.exists(): + logger.error("Checkpoint not found: %s", ckpt_path) + sys.exit(1) + + run_dir = ckpt_path.parent + onnx_path = Path(args.out) if args.out else run_dir / "model.onnx" + + export_checkpoint( + ckpt_path=ckpt_path, + onnx_path=onnx_path, + image_size=args.image_size, + opset=args.opset, + no_quantize=args.no_quantize, + ) print() print("Serve with:") print(f" onnxruntime → sess = ort.InferenceSession('{onnx_path}')") @@ -274,7 +369,9 @@ def main() -> None: def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description="Export ClimateVision model to ONNX") - p.add_argument("--checkpoint", required=True, help="Path to best_model.pth") + p.add_argument("--checkpoint", default=None, help="Path to best_model.pth") + p.add_argument("--all", action="store_true", + help="Export every enabled analysis type to models/unet_.onnx") p.add_argument("--out", default=None, help="ONNX output path (default: /model.onnx)") p.add_argument("--image-size", type=int, default=256, diff --git a/src/climatevision/inference/pipeline.py b/src/climatevision/inference/pipeline.py index 1d43d01..9ef899c 100644 --- a/src/climatevision/inference/pipeline.py +++ b/src/climatevision/inference/pipeline.py @@ -11,6 +11,7 @@ import json import logging +import time from pathlib import Path from typing import Any, Optional @@ -102,6 +103,8 @@ class _ONNXSegmenter: pipeline needs no changes whether the backing model is PyTorch or ONNX. """ + backend = "onnx" + def __init__(self, session: Any, n_channels: int, n_classes: int, input_name: str): self._session = session self.n_channels = n_channels @@ -121,7 +124,17 @@ def __call__(self, tensor: "torch.Tensor") -> "torch.Tensor": def _load_model(analysis_type: str = "deforestation") -> tuple[Any, torch.device]: - """Load (or return cached) U-Net model configured for the analysis type.""" + """ + Load (or return cached) segmentation model configured for the analysis type. + + Backend priority (issue #12): + 1. **ONNX Runtime** — when an exported ``.onnx`` model is present and + onnxruntime can load it. Faster and lower-memory for serving. + 2. **PyTorch ``.pth`` checkpoint** — the trained weights (with EMA + overlay), used when no servable ONNX model exists. + 3. **Untrained U-Net** — demo weights, so the API still responds when no + trained model has been provisioned. + """ if analysis_type in _model_cache: return _model_cache[analysis_type] @@ -130,8 +143,21 @@ def _load_model(analysis_type: str = "deforestation") -> tuple[Any, torch.device n_channels = model_cfg.get("in_channels", 4) n_classes = model_cfg.get("num_classes", 2) - model = UNet(n_channels=n_channels, n_classes=n_classes) + # 1. Prefer an exported ONNX model when one is present and loadable. + onnx_path = _find_onnx(analysis_type) + if onnx_path is not None: + onnx_model = _try_load_onnx(onnx_path, n_channels, n_classes) + if onnx_model is not None: + logger.info( + "Loaded %s model from ONNX %s (%d channels, %d classes).", + analysis_type, onnx_path, n_channels, n_classes, + ) + onnx_model.eval() + _model_cache[analysis_type] = (onnx_model, device) + return onnx_model, device + # 2. Fall back to the PyTorch checkpoint. + model = UNet(n_channels=n_channels, n_classes=n_classes) model_path = _find_best_checkpoint(analysis_type) if model_path is not None: checkpoint = torch.load(model_path, map_location=device) @@ -157,17 +183,7 @@ def _load_model(analysis_type: str = "deforestation") -> tuple[Any, torch.device checkpoint.get("val_iou", 0.0), ) else: - onnx_path = _find_onnx(analysis_type) - onnx_model = _try_load_onnx(onnx_path, n_channels, n_classes) if onnx_path else None - if onnx_model is not None: - logger.info( - "Loaded %s model from ONNX %s (%d channels, %d classes).", - analysis_type, onnx_path, n_channels, n_classes, - ) - onnx_model.eval() - _model_cache[analysis_type] = (onnx_model, device) - return onnx_model, device - + # 3. No servable model at all — fall back to untrained demo weights. logger.warning( "No trained model found for %s under %s — using untrained weights (demo).", analysis_type, @@ -216,6 +232,60 @@ def _try_load_onnx( return _ONNXSegmenter(session, n_channels, n_classes, input_name) +def benchmark_inference( + analysis_type: str = "deforestation", + *, + image_size: int = 256, + runs: int = 20, +) -> dict[str, Any]: + """ + Time the PyTorch and ONNX backends side by side for an analysis type. + + Loads an untrained PyTorch U-Net and, when available, the exported ONNX + model, then times a batch-1 forward pass on a dummy input. Useful for + verifying the speedup that motivates preferring ONNX at serve time + (issue #12). ``onnx_ms`` is ``None`` when no servable ONNX model exists. + """ + model_cfg = get_model_config(analysis_type) + n_channels = model_cfg.get("in_channels", 4) + n_classes = model_cfg.get("num_classes", 2) + dummy = torch.zeros(1, n_channels, image_size, image_size) + + def _time(model: Any) -> float: + with torch.no_grad(): + for _ in range(3): # warm-up + model(dummy) + t0 = time.perf_counter() + for _ in range(runs): + model(dummy) + return (time.perf_counter() - t0) / runs * 1000.0 + + pytorch_ms = _time(UNet(n_channels=n_channels, n_classes=n_classes).eval()) + + onnx_ms: Optional[float] = None + onnx_path = _find_onnx(analysis_type) + if onnx_path is not None: + onnx_model = _try_load_onnx(onnx_path, n_channels, n_classes) + if onnx_model is not None: + onnx_ms = _time(onnx_model) + + speedup = round(pytorch_ms / onnx_ms, 2) if onnx_ms else None + result = { + "analysis_type": analysis_type, + "backend_active": "onnx" if onnx_ms is not None else "pytorch", + "pytorch_ms": round(pytorch_ms, 2), + "onnx_ms": round(onnx_ms, 2) if onnx_ms is not None else None, + "speedup": speedup, + } + logger.info( + "Backend benchmark %s: pytorch=%.1f ms onnx=%s speedup=%s", + analysis_type, pytorch_ms, + f"{onnx_ms:.1f} ms" if onnx_ms is not None else "n/a", + f"{speedup}x" if speedup else "n/a", + ) + return result + + # --------------------------------------------------------------------------- # Sentinel-2 normalisation statistics (matches preprocessing.py) # Band order: [Red, Green, Blue, NIR] @@ -345,11 +415,19 @@ def run_inference( tensor = torch.FloatTensor(image.astype(np.float32).tolist()).unsqueeze(0) # (1, C, H, W) tensor = tensor.to(device) + backend = getattr(model, "backend", "pytorch") with torch.no_grad(): + _t0 = time.perf_counter() output = model(tensor) + inference_ms = (time.perf_counter() - _t0) * 1000.0 predictions = torch.argmax(output, dim=1) # (1, H, W) probabilities = torch.softmax(output, dim=1) # (1, n_classes, H, W) + logger.info( + "Inference %s: backend=%s latency=%.1f ms (batch=1, %dx%d)", + analysis_type, backend, inference_ms, h, w, + ) + output_check = _guard.check_output( predictions.cpu().numpy(), probabilities=probabilities.cpu().numpy(), @@ -381,6 +459,8 @@ def run_inference( "image_size": [h, w], "num_classes": n_classes, "mean_confidence": round(mean_confidence, 4), + "backend": backend, + "inference_ms": round(inference_ms, 2), **class_pixels, **class_percentages, } diff --git a/tests/test_onnx_inference.py b/tests/test_onnx_inference.py new file mode 100644 index 0000000..60555d3 --- /dev/null +++ b/tests/test_onnx_inference.py @@ -0,0 +1,170 @@ +"""Tests for the ONNX-preferred inference path and PyTorch fallback (issue #12).""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +import climatevision.inference.pipeline as pipeline +from climatevision.inference.pipeline import ( + UNet, + _ONNXSegmenter, + _load_model, + benchmark_inference, + run_inference, +) + + +@pytest.fixture(autouse=True) +def _clear_cache(): + """Each test starts with an empty model cache.""" + pipeline._model_cache.clear() + yield + pipeline._model_cache.clear() + + +def _fake_onnx_model(n_channels: int = 4, n_classes: int = 2): + """A stand-in that quacks like an ONNX-backed segmenter without onnxruntime.""" + model = UNet(n_channels=n_channels, n_classes=n_classes) + model.backend = "onnx" # type: ignore[attr-defined] + return model + + +# --------------------------------------------------------------------------- +# Backend selection priority +# --------------------------------------------------------------------------- + +def test_onnx_is_preferred_when_present(monkeypatch): + """When a servable ONNX model loads, _load_model must use it over the .pth.""" + monkeypatch.setattr(pipeline, "_find_onnx", lambda a: Path("models/unet_deforestation.onnx")) + monkeypatch.setattr(pipeline, "_try_load_onnx", lambda p, c, k: _fake_onnx_model(c, k)) + # A .pth also exists — ONNX must still win. + monkeypatch.setattr(pipeline, "_find_best_checkpoint", lambda a: Path("models/best_model.pth")) + + model, _ = _load_model("deforestation") + assert getattr(model, "backend", "pytorch") == "onnx" + + +def test_falls_back_to_pytorch_when_no_onnx(monkeypatch): + """No ONNX file present → PyTorch UNet backend.""" + monkeypatch.setattr(pipeline, "_find_onnx", lambda a: None) + monkeypatch.setattr(pipeline, "_find_best_checkpoint", lambda a: None) + + model, _ = _load_model("deforestation") + assert isinstance(model, UNet) + assert getattr(model, "backend", "pytorch") == "pytorch" + + +def test_falls_back_when_onnx_unloadable(monkeypatch): + """ONNX path present but unloadable (e.g. LFS pointer) → PyTorch fallback.""" + monkeypatch.setattr(pipeline, "_find_onnx", lambda a: Path("models/unet_deforestation.onnx")) + monkeypatch.setattr(pipeline, "_try_load_onnx", lambda p, c, k: None) + monkeypatch.setattr(pipeline, "_find_best_checkpoint", lambda a: None) + + model, _ = _load_model("deforestation") + assert isinstance(model, UNet) + assert getattr(model, "backend", "pytorch") == "pytorch" + + +# --------------------------------------------------------------------------- +# Latency / backend surfaced in the response +# --------------------------------------------------------------------------- + +def test_run_inference_reports_backend_and_latency(monkeypatch): + monkeypatch.setattr(pipeline, "_find_onnx", lambda a: None) + monkeypatch.setattr(pipeline, "_find_best_checkpoint", lambda a: None) + + image = np.zeros((4, 64, 64), dtype=np.float32) + result = run_inference(image, analysis_type="deforestation") + inf = result["inference"] + + assert inf["backend"] == "pytorch" + assert "inference_ms" in inf + assert isinstance(inf["inference_ms"], float) + assert inf["inference_ms"] >= 0.0 + + +def test_onnx_segmenter_labels_backend(): + assert _ONNXSegmenter.backend == "onnx" + assert getattr(UNet(n_channels=4, n_classes=2), "backend", "pytorch") == "pytorch" + + +# --------------------------------------------------------------------------- +# Benchmark helper +# --------------------------------------------------------------------------- + +def test_benchmark_inference_without_onnx(monkeypatch): + monkeypatch.setattr(pipeline, "_find_onnx", lambda a: None) + + result = benchmark_inference("deforestation", image_size=64, runs=2) + assert result["backend_active"] == "pytorch" + assert result["pytorch_ms"] > 0 + assert result["onnx_ms"] is None + assert result["speedup"] is None + + +# --------------------------------------------------------------------------- +# End-to-end with a real ONNX export +# --------------------------------------------------------------------------- + +def test_end_to_end_onnx_roundtrip(monkeypatch, tmp_path): + """Export a small UNet to ONNX, then confirm the pipeline loads and serves it.""" + ort = pytest.importorskip("onnxruntime") + import torch + + onnx_path = tmp_path / "unet_deforestation.onnx" + model = UNet(n_channels=4, n_classes=2).eval() + dummy = torch.zeros(1, 4, 64, 64) + export_kwargs = dict( + input_names=["image"], + output_names=["logits"], + opset_version=17, + dynamic_axes={"image": {0: "batch", 2: "h", 3: "w"}, + "logits": {0: "batch", 2: "h", 3: "w"}}, + ) + try: + # dynamo=False forces the legacy TorchScript exporter, which does not + # require the optional onnxscript dependency. + try: + torch.onnx.export(model, dummy, str(onnx_path), dynamo=False, **export_kwargs) + except TypeError: # older torch without the dynamo kwarg + torch.onnx.export(model, dummy, str(onnx_path), **export_kwargs) + except Exception as exc: # optional export tooling (onnx/onnxscript) may be absent + if "onnx" in str(exc).lower(): + pytest.skip(f"ONNX export tooling unavailable: {exc}") + raise + assert onnx_path.exists() + + monkeypatch.setattr(pipeline, "_find_onnx", lambda a: onnx_path) + monkeypatch.setattr(pipeline, "_find_best_checkpoint", lambda a: None) + + loaded, _ = _load_model("deforestation") + assert isinstance(loaded, _ONNXSegmenter) + + image = np.zeros((4, 64, 64), dtype=np.float32) + result = run_inference(image, analysis_type="deforestation") + assert result["inference"]["backend"] == "onnx" + assert result["inference"]["image_size"] == [64, 64] + + +# --------------------------------------------------------------------------- +# Export discovery +# --------------------------------------------------------------------------- + +def test_discover_analysis_checkpoints_covers_enabled_types(): + import importlib.util + + root = Path(pipeline.__file__).resolve().parents[3] + spec = importlib.util.spec_from_file_location( + "export_model", root / "scripts" / "export_model.py" + ) + export_model = importlib.util.module_from_spec(spec) + spec.loader.exec_module(export_model) + + targets = {t[0]: (t[1], t[2]) for t in export_model._discover_analysis_checkpoints()} + assert "deforestation" in targets + ckpt, n_classes = targets["deforestation"] + assert n_classes == 2 + assert str(ckpt).endswith("unet_deforestation.pth")