v4:Two-stage cascade (WNet)

checkpoints_v4/history.csv

@@ -0,0 +1,21 @@
+epoch,train_mse,val_mse
+1,0.000914,0.000824
+2,0.000637,0.000939
+3,0.000570,0.000672
+4,0.000504,0.000629
+5,0.000455,0.000595
+6,0.000416,0.000640
+7,0.000391,0.000564
+8,0.000365,0.000544
+9,0.000344,0.000527
+10,0.000327,0.000515
+11,0.000309,0.000534
+12,0.000297,0.000504
+13,0.000284,0.000517
+14,0.000271,0.000462
+15,0.000263,0.000468
+16,0.000250,0.000450
+17,0.000243,0.000460
+18,0.000234,0.000448
+19,0.000227,0.000446
+20,0.000220,0.000443

src/export_onnx.py

@@ -1,14 +1,16 @@
 """
-export_onnx.py — export a trained U-Net to ONNX, auto-detecting channels (v1=2, v2=3).
+export_onnx.py — export a trained model to ONNX, auto-detecting config from the checkpoint.
 
-  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
+  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
+  python src\\export_onnx.py --ckpt checkpoints_v3\\best.pt     --out web\\radio_unet_v3.onnx
+  python src\\export_onnx.py --ckpt checkpoints_v4\\best.pt     --out web\\radio_unet_v4.onnx
 """
 import argparse
 import numpy as np
 import torch
 
-from model import UNet
+from model import UNet, WNet
 from dataset import RadioMapSeerDataset
 
 
@@ -22,12 +24,16 @@ 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")
+    arch = cargs.get("arch", "unet")
     in_ch = 3 if add_dist else 2
 
-    model = UNet(in_channels=in_ch, out_channels=1, base=base)
+    Net = WNet if arch == "wnet" else UNet
+    model = Net(in_channels=in_ch, out_channels=1, base=base, up_mode=up_mode)
     model.load_state_dict(ckpt["model_state"])
     model.eval()
-    print(f"Loaded {args.ckpt} | in_channels={in_ch} | distance={add_dist}")
+    print(f"Loaded {args.ckpt} | arch={arch} | in_channels={in_ch} | "
+          f"distance={add_dist} | up_mode={up_mode}")
 
     dummy = torch.randn(1, in_ch, 256, 256)
     torch.onnx.export(model, dummy, args.out,

src/model.py

@@ -89,17 +89,44 @@ class UNet(nn.Module):
         x = self.up4(x, x1)
         return torch.sigmoid(self.outc(x))
 
+class WNet(nn.Module):
+    """Two cascaded U-Nets (RadioUNet-style): a coarse predictor + a refiner.
+
+    unet1: input channels         -> coarse coverage map (1 ch)
+    unet2: input channels + coarse -> refined coverage map (1 ch)
+
+    The refiner sees a global coverage estimate as an extra input channel,
+    so every pixel has long-range context -> larger effective receptive field.
+    forward() returns the refined map by default (so eval/ONNX stay single-output);
+    pass return_coarse=True during training for deep supervision.
+    """
+    def __init__(self, in_channels=3, out_channels=1, base=64, up_mode="resize"):
+        super().__init__()
+        self.unet1 = UNet(in_channels,     out_channels, base=base, up_mode=up_mode)
+        self.unet2 = UNet(in_channels + 1, out_channels, base=base, up_mode=up_mode)
+
+    def forward(self, x, return_coarse=False):
+        coarse = self.unet1(x)
+        refined = self.unet2(torch.cat([x, coarse], dim=1))
+        if return_coarse:
+            return refined, coarse
+        return refined
 
 if __name__ == "__main__":
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = UNet(in_channels=2, out_channels=1, base=64).to(device)
-    n_params = sum(p.numel() for p in model.parameters())
     print(f"Device: {device}")
-    print(f"Parameters: {n_params:,}")
 
-    x = torch.randn(4, 2, 256, 256, device=device)
-    y = model(x)
-    print("Input :", tuple(x.shape))
-    print("Output:", tuple(y.shape), "| range:", round(float(y.min()), 3), "-", round(float(y.max()), 3))
+    unet = UNet(in_channels=3, out_channels=1, base=64, up_mode="resize").to(device)
+    print(f"UNet params: {sum(p.numel() for p in unet.parameters()):,}")
+
+    wnet = WNet(in_channels=3, out_channels=1, base=64, up_mode="resize").to(device)
+    print(f"WNet params: {sum(p.numel() for p in wnet.parameters()):,}")
+
+    x = torch.randn(2, 3, 256, 256, device=device)
+    refined, coarse = wnet(x, return_coarse=True)
+    print("Input  :", tuple(x.shape))
+    print("Refined:", tuple(refined.shape),
+          "| range:", round(float(refined.min()), 3), "-", round(float(refined.max()), 3))
+    print("Coarse :", tuple(coarse.shape))
     if device == "cuda":
-        print(f"Peak VRAM (forward, batch 4): {torch.cuda.max_memory_allocated()/1e9:.2f} GB")
+        print(f"Peak VRAM (WNet forward, batch 2): {torch.cuda.max_memory_allocated()/1e9:.2f} GB")

src/train.py

@@ -1,14 +1,14 @@
 """
-train.py — train the U-Net on RadioMapSeer coverage maps.
+train.py — train the radio-coverage models on RadioMapSeer.
 
-v1: 2-channel input (buildings, transmitter).
-v2: add --distance for a 3rd channel (distance-to-Tx) that fixes the square cutoff.
+v1: 2-channel UNet.  v2: --distance (3-channel).  v3: --up-mode resize.
+v4: --arch wnet  (two cascaded U-Nets: coarse + refine, deep supervision).
 
 Examples:
-  # v1
-  python src\\train.py --epochs 30 --batch-size 32 --num-workers 8
-  # v2 (separate checkpoint dir so v1 stays intact)
-  python src\\train.py --distance --epochs 30 --batch-size 32 --num-workers 8 --ckpt-dir checkpoints_v2
+  # v3 (single U-Net, resize-conv)
+  python src\\train.py --distance --up-mode resize --epochs 30 --batch-size 32 --num-workers 8 --ckpt-dir checkpoints_v3
+  # v4 (WNet), warm-starting stage 1 from the trained v3
+  python src\\train.py --arch wnet --distance --up-mode resize --init-unet1-from checkpoints_v3\\best.pt --epochs 25 --batch-size 16 --num-workers 8 --ckpt-dir checkpoints_v4
 """
 
 import argparse
@@ -21,12 +21,13 @@ import torch.nn as nn
 from torch.utils.data import DataLoader
 
 from dataset import RadioMapSeerDataset
-from model import UNet
+from model import UNet, WNet
 
 torch.backends.cudnn.benchmark = True   # autotune convs (input size is fixed)
 
 
 def evaluate(net, loader, criterion, device, amp):
+    """Validation MSE on the final output (WNet returns the refined map by default)."""
     net.eval()
     total, n = 0.0, 0
     with torch.no_grad():
@@ -50,16 +51,25 @@ def main():
     p.add_argument("--num-workers", type=int, default=4)
     p.add_argument("--ckpt-dir", default="checkpoints")
     p.add_argument("--distance", action="store_true",
-                   help="v2: add a distance-to-transmitter input channel")
+                   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)")
+                   help="v3+: 'resize' = bilinear upsample + conv (removes checkerboard)")
+    p.add_argument("--arch", default="unet", choices=["unet", "wnet"],
+                   help="v4: 'wnet' = two cascaded U-Nets (coarse + refine)")
+    p.add_argument("--aux-weight", type=float, default=0.4,
+                   help="WNet deep-supervision weight on the coarse output")
+    p.add_argument("--init-from", default=None,
+                   help="warm-start: load full model weights from a checkpoint")
+    p.add_argument("--init-unet1-from", default=None,
+                   help="WNet: warm-start stage-1 (unet1) from a single-UNet checkpoint (e.g. v3)")
     args = p.parse_args()
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     amp = (device == "cuda")
-    print(f"Device: {device} | AMP: {amp} | distance: {args.distance} | compile: {args.compile}")
+    print(f"Device: {device} | AMP: {amp} | arch: {args.arch} | distance: {args.distance} | "
+          f"up_mode: {args.up_mode} | compile: {args.compile}")
 
     train_ds = RadioMapSeerDataset(split="train", img_size=args.img_size,
                                    max_samples=args.max_samples, add_distance=args.distance)
@@ -74,8 +84,26 @@ 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, up_mode=args.up_mode).to(device)
+    if args.arch == "wnet":
+        model = WNet(in_channels=in_ch, out_channels=1, base=args.base, up_mode=args.up_mode).to(device)
+    else:
+        model = UNet(in_channels=in_ch, out_channels=1, base=args.base, up_mode=args.up_mode).to(device)
+
+    # WNet: warm-start stage 1 from an existing single-UNet checkpoint (same in_ch + up_mode)
+    if args.arch == "wnet" and args.init_unet1_from:
+        ck1 = torch.load(args.init_unet1_from, map_location=device, weights_only=False)
+        model.unet1.load_state_dict(ck1["model_state"])
+        print(f"Warm-started unet1 from {args.init_unet1_from} "
+              f"(epoch {ck1.get('epoch','?')}, val {ck1.get('val_loss','?')})")
+
     net = torch.compile(model) if args.compile else model
+
+    if args.init_from:
+        ckpt0 = torch.load(args.init_from, map_location=device, weights_only=False)
+        model.load_state_dict(ckpt0["model_state"])
+        print(f"Warm-started from {args.init_from} "
+              f"(epoch {ckpt0.get('epoch','?')}, val {ckpt0.get('val_loss','?')})")
+
     criterion = nn.MSELoss()
     optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
     scaler = torch.amp.GradScaler(enabled=amp)
@@ -96,20 +124,26 @@ def main():
             x, y = x.to(device), y.to(device)
             optimizer.zero_grad(set_to_none=True)
             with torch.autocast(device_type=device, dtype=torch.float16, enabled=amp):
-                loss = criterion(net(x), y)
+                if args.arch == "wnet":
+                    refined, coarse = net(x, return_coarse=True)
+                    main_loss = criterion(refined, y)                       # the real metric
+                    loss = main_loss + args.aux_weight * criterion(coarse, y)  # + deep supervision
+                else:
+                    main_loss = criterion(net(x), y)
+                    loss = main_loss
             scaler.scale(loss).backward()
             scaler.step(optimizer)
             scaler.update()
-            running += loss.item() * x.size(0)
+            running += main_loss.item() * x.size(0)   # log refined MSE (comparable to val)
             n += x.size(0)
             if i % 50 == 0:
-                print(f"  epoch {epoch} step {i}/{len(train_loader)} | loss {running/n:.5f}")
+                print(f"  epoch {epoch} step {i}/{len(train_loader)} | loss {running/n:.6f}")
 
         train_loss = running / n
         val_loss = evaluate(net, val_loader, criterion, device, amp)
         dt = time.time() - t0
-        print(f"[epoch {epoch}/{args.epochs}] train MSE {train_loss:.5f} | "
-              f"val MSE {val_loss:.5f} | val RMSE {val_loss**0.5:.5f} | {dt:.0f}s")
+        print(f"[epoch {epoch}/{args.epochs}] train MSE {train_loss:.6f} | "
+              f"val MSE {val_loss:.6f} | val RMSE {val_loss**0.5:.5f} | {dt:.0f}s")
         with open(history_path, "a", newline="") as f:
             csv.writer(f).writerow([epoch, f"{train_loss:.6f}", f"{val_loss:.6f}"])
 
@@ -118,13 +152,13 @@ def main():
             torch.save({"epoch": epoch, "model_state": model.state_dict(),
                         "val_loss": val_loss, "args": vars(args)},
                        ckpt_dir / "best.pt")
-            print(f"           -> saved best.pt (val MSE {val_loss:.5f})")
+            print(f"           -> saved best.pt (val MSE {val_loss:.6f})")
 
         if device == "cuda":
             print(f"           peak VRAM: {torch.cuda.max_memory_allocated()/1e9:.2f} GB")
             torch.cuda.reset_peak_memory_stats()
 
-    print(f"\nDone. Best val MSE: {best_val:.5f}  (RMSE {best_val**0.5:.5f})")
+    print(f"\nDone. Best val MSE: {best_val:.6f}  (RMSE {best_val**0.5:.5f})")
 
 
 if __name__ == "__main__":

src/visualize.py

@@ -16,7 +16,8 @@ import torch
 import matplotlib.pyplot as plt
 
 from dataset import RadioMapSeerDataset
-from model import UNet
+from model import UNet, WNet
+
 
 
 def main():
@@ -36,7 +37,13 @@ def main():
     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, up_mode=up_mode).to(device)
+    arch = cargs.get("arch", "unet")
+    if arch == "wnet":
+        model = WNet(in_channels=in_ch, out_channels=1, base=base, up_mode=up_mode).to(device)
+    else:
+        arch = cargs.get("arch", "unet")
+        Net = WNet if arch == "wnet" else UNet
+        model = Net(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} | "

web/index.html

@@ -56,7 +56,8 @@
   <div class="vers">
     <button class="ver-btn" data-v="v1" data-i18n="btn_v1"></button>
     <button class="ver-btn" data-v="v2" data-i18n="btn_v2"></button>
-    <button class="ver-btn" data-v="v3" data-i18n="btn_v3" disabled></button>
+    <button class="ver-btn" data-v="v3" data-i18n="btn_v3"></button>
+    <button class="ver-btn" data-v="v4" data-i18n="btn_v4"></button>
   </div>
   <div class="notice" id="verNotice"></div>
   <div class="metricbar">
@@ -97,18 +98,19 @@ const T = {
     title: "Predicting Wireless Coverage with Deep Learning",
     subtitle: "A U-Net predicts a base station's signal coverage across a city — running live in your browser. Switch between model versions to see how it improved.",
     ver_h: "Model version",
-    btn_v1: "v1 · baseline", btn_v2: "v2 · distance fix", btn_v3: "v3 · soon",
-    m_error: "RMSE on unseen cities",
+    btn_v1: "v1 · baseline", btn_v2: "v2 · distance fix", btn_v3: "v3 · checkerboard fix", btn_v4: "v4 · WNet refine",
+    m_error: "RMSE on held-out cities",
     notice_v1: "Baseline (2-channel input: buildings + transmitter). Known issue: coverage cuts off in a square around the transmitter — the far field stays blank because the network can't propagate the source that far.",
     notice_v2: "Adds a distance-to-transmitter input channel (3-channel). Fixed: the square cutoff — coverage now spans the whole map, with far-field rays and distant building shadows. Known issue: faint checkerboard texture in smooth areas, from transposed-convolution upsampling.",
-    notice_v3: "In progress. Replaces transposed-convolution upsampling with resize-convolution to remove the checkerboard texture. Not yet available.",
+    notice_v3: "Replaces transposed-convolution upsampling with resize-convolution. Fixed: the checkerboard texture — smooth regions are clean. Known issue: along a long unobstructed corridor the far-field beam fades before reaching the edge — a residual receptive-field limit.",
+    notice_v4: "Two cascaded U-Nets — a coarse predictor feeds a refiner. Because the refiner's input already contains a global coverage estimate, every pixel gets long-range context (a larger effective receptive field). Fixed: the fading corridor — coverage stays bright to the edge. Lowest error of all four versions.",
     demo_h: "Try it yourself",
     select_label: "City layout:", click_hint: "Click the map to place a transmitter.",
     input_title: "City map — click to place transmitter", output_title: "Predicted coverage",
     legend_weak: "weak", legend_strong: "strong",
     how_h: "How it works",
     how_p: "The model takes the building layout and the transmitter location as image channels and outputs a coverage heatmap, learned from the public RadioMapSeer dataset (5.9 GHz urban propagation). It was given no propagation equations — it learned signal falloff and building shadowing from data. The model runs entirely in your browser via ONNX Runtime Web; nothing is sent to a server.",
-    footer: "Dataset: RadioMapSeer (Yapar et al., 2022), CC BY 4.0. Independent portfolio project. Error reported on held-out test cities.",
+    footer: "Dataset: RadioMapSeer (Yapar et al., 2022), CC BY 4.0. Independent portfolio project. Errors are on held-out cities the model never trained on.",
     st_loading: "Loading model…", st_ready: "Pick a layout and click to place a transmitter.",
     st_running: "Running inference…",
     st_done: (ms) => "Predicted in " + ms + " ms. Click again to move the transmitter.",
@@ -116,37 +118,39 @@ const T = {
   },
   zh: {
     title: "基于深度学习的无线信号覆盖预测",
-    subtitle: "一个 U-Net 在你的浏览器中实时预测基站的信号覆盖。切换不同模型版本，看看它是如何一步步改进的。",
+    subtitle: "一个 U-Net 在你的浏览器中实时预测基站的信号覆盖。切换不同模型版本,看看它是如何一步步改进的。",
     ver_h: "模型版本",
-    btn_v1: "v1 · 基线", btn_v2: "v2 · 距离修复", btn_v3: "v3 · 即将推出",
-    m_error: "未见城市上的 RMSE",
-    notice_v1: "基线模型（2 通道输入：建筑 + 发射机）。已知问题：覆盖在发射机周围呈方形截断——远场为空白，因为网络无法将信号源的影响传播到远处。",
-    notice_v2: "新增「到发射机的距离」输入通道（3 通道）。已修复：方形截断——覆盖现已贯穿整张地图，远场射线与远处建筑阴影都能呈现。已知问题：平滑区域出现轻微棋盘格纹理，来自转置卷积上采样。",
-    notice_v3: "开发中。用 resize 卷积替换转置卷积上采样，以消除棋盘格纹理。暂未上线。",
+    btn_v1: "v1 · 基线", btn_v2: "v2 · 距离修复", btn_v3: "v3 · 棋盘格修复", btn_v4: "v4 · WNet 精修",
+    m_error: "留出城市上的 RMSE",
+    notice_v1: "基线模型(2 通道输入:建筑 + 发射机)。已知问题:覆盖在发射机周围呈方形截断——远场为空白,因为网络无法将信号源的影响传播到远处。",
+    notice_v2: "新增「到发射机的距离」输入通道(3 通道)。已修复:方形截断——覆盖现已贯穿整张地图,远场射线与远处建筑阴影都能呈现。已知问题:平滑区域出现轻微棋盘格纹理,来自转置卷积上采样。",
+    notice_v3: "用 resize 卷积(双线性上采样 + 卷积)替换转置卷积。已修复:棋盘格纹理消失,平滑区域变得干净。已知问题:在无遮挡的长走廊中,远场光束在到达边缘前逐渐变暗——残留的感受野限制。",
+    notice_v4: "两个级联 U-Net——粗预测网络的输出送入精修网络。由于精修网络的输入已包含一张全局覆盖估计,每个像素都获得了长程上下文(更大的有效感受野)。已修复:长走廊衰减,覆盖一直明亮延伸到边缘。四个版本中误差最低。",
     demo_h: "在线体验",
-    select_label: "城市布局：", click_hint: "点击地图以放置发射机。",
+    select_label: "城市布局:", click_hint: "点击地图以放置发射机。",
     input_title: "城市地图——点击放置发射机", output_title: "预测覆盖",
     legend_weak: "弱", legend_strong: "强",
     how_h: "实现原理",
-    how_p: "模型以建筑布局和发射机位置作为图像通道输入，输出覆盖热力图，训练自公开的 RadioMapSeer 数据集（5.9 GHz 城市传播）。它没有被告知任何传播公式，完全从数据中学会了信号衰减与建筑遮蔽。模型通过 ONNX Runtime Web 完全在你的浏览器中运行，不会向服务器发送任何数据。",
-    footer: "数据集：RadioMapSeer（Yapar 等，2022），CC BY 4.0 许可。独立作品集项目。误差基于留出的测试城市。",
+    how_p: "模型以建筑布局和发射机位置作为图像通道输入,输出覆盖热力图,训练自公开的 RadioMapSeer 数据集(5.9 GHz 城市传播)。它没有被告知任何传播公式,完全从数据中学会了信号衰减与建筑遮蔽。模型通过 ONNX Runtime Web 完全在你的浏览器中运行,不向服务器发送任何数据。",
+    footer: "数据集:RadioMapSeer(Yapar 等,2022),CC BY 4.0 许可。独立作品集项目。误差基于模型从未训练过的留出城市。",
     st_loading: "正在加载模型…", st_ready: "选择布局并点击放置发射机。",
     st_running: "正在推理…",
-    st_done: (ms) => "推理完成，用时 " + ms + " 毫秒。再次点击可移动发射机。",
-    st_error: (m) => "错误：" + m
+    st_done: (ms) => "推理完成,用时 " + ms + " 毫秒。再次点击可移动发射机。",
+    st_error: (m) => "错误:" + m
   }
 };
 
 const VERSIONS = {
   v1: { model: "radio_unet_v1.onnx", channels: 2, available: true },
   v2: { model: "radio_unet_v2.onnx", channels: 3, available: true },
-  v3: { model: "radio_unet_v3.onnx", channels: 3, available: false },
+  v3: { model: "radio_unet_v3.onnx", channels: 3, available: true },
+  v4: { model: "radio_unet_v4.onnx", channels: 3, available: true },
 };
-const METRICS = { v1: "0.052", v2: "0.034", v3: "—" };
+const METRICS = { v1: "0.052", v2: "0.034", v3: "0.030", v4: "0.021" };
 
 const SIZE = 256;
 let currentLang = (navigator.language || "en").toLowerCase().startsWith("zh") ? "zh" : "en";
-let activeVersion = "v2";
+let activeVersion = "v4";
 const sessions = {};
 let buildings = null, txRC = null;
 
@@ -221,7 +225,7 @@ function buildInput(){
   input.set(buildings, 0);                         // channel 0: buildings
   const [r,c] = txRC;
   input[SIZE*SIZE + r*SIZE + c] = 1.0;             // channel 1: transmitter
-  if (ch === 3){                                   // channel 2: distance to Tx (v2/v3)
+  if (ch === 3){                                   // channel 2: distance to Tx (v2/v3/v4)
     const off = 2*SIZE*SIZE, norm = Math.SQRT2 * SIZE;
     for(let yy=0; yy<SIZE; yy++)
       for(let xx=0; xx<SIZE; xx++)