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