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

src/model.py

@@ -42,19 +42,26 @@ class Down(nn.Module):
 
 class Up(nn.Module):
     """Upsample, concatenate the skip connection, then DoubleConv."""
-    def __init__(self, in_ch, skip_ch, out_ch):
+    def __init__(self, in_ch, skip_ch, out_ch, up_mode="deconv"):
         super().__init__()
-        self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size=2, stride=2)
+        if up_mode == "resize":
+            # bilinear upsample + 1x1 conv — uniform, so no checkerboard artifacts
+            self.up = nn.Sequential(
+                nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
+                nn.Conv2d(in_ch, in_ch // 2, kernel_size=1),
+            )
+        else:  # "deconv" (v1/v2): transposed convolution
+            self.up = nn.ConvTranspose2d(in_ch, in_ch // 2, kernel_size=2, stride=2)
         self.conv = DoubleConv(in_ch // 2 + skip_ch, out_ch)
 
     def forward(self, x, skip):
         x = self.up(x)
-        x = torch.cat([skip, x], dim=1)   # skip connection
+        x = torch.cat([skip, x], dim=1)
         return self.conv(x)
 
 
 class UNet(nn.Module):
-    def __init__(self, in_channels=2, out_channels=1, base=64):
+    def __init__(self, in_channels=2, out_channels=1, base=64, up_mode="deconv"):
         super().__init__()
         self.inc   = DoubleConv(in_channels, base)   # 256
         self.down1 = Down(base,     base * 2)        # 128
@@ -62,10 +69,10 @@ class UNet(nn.Module):
         self.down3 = Down(base * 4, base * 8)        # 32
         self.down4 = Down(base * 8, base * 8)        # 16  (bottleneck)
 
-        self.up1 = Up(base * 8, base * 8, base * 4)
-        self.up2 = Up(base * 4, base * 4, base * 2)
-        self.up3 = Up(base * 2, base * 2, base)
-        self.up4 = Up(base,     base,     base)
+        self.up1 = Up(base * 8, base * 8, base * 4, up_mode)
+        self.up2 = Up(base * 4, base * 4, base * 2, up_mode)
+        self.up3 = Up(base * 2, base * 2, base,     up_mode)
+        self.up4 = Up(base,     base,     base,     up_mode)
 
         self.outc = nn.Conv2d(base, out_channels, kernel_size=1)
 

src/train.py

@@ -53,6 +53,8 @@ def main():
                    help="v2: add a distance-to-transmitter input channel")
     p.add_argument("--compile", action="store_true",
                    help="wrap model in torch.compile (optional; flaky on Windows)")
+    p.add_argument("--up-mode", default="deconv", choices=["deconv", "resize"],
+                   help="v3: 'resize' = bilinear upsample + conv (removes checkerboard)")
     args = p.parse_args()
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -72,7 +74,7 @@ def main():
     val_loader = DataLoader(val_ds, shuffle=False, **dl_kwargs)
 
     in_ch = 3 if args.distance else 2
-    model = UNet(in_channels=in_ch, out_channels=1, base=args.base).to(device)
+    model = UNet(in_channels=in_ch, out_channels=1, base=args.base, up_mode=args.up_mode).to(device)
     net = torch.compile(model) if args.compile else model
     criterion = nn.MSELoss()
     optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)

src/visualize.py

@@ -34,8 +34,9 @@ def main():
     cargs = ckpt.get("args", {})
     base = cargs.get("base", 64)
     add_dist = bool(cargs.get("distance", False))
+    up_mode = cargs.get("up_mode", "deconv")
     in_ch = 3 if add_dist else 2
-    model = UNet(in_channels=in_ch, out_channels=1, base=base).to(device)
+    model = UNet(in_channels=in_ch, out_channels=1, base=base, up_mode=up_mode).to(device)
     model.load_state_dict(ckpt["model_state"])
     model.eval()
     print(f"Loaded {args.ckpt} | in_channels={in_ch} | distance={add_dist} | "