#!/usr/bin/env python3 """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 # noqa: E402 from unified.model import UnifiedIMUNet, vector_huber # noqa: E402 from score_val import score_breakdown # noqa: E402 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): # (N,200,6) -> (N,6,200) normalised 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()