v3: resize-convolution upsampling — removes checkerboard artifacts

src/export_onnx.py

@@ -1,45 +1,50 @@
 """
-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
+export_onnx.py — export a trained U-Net to ONNX, auto-detecting channels (v1=2, v2=3).
+
+  python src\\export_onnx.py --ckpt checkpoints\\best.pt     --out web\\radio_unet_v1.onnx
+  python src\\export_onnx.py --ckpt checkpoints_v2\\best.pt  --out web\\radio_unet_v2.onnx
 """
+import argparse
 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()
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ckpt", default="checkpoints/best.pt")
+    p.add_argument("--out", default="web/radio_unet.onnx")
+    args = p.parse_args()
 
-# --- 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}")
+    ckpt = torch.load(args.ckpt, map_location="cpu", weights_only=False)
+    cargs = ckpt.get("args", {})
+    base = cargs.get("base", 64)
+    add_dist = bool(cargs.get("distance", False))
+    in_ch = 3 if add_dist else 2
 
-# --- verify ONNX output matches PyTorch on a real test sample ---
-import onnxruntime as ort
+    model = UNet(in_channels=in_ch, out_channels=1, base=base)
+    model.load_state_dict(ckpt["model_state"])
+    model.eval()
+    print(f"Loaded {args.ckpt} | in_channels={in_ch} | distance={add_dist}")
 
-ds = RadioMapSeerDataset(split="test", img_size=256)
-x, _ = ds[0]
-xb = x.unsqueeze(0).numpy().astype(np.float32)
+    dummy = torch.randn(1, in_ch, 256, 256)
+    torch.onnx.export(model, dummy, args.out,
+                      input_names=["input"], output_names=["coverage"],
+                      opset_version=17)
+    print(f"Exported {args.out}")
 
-with torch.no_grad():
-    torch_out = model(x.unsqueeze(0)).numpy()
+    import onnxruntime as ort
+    ds = RadioMapSeerDataset(split="test", img_size=256, add_distance=add_dist)
+    x, _ = ds[0]
+    xb = x.unsqueeze(0).numpy().astype(np.float32)
+    with torch.no_grad():
+        t_out = model(x.unsqueeze(0)).numpy()
+    sess = ort.InferenceSession(args.out, providers=["CPUExecutionProvider"])
+    o_out = sess.run(["coverage"], {"input": xb})[0]
+    print(f"Max |PyTorch - ONNX|: {np.max(np.abs(t_out - o_out)):.2e}")
 
-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")
+if __name__ == "__main__":
+    main()