| |
| """Train UnifiedIMUNet locally and score every epoch with the OFFICIAL scorer. |
| |
| Faithful to sample_code/ts-tcmpio-unified-temporal-resnet.ipynb (same |
| architecture, loss, optimiser, schedule, sampler weights and seed) with two |
| deliberate changes: |
| |
| * the 0.39 GB train split lives on the GPU and batches are drawn with |
| torch.multinomial, which is what WeightedRandomSampler(replacement=True) |
| does -- same distribution, no dataloader overhead; |
| * every epoch also reports the real TartanIMU Score on val via |
| local_eval/score_val.py, not just macro AVE, so we can see whether the |
| notebook's model-selection signal (AVE) tracks the thing being graded. |
| |
| Checkpoint selection follows the notebook (best macro AVE); --select score |
| switches it to the competition score instead. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import torch.nn.functional as F |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path[:0] = [str(ROOT), str(ROOT / "local_eval")] |
| from unified.data import PLATFORMS, build_windows, load_index, sample_weights |
| from unified.model import UnifiedIMUNet, vector_huber |
| from score_val import score_breakdown |
|
|
|
|
| def macro_ave(pred: np.ndarray, target: np.ndarray, meta: pd.DataFrame) -> float: |
| s = meta[["platform", "traj_id"]].copy() |
| s["error"] = np.linalg.norm(pred - target, axis=1) |
| return float(s.groupby(["platform", "traj_id"])["error"].mean().groupby("platform").mean().mean()) |
|
|
|
|
| @torch.inference_mode() |
| def predict(model, X: torch.Tensor, batch: int = 1024) -> np.ndarray: |
| model.eval() |
| out = [] |
| for i in range(0, len(X), batch): |
| with torch.autocast("cuda", dtype=torch.float16, enabled=X.is_cuda): |
| v, _ = model(X[i:i + batch]) |
| out.append(v.float().cpu().numpy()) |
| return np.concatenate(out) |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--epochs", type=int, default=25) |
| ap.add_argument("--patience", type=int, default=6) |
| ap.add_argument("--batch_size", type=int, default=512) |
| ap.add_argument("--lr", type=float, default=2e-3) |
| ap.add_argument("--weight_decay", type=float, default=1e-4) |
| ap.add_argument("--platform_loss_weight", type=float, default=0.05) |
| ap.add_argument("--seed", type=int, default=42) |
| ap.add_argument("--select", choices=["ave", "score"], default="ave") |
| ap.add_argument("--sampling", choices=["balanced", "uniform", "drone2x"], default="balanced", |
| help="balanced: notebook's platform-equal weights; uniform: raw window " |
| "frequency (drone drops to 17%%); drone2x: balanced with drone doubled") |
| ap.add_argument("--tag", default="run") |
| ap.add_argument("--ema", type=float, default=0.0, |
| help="EMA decay for weight averaging (0 = off). Produces ONE set of " |
| "weights, unlike a multi-seed ensemble, so it stays inside the " |
| "'single model with one shared set of weights' rule.") |
| ap.add_argument("--solution", default="local_eval/val_solution.csv") |
| a = ap.parse_args() |
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| torch.manual_seed(a.seed) |
| torch.cuda.manual_seed_all(a.seed) |
| torch.backends.cudnn.benchmark = True |
|
|
| t0 = time.time() |
| idx = load_index() |
| train, val, test = idx["train"], idx["val"], idx["test"] |
| Xtr, Xva, Xte = (build_windows(f) for f in (train, val, test)) |
| print(f"windows: train {len(Xtr):,} val {len(Xva):,} test {len(Xte):,}" |
| f" ({time.time() - t0:.0f}s to load)") |
|
|
| mean = Xtr.mean(axis=(0, 1), dtype=np.float64).astype(np.float32) |
| std = Xtr.std(axis=(0, 1), dtype=np.float64).astype(np.float32) |
|
|
| def to_gpu(X): |
| t = torch.from_numpy((X - mean) / std).to(device) |
| return t.permute(0, 2, 1).contiguous() |
|
|
| Xtr_g, Xva_g, Xte_g = to_gpu(Xtr), to_gpu(Xva), to_gpu(Xte) |
| ytr = torch.from_numpy(train[["vx", "vy", "vz"]].to_numpy(np.float32)).to(device) |
| ptr = torch.from_numpy(train["platform_id"].to_numpy(np.int64)).to(device) |
| yva = val[["vx", "vy", "vz"]].to_numpy(np.float32) |
| if a.sampling == "uniform": |
| wnp = np.ones(len(train), dtype=np.float64) |
| else: |
| wnp = sample_weights(train) |
| if a.sampling == "drone2x": |
| wnp = wnp * np.where(train["platform"].to_numpy() == "drone", 2.0, 1.0) |
| w = torch.as_tensor(wnp, dtype=torch.double, device=device) |
| share = (pd.Series(wnp).groupby(train["platform"]).sum() / wnp.sum() * 100).round(1) |
| print("抽樣佔比 %:", share.to_dict()) |
|
|
| solution = pd.read_csv(a.solution) |
| model = UnifiedIMUNet().to(device) |
| n_par = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| steps = len(Xtr_g) // a.batch_size |
| opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=a.weight_decay) |
| sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=a.lr, epochs=a.epochs, |
| steps_per_epoch=steps, pct_start=0.1) |
| scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda") |
|
|
| ema_model = None |
| if a.ema > 0: |
| from torch.optim.swa_utils import AveragedModel, get_ema_multi_avg_fn |
| ema_model = AveragedModel(model, multi_avg_fn=get_ema_multi_avg_fn(a.ema), |
| use_buffers=True) |
|
|
| ckpt_path = ROOT / f"unified/ckpt_{a.tag}.pt" |
| print(f"{n_par/1e6:.2f}M params | {steps} steps/epoch | select on {a.select} | {device}\n") |
| history, best, waiting = [], float("inf"), 0 |
|
|
| for epoch in range(1, a.epochs + 1): |
| model.train() |
| order = torch.multinomial(w, steps * a.batch_size, replacement=True) |
| total = 0.0 |
| for i in range(steps): |
| sel = order[i * a.batch_size:(i + 1) * a.batch_size] |
| opt.zero_grad(set_to_none=True) |
| with torch.autocast("cuda", dtype=torch.float16, enabled=device.type == "cuda"): |
| v, logits = model(Xtr_g[sel]) |
| loss = vector_huber(v, ytr[sel]) + a.platform_loss_weight * F.cross_entropy(logits, ptr[sel]) |
| scaler.scale(loss).backward() |
| scaler.unscale_(opt) |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| scaler.step(opt) |
| scaler.update() |
| sched.step() |
| total += loss.item() |
|
|
| if ema_model is not None: |
| ema_model.update_parameters(model) |
| pred = predict(model, Xva_g) |
| ave = macro_ave(pred, yva, val) |
| sub = pd.DataFrame({"window_id": val["window_id"], "vx": pred[:, 0], |
| "vy": pred[:, 1], "vz": pred[:, 2]}) |
| score, table = score_breakdown(solution, sub) |
| per = {r.platform: r.score for r in table.itertuples()} |
| history.append({"epoch": epoch, "loss": total / steps, "macro_ave": ave, "score": score, **per}) |
|
|
| watched = ave if a.select == "ave" else score |
| star = "" |
| if watched < best - 1e-4: |
| best, waiting, star = watched, 0, " *" |
| torch.save({"model": model.state_dict(), "imu_mean": mean, "imu_std": std, |
| "epoch": epoch, "macro_ave": ave, "score": score}, ckpt_path) |
| else: |
| waiting += 1 |
| line = (f"ep{epoch:3d} loss {total/steps:.4f} AVE {ave:.4f} Score {score:.4f} " |
| + " ".join(f"{p} {per[p]:.3f}" for p in PLATFORMS)) |
| if ema_model is not None: |
| ep_pred = predict(ema_model.module, Xva_g) |
| ep_sub = pd.DataFrame({"window_id": val["window_id"], "vx": ep_pred[:, 0], |
| "vy": ep_pred[:, 1], "vz": ep_pred[:, 2]}) |
| ema_score, _ = score_breakdown(solution, ep_sub) |
| history[-1]["ema_score"] = ema_score |
| line += f" EMA {ema_score:.4f}" |
| print(line + star) |
| if waiting >= a.patience: |
| print(f"early stop (patience {a.patience})") |
| break |
|
|
| hist = pd.DataFrame(history) |
| hist.to_csv(ROOT / f"unified/history_{a.tag}.csv", index=False) |
| ck = torch.load(ckpt_path, map_location=device, weights_only=False) |
| model.load_state_dict(ck["model"]) |
| print(f"\nbest epoch {ck['epoch']}: AVE {ck['macro_ave']:.4f} Score {ck['score']:.4f}") |
|
|
| variants = [("", model)] |
| if ema_model is not None: |
| torch.save({"model": ema_model.module.state_dict(), "imu_mean": mean, "imu_std": std, |
| "ema": a.ema}, ROOT / f"unified/ckpt_{a.tag}_ema.pt") |
| variants.append(("_ema", ema_model.module)) |
| for suffix, net in variants: |
| for name, Xg, frame in (("val", Xva_g, val), ("test", Xte_g, test)): |
| p = predict(net, Xg) |
| out = ROOT / f"local_eval/sub_{name}_{a.tag}{suffix}.csv" |
| pd.DataFrame({"window_id": frame["window_id"], "vx": p[:, 0], "vy": p[:, 1], |
| "vz": p[:, 2]}).to_csv(out, index=False) |
| print(f"wrote {out.relative_to(ROOT)} rows={len(p)}") |
| print(f"total {time.time() - t0:.0f}s") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
