v1: U-Net radio coverage prediction (2-channel input, baseline)

src/export_onnx.py

@@ -0,0 +1,45 @@
+"""
+export_onnx.py — export the trained U-Net to ONNX and verify it matches PyTorch.
+Run from the project root:  python src\\export_onnx.py
+"""
+import numpy as np
+import torch
+
+from model import UNet
+from dataset import RadioMapSeerDataset
+
+CKPT = "checkpoints/best.pt"
+OUT = "web/radio_unet.onnx"
+
+# --- load the trained model ---
+ckpt = torch.load(CKPT, map_location="cpu", weights_only=False)
+base = ckpt.get("args", {}).get("base", 64)
+model = UNet(2, 1, base=base)
+model.load_state_dict(ckpt["model_state"])
+model.eval()
+
+# --- export to ONNX (fixed 1x2x256x256 input) ---
+dummy = torch.randn(1, 2, 256, 256)
+torch.onnx.export(
+    model, dummy, OUT,
+    input_names=["input"], output_names=["coverage"],
+    opset_version=17,
+)
+print(f"Exported {OUT}")
+
+# --- verify ONNX output matches PyTorch on a real test sample ---
+import onnxruntime as ort
+
+ds = RadioMapSeerDataset(split="test", img_size=256)
+x, _ = ds[0]
+xb = x.unsqueeze(0).numpy().astype(np.float32)
+
+with torch.no_grad():
+    torch_out = model(x.unsqueeze(0)).numpy()
+
+sess = ort.InferenceSession(OUT, providers=["CPUExecutionProvider"])
+onnx_out = sess.run(["coverage"], {"input": xb})[0]
+
+max_diff = float(np.max(np.abs(torch_out - onnx_out)))
+print(f"Max |PyTorch - ONNX| on a test sample: {max_diff:.2e}")
+print("OK — outputs match" if max_diff < 1e-3 else "WARNING: large mismatch, stop and check")