| |
| |
| """CPU-reference goldens of transfusion-p150 (research venv: torch + onnxruntime, CPU only, <= 4 threads). |
| |
| cd bundles/transfusion-p150 && OMP_NUM_THREADS=4 PYTHONPATH=code \\ |
| /home/ubuntu/experiments/tt-models/tools/research-venv/bin/python -I code/scripts/make_goldens.py \\ |
| --weights /home/ubuntu/experiments/tt-models/assets/transfusion/hf_lidar_transfusion \\ |
| --research /home/ubuntu/experiments/tt-models/research/transfusion \\ |
| --out /home/ubuntu/experiments/tt-models/research/transfusion/goldens |
| |
| For every frame the fp32 torch reference (``reference/pipeline.py``: host pre-processing, the deployed network rebuilt |
| from the ONNX initializers, host post-processing) runs once; ONNX Runtime runs the deployed ONNX on the same voxel |
| tensors twice (plain: the bitwise output golden; with the module intermediates exposed: the agreement check), and the |
| reference runs once more teacher-forced with ORT's top-k indices. The agreement (PCC per module tap before the |
| selection, top-500 set overlap and strong-proposal coverage, teacher-forced PCC after it, detection match) is recorded |
| with the golden. Written: |
| |
| - ``<out>/<frame>.io.npz``: the ONNX inputs (``voxels``, ``num_points``, ``coors``), the decorated features of the |
| kept pillars (``features``), the reference outputs (``cls_score0``, ``bbox_pred0``, ``dir_cls_pred0``) and the |
| plain ORT outputs (``ort_cls_score0``, ...; float32); |
| - ``<out>/<frame>.detections.json``: the reference output (``Detections3D.to_dict``), per-stage counts, pillar and |
| sampling statistics, the ORT agreement; |
| - ``<out>/<frame>.taps.npz`` (``--taps`` frames only; ~100-200 MB each): every tap of ``reference.model.MODULE_TAPS`` |
| plus ``pfn.layer0`` and ``decoder.key_pos_embed`` (float32; ``head.topk_indices`` int64): the per-module goldens of |
| the device PCC gates, including the indices that teacher-force the device decoder; |
| - ``<out>/manifest.json``: sha256 of every file, tool versions, inputs; |
| - in the bundle (small, shippable: Apache-2.0 / CC BY 4.0 inputs only): ``tests/goldens/voxels.json`` (voxel tensor |
| hashes of the samples), ``tests/goldens/detections/<sample>.json`` and the smoke reference |
| ``samples/<stem>.reference.json``. The PandaSet sample is not shipped until the user approves PandaSet samples |
| (research/DATASETS.md N5): its input is read from, and its small golden and reference are written to, the |
| git-ignored ``staging_samples_pandaset/`` of the bundle (``goldens/`` there) unless it already sits in ``samples/``. |
| |
| ``--set selection`` (fix round 1) adds ``<out>/<frame>.selection.npz`` to every development frame above: the |
| reference's top-500 (``head.topk_indices`` int64, ``head.topk_values``) from a fresh reference run, whose outputs must |
| equal the stored ``io.npz`` bit for bit (nothing else is rewritten): the proposal identities of the flip-free |
| proposal-level agreement (``tests/_agreement.py``). |
| |
| ``--set heldout`` (fix round 1) writes compact goldens of the HELD-OUT frames (``tests/_data.py HELDOUT``: 34 public |
| frames never used for development or precision choices, converted by the research converter into |
| ``research/transfusion/heldout/inputs``) to ``<out>/heldout/``: ``<frame>.io.npz`` (reference outputs, top-500 |
| indices / values, plain ORT outputs) and ``<frame>.detections.json`` (output, counts, proposals and the same ORT |
| agreement record as the development goldens: pre-selection PCC, top-k sets, teacher-forced PCC, detections), plus |
| ``heldout/manifest.json``. No module taps are stored. |
| |
| Frames: the shipped samples (``pandaset_019_f40``, ``test_pcd``), the 16 public frames of |
| research/transfusion/public_data (PandaSet CC BY 4.0; nuScenes / Argoverse 2 non-commercial: research only) and the |
| Autoware sample-rosbag frames 50 / 150 / 250 (2 sweeps; frame 150 has 62,527 pillars; license unstated: research |
| only), plus frame 150 through the ``autoware_compat`` shuffle emulation (offset 12345, the research's). |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import platform |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| HERE = Path(__file__).resolve().parent |
| CODE = HERE.parent |
| sys.path.insert(0, str(CODE)) |
|
|
| from tt_transfusion import __version__ |
| from tt_transfusion import io as tio |
| from tt_transfusion.reference.model import MODULE_TAPS, SELECTION_TAPS |
| from tt_transfusion.reference.ort import OrtTransFusion |
| from tt_transfusion.reference.pipeline import TransFusionReference |
| from tt_transfusion.ttaw import __version__ as ttaw_version |
| from tt_transfusion.ttaw.golden import TapRegistry, save_goldens |
| from tt_transfusion.ttaw.metrics import match_detections, pcc |
|
|
| SAMPLES = CODE / "tt_transfusion" / "samples" |
| BUNDLE_GOLDENS = CODE / "tt_transfusion" / "tests" / "goldens" |
| PANDASET_STAGING = CODE.parent / "staging_samples_pandaset" |
| SHIPPED = ("pandaset_019_f40", "test_pcd") |
|
|
|
|
| def sample_home(name: str) -> Path: |
| """Where a sample lives: ``samples/`` (shipped), or the PandaSet staging copy until it ships.""" |
| if name.startswith("pandaset") and not (SAMPLES / f"{name}.npz").is_file(): |
| return PANDASET_STAGING |
| return SAMPLES |
|
|
|
|
| def small_golden_home(name: str) -> Path: |
| return PANDASET_STAGING / "goldens" if sample_home(name) == PANDASET_STAGING else BUNDLE_GOLDENS / "detections" |
| PUBLIC = ("ps019_f10", "ps019_f40", "ps019_f70", "ps005_f40", "ps090_f10", "ps090_f40", "ps090_f70", "ps029_f40", |
| "ns0103_k10", "ns0103_k20", "ns0103_k30", "ns0916_k10", "ns0916_k25", "av2_280269f9_s080", |
| "av2_51bbdd4d_s080", "av2_e13c06cb_s080") |
| ROSBAG = ("frame_0050", "frame_0150", "frame_0250") |
| FIELDS5 = ("x", "y", "z", "intensity", "time_lag") |
| COMPAT = {"mode": "autoware_compat", "offset": 12345} |
| TAP_SET = tuple(MODULE_TAPS) + ("pfn.layer0", "decoder.key_pos_embed") |
| STRONG = 0.1 |
|
|
|
|
| def sha(a: np.ndarray) -> str: |
| return hashlib.sha256(np.ascontiguousarray(a).tobytes()).hexdigest() |
|
|
|
|
| def file_sha(p: Path) -> str: |
| return hashlib.sha256(p.read_bytes()).hexdigest() |
|
|
|
|
| def frames(research: Path): |
| """(name, source, sampling, license, shippable).""" |
| for name in SHIPPED: |
| lic = "CC BY 4.0 + PandaSet Dataset Terms" if name.startswith("pandaset") else "Apache-2.0 (autoware_universe)" |
| yield name, sample_home(name) / f"{name}.npz", None, lic, True |
| for name in PUBLIC: |
| path = research / "public_data" / "inputs" / f"{name}.npy" |
| if not path.is_file(): |
| print(f"skip {name}: {path} missing") |
| continue |
| meta = json.loads(path.with_suffix(".json").read_text()) |
| lic = meta.get("license") or ("CC BY 4.0 + PandaSet Dataset Terms" if name.startswith("ps") else |
| "CC BY-NC-SA 4.0 (research only)") |
| yield f"public_{name}", tio.PointCloud(np.load(path, allow_pickle=False), FIELDS5), None, lic, False |
| for name in ROSBAG: |
| path = research / "sample_frames" / f"{name}_sweeps2.npy" |
| if not path.is_file(): |
| print(f"skip {name}: {path} missing") |
| continue |
| cloud = tio.PointCloud(np.load(path, allow_pickle=False), FIELDS5) |
| lic = "Autoware sample rosbag (license unstated): research only" |
| yield f"rosbag_{name[-4:]}", cloud, None, lic, False |
| if name == "frame_0150": |
| yield "rosbag_0150_compat", cloud, COMPAT, lic, False |
|
|
|
|
| def agreement(rt, ro_tapped, run, run_o, tt: TapRegistry, to: TapRegistry) -> dict: |
| """Pre-selection PCC per tap, top-k set agreement, teacher-forced PCC after the selection, detection match.""" |
| pre = {k: float(pcc(tt[k], to[k].reshape(tt[k].shape))) for k in to.names() |
| if k in tt and k not in SELECTION_TAPS} |
| mine, ref = tt["head.topk_indices"].ravel(), to["head.topk_indices"].ravel() |
| strong = ref[to["head.topk_values"].ravel() >= STRONG] |
| tf = TapRegistry(include=SELECTION_TAPS) |
| rt.network(run.prepared, tf, topk_indices=to["head.topk_indices"]) |
| post = {k: float(pcc(tf[k], to[k].reshape(tf[k].shape))) for k in SELECTION_TAPS |
| if k != "head.topk_indices" and k in to} |
| a, b = run.output, run_o.output |
| m = match_detections(a.boxes, a.label_ids, a.scores, b.boxes, b.label_ids, b.scores, max_dist=0.01) |
| return {"pre_selection_min_pcc": min(pre.values()), "pre_selection": pre, |
| "topk": {"shared": int(len(set(mine.tolist()) & set(ref.tolist()))), "k": int(ref.size), |
| "positional_equal": float(np.mean(mine == ref)), "strong": int(len(strong)), |
| "strong_covered": bool(set(strong.tolist()) <= set(mine.tolist()))}, |
| "teacher_forced_min_pcc": min(post.values()), "teacher_forced": post, |
| "detections": {"counts_equal": run.counts == run_o.counts, "ort_counts": run_o.counts, |
| "recall": m.recall, "precision": m.precision, |
| "max_score_err": max(m.score_errors, default=0.0)}} |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) |
| ap.add_argument("--weights", type=Path, required=True) |
| ap.add_argument("--research", type=Path, required=True) |
| ap.add_argument("--out", type=Path, required=True) |
| ap.add_argument("--taps", default="pandaset_019_f40,public_ns0103_k20,rosbag_0150", |
| help="frames whose module taps are saved") |
| ap.add_argument("--threads", type=int, default=4) |
| ap.add_argument("--only", default="", help="comma list of frame names (default: all)") |
| ap.add_argument("--set", default="dev", choices=("dev", "selection", "heldout"), |
| help="dev: the development goldens (default); selection: add <frame>.selection.npz to them; " |
| "heldout: the compact held-out goldens in <out>/heldout") |
| a = ap.parse_args() |
| if a.set == "selection": |
| return selection_main(a) |
| if a.set == "heldout": |
| return heldout_main(a) |
| import onnxruntime |
| import torch |
|
|
| torch.set_num_threads(a.threads) |
| taps_for = set(filter(None, a.taps.split(","))) |
| only = set(filter(None, a.only.split(","))) |
| a.out.mkdir(parents=True, exist_ok=True) |
| (BUNDLE_GOLDENS / "detections").mkdir(parents=True, exist_ok=True) |
| manifest_path = a.out / "manifest.json" |
| manifest = json.loads(manifest_path.read_text()) if manifest_path.is_file() else {"files": {}} |
| rt = TransFusionReference(a.weights, backend="torch", threads=a.threads) |
| ro = TransFusionReference(a.weights, backend="ort", threads=a.threads, intermediates=TAP_SET) |
| plain = OrtTransFusion(a.weights, threads=a.threads, weights=rt.weights) |
| manifest.update({"bundle_version": __version__, "ttaw": ttaw_version, "torch": torch.__version__, |
| "onnxruntime": onnxruntime.__version__, "numpy": np.__version__, |
| "python": platform.python_version(), "onnx_sha256": rt.weights.sha256, |
| "weights": "AutowareFoundation/lidar_transfusion@8f07eb93 (v2.1)", |
| "config": rt.config.describe(), "written": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), |
| "command": " ".join(sys.argv)}) |
| voxels = {} |
| if (BUNDLE_GOLDENS / "voxels.json").is_file(): |
| voxels = json.loads((BUNDLE_GOLDENS / "voxels.json").read_text()) |
| for name, src, sampling, lic, shippable in frames(a.research): |
| if only and name not in only: |
| continue |
| t0 = time.perf_counter() |
| tt, to = TapRegistry(include=TAP_SET), TapRegistry(include=TAP_SET) |
| run = rt.run(src, sampling=sampling, taps=tt) |
| run_o = ro.run(src, sampling=sampling, taps=to) |
| p = run.prepared |
| n = p.num_pillars |
| assert np.array_equal(p.features, run_o.prepared.features), name |
| inputs = p.onnx_inputs() |
| ort_out = plain.run(inputs) |
| agree = agreement(rt, ro, run, run_o, tt, to) |
| t1 = time.perf_counter() |
| io_path = a.out / f"{name}.io.npz" |
| save_goldens(io_path, {**inputs, "features": p.features[:n], **run.outputs, |
| **{f"ort_{k}": v for k, v in ort_out.items()}}, |
| meta={"frame": name, "num_pillars": n, "license": lic, "shippable": shippable, |
| "sampling": p.sampling, "points_sha256": sha(p.points)}, compress=True) |
| files = [io_path] |
| if name in taps_for: |
| taps_path = a.out / f"{name}.taps.npz" |
| tt.save(taps_path, meta={"frame": name, "taps": tt.names(), "sampling": p.sampling}, compress=True) |
| files.append(taps_path) |
| body = run.output.to_dict() |
| rec = {"frame": name, "license": lic, "shippable": shippable, |
| "reference": "tt_transfusion.reference torch fp32 (ONNX initializers, BN folded)", |
| "counts": run.counts, "pillars": {k: v for k, v in p.stats().items() if k != "densify"}, |
| "densify": p.densify, "ort_agreement": agree, |
| "proposals": [int(v) for v in run.objects.extra["proposal"]], "output": body} |
| det_path = a.out / f"{name}.detections.json" |
| det_path.write_text(json.dumps(rec, indent=1)) |
| files.append(det_path) |
| for f in files: |
| manifest["files"][f.name] = {"sha256": file_sha(f), "bytes": f.stat().st_size} |
| if name in SHIPPED: |
| small = {k: rec[k] for k in ("frame", "license", "counts", "pillars", "proposals", "output")} |
| small["ort_agreement"] = {"pre_selection_min_pcc": agree["pre_selection_min_pcc"], |
| "topk": agree["topk"], "teacher_forced_min_pcc": agree["teacher_forced_min_pcc"], |
| "detections": agree["detections"]} |
| small_golden_home(name).mkdir(parents=True, exist_ok=True) |
| (small_golden_home(name) / f"{name}.json").write_text(json.dumps(small, indent=1)) |
| (sample_home(name) / f"{name}.reference.json").write_text(json.dumps(body, indent=1)) |
| voxels[f"{name}.npz"] = {"num_pillars": n, "num_pillars_total": p.pillars.num_pillars_total, |
| "points_sha256": sha(p.points), "features_sha256": sha(p.features[:n]), |
| "coords_sha256": sha(p.coords[:n]), |
| "num_points_sha256": sha(p.pillars.num_points[:n]), |
| "canvas_index_sha256": sha(p.canvas_index)} |
| print(f"{name}: pillars {n}/{p.pillars.num_pillars_total} counts {run.counts} | ort pre-sel min pcc " |
| f"{agree['pre_selection_min_pcc']:.10f}, top-k shared {agree['topk']['shared']}/{agree['topk']['k']} " |
| f"strong {agree['topk']['strong']} covered {agree['topk']['strong_covered']}, teacher-forced min pcc " |
| f"{agree['teacher_forced_min_pcc']:.10f}, detections recall/precision " |
| f"{agree['detections']['recall']:.3f}/{agree['detections']['precision']:.3f} " |
| f"({t1 - t0:.1f} s + {time.perf_counter() - t1:.1f} s io)", flush=True) |
| manifest_path.write_text(json.dumps(manifest, indent=1)) |
| if voxels: |
| (BUNDLE_GOLDENS / "voxels.json").write_text(json.dumps(voxels, indent=1)) |
| return 0 |
|
|
|
|
| SEL_TAPS = ("head.topk_indices", "head.topk_values") |
|
|
|
|
| def _manifest(path: Path, a, rt, extra: dict) -> dict: |
| import onnxruntime |
| import torch |
|
|
| m = json.loads(path.read_text()) if path.is_file() else {"files": {}} |
| m.setdefault("files", {}) |
| m.update({"bundle_version": __version__, "ttaw": ttaw_version, "torch": torch.__version__, |
| "onnxruntime": onnxruntime.__version__, "numpy": np.__version__, "python": platform.python_version(), |
| "onnx_sha256": rt.weights.sha256, "weights": "AutowareFoundation/lidar_transfusion@8f07eb93 (v2.1)", |
| "config": rt.config.describe(), "written": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), |
| "command": " ".join(sys.argv), **extra}) |
| return m |
|
|
|
|
| def selection_main(a) -> int: |
| """``--set selection``: ``<frame>.selection.npz`` (the reference top-500) next to every development golden; the |
| fresh reference outputs must equal the stored ones bit for bit.""" |
| import torch |
|
|
| torch.set_num_threads(a.threads) |
| only = set(filter(None, a.only.split(","))) |
| rt = TransFusionReference(a.weights, backend="torch", threads=a.threads) |
| manifest_path = a.out / "manifest.json" |
| manifest = json.loads(manifest_path.read_text()) |
| for name, src, sampling, _lic, _ship in frames(a.research): |
| if only and name not in only: |
| continue |
| io_path = a.out / f"{name}.io.npz" |
| if not io_path.is_file(): |
| print(f"skip {name}: {io_path} missing") |
| continue |
| taps = TapRegistry(include=SEL_TAPS) |
| run = rt.run(src, sampling=sampling, taps=taps) |
| with np.load(io_path, allow_pickle=False) as z: |
| for k, v in run.outputs.items(): |
| if not np.array_equal(np.asarray(v), z[k]): |
| raise SystemExit(f"{name}: the reference output {k} differs from the stored golden") |
| path = a.out / f"{name}.selection.npz" |
| save_goldens(path, {k: taps[k] for k in SEL_TAPS}, meta={"frame": name, "io_sha256": file_sha(io_path)}, |
| compress=True) |
| manifest["files"][path.name] = {"sha256": file_sha(path), "bytes": path.stat().st_size} |
| print(f"{name}: selection.npz written (strong {int((taps['head.topk_values'] >= STRONG).sum())}), outputs " |
| f"== stored golden", flush=True) |
| manifest["selection_written"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) |
| manifest_path.write_text(json.dumps(manifest, indent=1)) |
| return 0 |
|
|
|
|
| def heldout_main(a) -> int: |
| """``--set heldout``: compact goldens of the held-out frames (module docstring).""" |
| import torch |
|
|
| sys.path.insert(0, str(CODE)) |
| from tt_transfusion.tests._data import HELDOUT |
|
|
| torch.set_num_threads(a.threads) |
| only = set(filter(None, a.only.split(","))) |
| out = a.out / "heldout" |
| out.mkdir(parents=True, exist_ok=True) |
| rt = TransFusionReference(a.weights, backend="torch", threads=a.threads) |
| ro = TransFusionReference(a.weights, backend="ort", threads=a.threads, intermediates=TAP_SET) |
| plain = OrtTransFusion(a.weights, threads=a.threads, weights=rt.weights) |
| inputs_dir = a.research / "heldout" / "inputs" |
| manifest_path = out / "manifest.json" |
| manifest = _manifest(manifest_path, a, rt, {"set": "heldout", "inputs": str(inputs_dir)}) |
| for fid in HELDOUT: |
| name = f"heldout_{fid}" |
| if only and name not in only and fid not in only: |
| continue |
| path = inputs_dir / f"{fid}.npy" |
| if not path.is_file(): |
| print(f"skip {fid}: {path} missing") |
| continue |
| t0 = time.perf_counter() |
| cloud = tio.PointCloud(np.load(path, allow_pickle=False), FIELDS5) |
| taps, to = TapRegistry(include=TAP_SET), TapRegistry(include=TAP_SET) |
| run = rt.run(cloud, taps=taps) |
| run_o = ro.run(cloud, taps=to) |
| p = run.prepared |
| assert np.array_equal(p.features, run_o.prepared.features), name |
| inputs = p.onnx_inputs() |
| ort_out = plain.run(inputs) |
| agree = agreement(rt, ro, run, run_o, taps, to) |
| m = agree["detections"] |
| io_path = out / f"{name}.io.npz" |
| save_goldens(io_path, {**run.outputs, **{k: taps[k] for k in SEL_TAPS}, |
| **{f"ort_{k}": v for k, v in ort_out.items()}}, |
| meta={"frame": name, "num_pillars": p.num_pillars, "points_sha256": sha(p.points), |
| "input_sha256": file_sha(path)}, compress=True) |
| meta = json.loads(path.with_suffix(".json").read_text()) |
| lic = meta.get("license") or ("CC BY 4.0 + PandaSet Dataset Terms" if fid.startswith("ps") else |
| "CC BY-NC-SA 4.0 (research only)") |
| rec = {"frame": name, "license": lic, "shippable": False, "set": "heldout", |
| "reference": "tt_transfusion.reference torch fp32 (ONNX initializers, BN folded)", |
| "counts": run.counts, "pillars": {k: v for k, v in p.stats().items() if k != "densify"}, |
| "ort_agreement": agree, "proposals": [int(v) for v in run.objects.extra["proposal"]], |
| "output": run.output.to_dict()} |
| det_path = out / f"{name}.detections.json" |
| det_path.write_text(json.dumps(rec, indent=1)) |
| for f in (io_path, det_path): |
| manifest["files"][f.name] = {"sha256": file_sha(f), "bytes": f.stat().st_size} |
| manifest_path.write_text(json.dumps(manifest, indent=1)) |
| print(f"{name}: pillars {p.num_pillars} counts {run.counts} | ort pre-sel min pcc " |
| f"{agree['pre_selection_min_pcc']:.10f}, top-k shared {agree['topk']['shared']}/{agree['topk']['k']} " |
| f"strong {agree['topk']['strong']} covered {agree['topk']['strong_covered']}, teacher-forced min pcc " |
| f"{agree['teacher_forced_min_pcc']:.10f}, detections recall/precision {m['recall']:.3f}/" |
| f"{m['precision']:.3f} counts equal {m['counts_equal']} ({time.perf_counter() - t0:.1f} s)", flush=True) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
