feat(inference): prefer ONNX Runtime with PyTorch fallback

scripts/export_model.py

@@ -1,15 +1,19 @@
 """
 Export the trained ClimateVision model to ONNX for production serving.
 
-Produces:
-  - <run_dir>/model.onnx           — standard ONNX graph
-  - <run_dir>/model_quantized.onnx — INT8 quantized (CPU-optimised)
-  - <run_dir>/export_info.json     — metadata (opset, input shape, benchmark)
+Produces (per exported model ``<name>``):
+  - <dir>/<name>.onnx              — standard ONNX graph
+  - <dir>/<name>_quantized.onnx    — INT8 quantized (CPU-optimised)
+  - <dir>/<name>_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_<type>.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_<type>.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 <path> 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_<type>.onnx")
     p.add_argument("--out",          default=None,
                    help="ONNX output path (default: <checkpoint_dir>/model.onnx)")
     p.add_argument("--image-size",   type=int, default=256,