|
| 1 | +"""Swappable-planner benchmark harness — gru vs flow_matching vs bezier. |
| 2 | +
|
| 3 | +Compares the trajectory planners registered in ``PLANNER_REGISTRY`` under |
| 4 | +IDENTICAL conditions (same embed_dim / horizon / inputs / device), as requested |
| 5 | +by @RyotaYamada in #56 (Zain to lead the Bézier-vs-Flow-Matching decision). |
| 6 | +
|
| 7 | +What this harness measures NOW (no trained checkpoint, no dataset, no simulator): |
| 8 | + * inference latency (p50 / p99 / jitter, ms) |
| 9 | + * parameter count |
| 10 | + * architectural smoothness on the (acceleration, curvature) unicycle output: |
| 11 | + - jerk proxy = Var(Δ acceleration) (lower = smoother) |
| 12 | + - curvature change = Var(Δ curvature) (lower = smoother) |
| 13 | + These hold even with random weights (Bézier's Bernstein basis is smooth by |
| 14 | + construction), so they are a fair *architectural* comparison. |
| 15 | +
|
| 16 | +What it does NOT measure yet (TODO — require a trained checkpoint + data/sim): |
| 17 | + * ADE / FDE -> needs ground-truth trajectories (KITScenes / L2D) |
| 18 | + * off-road rate, collision / -> needs map + agents + a metric/sim |
| 19 | + near-collision |
| 20 | + * closed-loop stability -> needs NAVSIM / Bench2Drive / HUGSIM |
| 21 | +
|
| 22 | +Run: |
| 23 | + env -u PYTHONPATH PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \ |
| 24 | + python Model/speed_benchmark/planner_benchmark.py |
| 25 | +""" |
| 26 | + |
| 27 | +import json |
| 28 | +import os |
| 29 | +import sys |
| 30 | +import time |
| 31 | + |
| 32 | +import torch |
| 33 | + |
| 34 | +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) |
| 35 | + |
| 36 | +from model_components.trajectory_planning import PLANNER_REGISTRY, build_planner # noqa: E402 |
| 37 | + |
| 38 | +CONFIG = dict(embed_dim=256, num_timesteps=64, num_signals=2, |
| 39 | + egomotion_dim=256, visual_history_dim=896) |
| 40 | +PLANNERS = ["gru", "flow_matching", "bezier"] |
| 41 | +WARMUP, ITERS, BATCH, H, W = 10, 50, 1, 8, 8 |
| 42 | + |
| 43 | + |
| 44 | +def _make_inputs(device): |
| 45 | + bev = torch.randn(BATCH, CONFIG["embed_dim"], H, W, device=device) |
| 46 | + vis = torch.randn(BATCH, CONFIG["visual_history_dim"], device=device) |
| 47 | + ego = torch.randn(BATCH, CONFIG["egomotion_dim"], device=device) |
| 48 | + return bev, vis, ego |
| 49 | + |
| 50 | + |
| 51 | +def _smoothness(traj): |
| 52 | + """traj [B, T*S] -> Var(Δaccel), Var(Δcurvature) on the (accel, curv) channels.""" |
| 53 | + t = traj.view(traj.shape[0], CONFIG["num_timesteps"], CONFIG["num_signals"]) |
| 54 | + accel, curv = t[..., 0], t[..., 1] |
| 55 | + jerk = (accel[:, 1:] - accel[:, :-1]).var().item() |
| 56 | + dcurv = (curv[:, 1:] - curv[:, :-1]).var().item() |
| 57 | + return float(jerk), float(dcurv) |
| 58 | + |
| 59 | + |
| 60 | +def _bench_one(name, device): |
| 61 | + torch.manual_seed(0) |
| 62 | + planner = build_planner(name, **CONFIG).to(device).eval() |
| 63 | + n_params = sum(p.numel() for p in planner.parameters()) |
| 64 | + bev, vis, ego = _make_inputs(device) |
| 65 | + with torch.no_grad(): |
| 66 | + for _ in range(WARMUP): |
| 67 | + out = planner(bev, vis, ego) |
| 68 | + if device.type == "cuda": |
| 69 | + torch.cuda.synchronize() |
| 70 | + times = [] |
| 71 | + for _ in range(ITERS): |
| 72 | + t0 = time.perf_counter() |
| 73 | + out = planner(bev, vis, ego) |
| 74 | + if device.type == "cuda": |
| 75 | + torch.cuda.synchronize() |
| 76 | + times.append((time.perf_counter() - t0) * 1000.0) |
| 77 | + jerk, dcurv = _smoothness(out[0]) |
| 78 | + times.sort() |
| 79 | + p50 = times[len(times) // 2] |
| 80 | + p99 = times[min(len(times) - 1, int(round(len(times) * 0.99)) - 1)] |
| 81 | + return { |
| 82 | + "planner": name, |
| 83 | + "params": n_params, |
| 84 | + "latency_p50_ms": round(p50, 3), |
| 85 | + "latency_p99_ms": round(p99, 3), |
| 86 | + "jitter_ms": round(p99 - p50, 3), |
| 87 | + "jerk_var_dAccel": jerk, |
| 88 | + "curv_var_dCurv": dcurv, |
| 89 | + } |
| 90 | + |
| 91 | + |
| 92 | +def main(): |
| 93 | + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 94 | + rows = [] |
| 95 | + for name in PLANNERS: |
| 96 | + if name not in PLANNER_REGISTRY: |
| 97 | + print(f"skip {name}: not in registry") |
| 98 | + continue |
| 99 | + try: |
| 100 | + rows.append(_bench_one(name, device)) |
| 101 | + except Exception as e: # noqa: BLE001 |
| 102 | + rows.append({"planner": name, "error": repr(e)}) |
| 103 | + |
| 104 | + print(f"\nDevice: {device} | torch {torch.__version__} | " |
| 105 | + f"batch={BATCH} warmup={WARMUP} iters={ITERS}\n") |
| 106 | + hdr = ["planner", "params", "latency_p50_ms", "latency_p99_ms", |
| 107 | + "jitter_ms", "jerk_var_dAccel", "curv_var_dCurv"] |
| 108 | + print("| " + " | ".join(hdr) + " |") |
| 109 | + print("|" + "|".join(["---"] * len(hdr)) + "|") |
| 110 | + for r in rows: |
| 111 | + if "error" in r: |
| 112 | + print(f"| {r['planner']} | ERROR: {r['error']} |") |
| 113 | + continue |
| 114 | + print("| " + " | ".join( |
| 115 | + f"{r[h]:.3e}" if h.startswith(("jerk", "curv")) else str(r[h]) |
| 116 | + for h in hdr) + " |") |
| 117 | + |
| 118 | + out_dir = os.path.join(os.path.dirname(__file__), "results") |
| 119 | + os.makedirs(out_dir, exist_ok=True) |
| 120 | + stamp = time.strftime("%Y%m%d_%H%M%S") |
| 121 | + out_path = os.path.join(out_dir, f"planner_benchmark_{stamp}.json") |
| 122 | + with open(out_path, "w") as f: |
| 123 | + json.dump({"device": str(device), "torch": torch.__version__, |
| 124 | + "config": CONFIG, "rows": rows}, f, indent=2) |
| 125 | + print(f"\nSaved: {out_path}") |
| 126 | + print("\nNOTE: ADE/FDE, off-road, collision and closed-loop require a trained " |
| 127 | + "checkpoint + dataset/simulator — not computed here (see module docstring).") |
| 128 | + |
| 129 | + |
| 130 | +if __name__ == "__main__": |
| 131 | + main() |
0 commit comments