File size: 5,397 Bytes
4d9b003 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | # SPDX-License-Identifier: Apache-2.0
"""Where does the end-to-end error of a scene come from? (development tool; device, run under ``bin/devrun``)
bin/devrun -t 1200 -- python code/scripts/split_error.py --public scene-0103_kf14 --scenes straight_road
For each scene the ego / neighbour displacement vs the fp32 CPU reference of:
- ``device``: the full device plan;
- ``enc_only``: the device encoding (``encoder_taps``) + the CPU fp32 decoder and solver (encoder error alone);
- ``dec_only``: the CPU encoding + the device decoder (``decode_once`` replays) driven by the host solver (decoder
error alone, including its compounding over the 11 evaluations).
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
import numpy as np
HERE = Path(__file__).resolve()
sys.path.insert(0, str(HERE.parents[1]))
from tt_diffusion_planner.host import pipeline as hp # noqa: E402
from tt_diffusion_planner.host.solver import apply_prefix_constraint, dpm_solver_sample # noqa: E402
from tt_diffusion_planner.reference import config as C # noqa: E402
RES = HERE.parents[4] / "research" / "diffusion-planner"
def load_raw(scene: str, public: bool):
if public:
with np.load(RES / "public_data" / "inputs" / "nuscenes" / f"{scene}.npz", allow_pickle=False) as z:
return {k: np.array(z[f"raw/{k}"]) for k in C.INPUT_NAMES}
with np.load(RES / "goldens" / f"{scene}.npz", allow_pickle=False) as z:
return {k: np.array(z[f"in.{k}"]) for k in C.INPUT_NAMES}
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--public", nargs="*", default=[])
ap.add_argument("--scenes", nargs="*", default=[])
ap.add_argument("--json", default=None)
for opt in ("ln-fp32", "hidden-fp32", "split", "attn-fp32-acc", "attn-matmul"):
ap.add_argument(f"--{opt}", default=None, help="module globs (comma-separated); default: the knob")
args = ap.parse_args()
from tt_diffusion_planner.tt.config import globs
opts = {k: globs(getattr(args, k)) for k in ("ln_fp32", "hidden_fp32", "split", "attn_fp32_acc", "attn_matmul")
if getattr(args, k) is not None}
import torch
torch.set_num_threads(4)
from tt_diffusion_planner.device import close_device, open_device
from tt_diffusion_planner.reference.model import Decoder, Encoder, torch_params
from tt_diffusion_planner.reference.weights import find_weights_dir, load_weights
from tt_diffusion_planner.tt.model import TtDiffusionPlanner
w = load_weights(find_weights_dir())
params = {k: v[3] for k, v in hp.RUNTIME_PARAMS.items()}
P = torch_params(w.params)
enc_ref, dec_ref = Encoder(P), Decoder(P)
dev = open_device(allow_fallback=False)
report = {}
try:
tt = TtDiffusionPlanner(dev, w, debug=True, **opts)
report["options"] = tt.build.options()
tt.capture()
for scene, public in [(s, True) for s in args.public] + [(s, False) for s in args.scenes]:
prep = hp.prepare(load_raw(scene, public), w.normalization.observation)
cs = prep.decoder.current_states
def solve(model_fn):
res = dpm_solver_sample(prep.x_T, model_fn, lambda x: apply_prefix_constraint(x, cs))
return res.final_x
def cpu_plan(encoding):
kv = dec_ref.cross_kv(torch.from_numpy(np.asarray(encoding, np.float32)))
with torch.no_grad():
return solve(lambda x, t: dec_ref.forward(x, t, kv, prep.decoder.agent_valid).numpy())
def out(final_x0):
o = hp.make_output(final_x0, np.zeros(5, np.float32), prep, w.normalization, params)
return o.poses[:, :2], o.predicted_agents[..., :2]
with torch.no_grad():
enc = enc_ref.forward(prep.features).numpy()
ref_ego, ref_nb = out(cpu_plan(enc))
dev_enc = tt.encoder_taps(prep)["enc.encoding"]
cases = {
"device": tt.forward(prep)["final_x0"],
"enc_only": cpu_plan(dev_enc),
"dec_only": solve(lambda x, t: _decode(tt, prep, x, t, enc)),
}
res = {}
for name, fx in cases.items():
ego, nb = out(fx)
d = np.hypot(*(ego - ref_ego).T)
r = {"ego_max_m": round(float(d.max()), 4), "ego_mean_m": round(float(d.mean()), 4)}
if nb.shape[0]:
per_agent = np.hypot(*(nb - ref_nb).transpose(2, 0, 1)).max(1)
r["nb_median_max_m"] = round(float(np.median(per_agent)), 4)
res[name] = r
res["encoding_pcc"] = float(np.corrcoef(dev_enc[prep.features.token_valid].ravel(),
enc[prep.features.token_valid].ravel())[0, 1])
report[scene] = res
print(scene, json.dumps(res), flush=True)
tt.release()
finally:
close_device(dev)
if args.json:
Path(args.json).write_text(json.dumps(report, indent=1) + "\n")
return 0
def _decode(tt, prep, x, t, enc):
out = tt.decode_once(prep, x, float(t), encoding=enc)
out[:, 0] = 0.0 # the t = 0 slot (masked on the device) is overwritten by the prefix constraint anyway
return out
if __name__ == "__main__":
sys.exit(main())
|