diff --git a/outputs_v2/prediction_00.png b/outputs_v2/prediction_00.png new file mode 100644 index 0000000..b2bc753 Binary files /dev/null and b/outputs_v2/prediction_00.png differ diff --git a/outputs_v2/prediction_01.png b/outputs_v2/prediction_01.png new file mode 100644 index 0000000..a8fbcbb Binary files /dev/null and b/outputs_v2/prediction_01.png differ diff --git a/outputs_v2/prediction_02.png b/outputs_v2/prediction_02.png new file mode 100644 index 0000000..43d6b69 Binary files /dev/null and b/outputs_v2/prediction_02.png differ diff --git a/outputs_v2/prediction_03.png b/outputs_v2/prediction_03.png new file mode 100644 index 0000000..b37ae78 Binary files /dev/null and b/outputs_v2/prediction_03.png differ diff --git a/outputs_v2/prediction_04.png b/outputs_v2/prediction_04.png new file mode 100644 index 0000000..36c9be5 Binary files /dev/null and b/outputs_v2/prediction_04.png differ diff --git a/outputs_v2/prediction_05.png b/outputs_v2/prediction_05.png new file mode 100644 index 0000000..135e2bb Binary files /dev/null and b/outputs_v2/prediction_05.png differ diff --git a/src/dataset.py b/src/dataset.py index 067ff6a..c2a123c 100644 --- a/src/dataset.py +++ b/src/dataset.py @@ -41,9 +41,11 @@ class RadioMapSeerDataset(Dataset): } def __init__(self, data_root=DEFAULT_DATA_ROOT, split="train", - sim="DPM", img_size=256, max_samples=None): + sim="DPM", img_size=256, max_samples=None, add_distance=False): self.data_root = Path(data_root) self.img_size = img_size + self.add_distance = add_distance # v2: distance-to-Tx channel + self._build_cache = {} # cache shared building maps self.gain_dir = self.data_root / "gain" / sim self.build_dir = self.data_root / "png" / "buildings_complete" @@ -79,12 +81,26 @@ class RadioMapSeerDataset(Dataset): def __getitem__(self, idx): map_id, tx_id = self.samples[idx] - buildings = _load_gray(self.build_dir / f"{map_id}.png", self.img_size) - antenna = _load_gray(self.ant_dir / f"{map_id}_{tx_id}.png", self.img_size) - gain = _load_gray(self.gain_dir / f"{map_id}_{tx_id}.png", self.img_size) + # buildings are shared across a map's 80 transmitters -> cache to cut I/O + buildings = self._build_cache.get(map_id) + if buildings is None: + buildings = _load_gray(self.build_dir / f"{map_id}.png", self.img_size) + self._build_cache[map_id] = buildings - x = np.stack([buildings, antenna], axis=0) # [2, H, W] - y = gain[None, :, :] # [1, H, W] + antenna = _load_gray(self.ant_dir / f"{map_id}_{tx_id}.png", self.img_size) + gain = _load_gray(self.gain_dir / f"{map_id}_{tx_id}.png", self.img_size) + + channels = [buildings, antenna] + if self.add_distance: # v2 third channel + H = W = self.img_size + r, c = np.unravel_index(int(np.argmax(antenna)), antenna.shape) + yy, xx = np.mgrid[0:H, 0:W] + dist = np.sqrt((yy - r) ** 2 + (xx - c) ** 2).astype(np.float32) + dist /= (2 ** 0.5) * H # normalize to ~[0, 1] + channels.append(dist) + + x = np.stack(channels, axis=0) + y = gain[None, :, :] return torch.from_numpy(x), torch.from_numpy(y) diff --git a/src/train.py b/src/train.py index 15f9d94..760c7b8 100644 --- a/src/train.py +++ b/src/train.py @@ -1,15 +1,18 @@ """ train.py — train the U-Net on RadioMapSeer coverage maps. -Examples: - # smoke test: tiny subset, 1 epoch — verify the loop and measure VRAM/speed - python src\\train.py --max-samples 200 --epochs 1 --batch-size 16 +v1: 2-channel input (buildings, transmitter). +v2: add --distance for a 3rd channel (distance-to-Tx) that fixes the square cutoff. - # full run (size it after the smoke test) - python src\\train.py --epochs 30 --batch-size 16 +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 """ import argparse +import csv import time from pathlib import Path @@ -20,15 +23,17 @@ from torch.utils.data import DataLoader from dataset import RadioMapSeerDataset from model import UNet +torch.backends.cudnn.benchmark = True # autotune convs (input size is fixed) -def evaluate(model, loader, criterion, device, amp): - model.eval() + +def evaluate(net, loader, criterion, device, amp): + net.eval() total, n = 0.0, 0 with torch.no_grad(): for x, y in loader: x, y = x.to(device), y.to(device) with torch.autocast(device_type=device, dtype=torch.float16, enabled=amp): - loss = criterion(model(x), y) + loss = criterion(net(x), y) total += loss.item() * x.size(0) n += x.size(0) return total / n @@ -37,25 +42,28 @@ def evaluate(model, loader, criterion, device, amp): def main(): p = argparse.ArgumentParser() p.add_argument("--epochs", type=int, default=30) - p.add_argument("--batch-size", type=int, default=16) + p.add_argument("--batch-size", type=int, default=32) p.add_argument("--lr", type=float, default=1e-3) p.add_argument("--img-size", type=int, default=256) p.add_argument("--base", type=int, default=64) - p.add_argument("--max-samples", type=int, default=None, - help="Limit samples for quick smoke tests") + p.add_argument("--max-samples", type=int, default=None) 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") + p.add_argument("--compile", action="store_true", + help="wrap model in torch.compile (optional; flaky on Windows)") args = p.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" amp = (device == "cuda") - print(f"Device: {device} | AMP: {amp}") + print(f"Device: {device} | AMP: {amp} | distance: {args.distance} | compile: {args.compile}") train_ds = RadioMapSeerDataset(split="train", img_size=args.img_size, - max_samples=args.max_samples) + max_samples=args.max_samples, add_distance=args.distance) val_max = None if args.max_samples is None else max(20, args.max_samples // 5) val_ds = RadioMapSeerDataset(split="val", img_size=args.img_size, - max_samples=val_max) + max_samples=val_max, add_distance=args.distance) print(f"Train: {len(train_ds)} | Val: {len(val_ds)}") dl_kwargs = dict(batch_size=args.batch_size, num_workers=args.num_workers, @@ -63,7 +71,9 @@ def main(): train_loader = DataLoader(train_ds, shuffle=True, **dl_kwargs) val_loader = DataLoader(val_ds, shuffle=False, **dl_kwargs) - model = UNet(in_channels=2, out_channels=1, base=args.base).to(device) + in_ch = 3 if args.distance else 2 + model = UNet(in_channels=in_ch, out_channels=1, base=args.base).to(device) + net = torch.compile(model) if args.compile else model criterion = nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) scaler = torch.amp.GradScaler(enabled=amp) @@ -72,20 +82,19 @@ def main(): ckpt_dir.mkdir(exist_ok=True) best_val = float("inf") - import csv history_path = ckpt_dir / "history.csv" with open(history_path, "w", newline="") as f: csv.writer(f).writerow(["epoch", "train_mse", "val_mse"]) for epoch in range(1, args.epochs + 1): - model.train() + net.train() t0 = time.time() running, n = 0.0, 0 for i, (x, y) in enumerate(train_loader, 1): 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(model(x), y) + loss = criterion(net(x), y) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() @@ -95,7 +104,7 @@ def main(): print(f" epoch {epoch} step {i}/{len(train_loader)} | loss {running/n:.5f}") train_loss = running / n - val_loss = evaluate(model, val_loader, criterion, device, amp) + 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") diff --git a/src/visualize.py b/src/visualize.py index 05f42b3..5d9dd64 100644 --- a/src/visualize.py +++ b/src/visualize.py @@ -31,14 +31,17 @@ def main(): device = "cuda" if torch.cuda.is_available() else "cpu" ckpt = torch.load(args.ckpt, map_location=device, weights_only=False) - base = ckpt.get("args", {}).get("base", 64) - model = UNet(in_channels=2, out_channels=1, base=base).to(device) + cargs = ckpt.get("args", {}) + base = cargs.get("base", 64) + add_dist = bool(cargs.get("distance", False)) + in_ch = 3 if add_dist else 2 + model = UNet(in_channels=in_ch, out_channels=1, base=base).to(device) model.load_state_dict(ckpt["model_state"]) model.eval() - print(f"Loaded {args.ckpt} (epoch {ckpt.get('epoch', '?')}, " - f"val MSE {ckpt.get('val_loss', float('nan')):.5f})") + print(f"Loaded {args.ckpt} | in_channels={in_ch} | distance={add_dist} | " + f"val MSE {ckpt.get('val_loss', float('nan')):.5f}") - ds = RadioMapSeerDataset(split="test", img_size=args.img_size) # unseen cities + ds = RadioMapSeerDataset(split="test", img_size=args.img_size, add_distance=add_dist) rng = np.random.default_rng(args.seed) idxs = rng.choice(len(ds), size=args.n, replace=False)