| |
| """The CPU reference (``reference/``: ONNX initializers by node, BN folded where the port folds it, torch fp32) vs |
| ONNX Runtime on the deployed ONNX (PLAN.md 0.3 item 1; BUNDLE_CONVENTIONS.md section 10), on identical voxel tensors |
| (the host pre-processing of ``host/``): |
| |
| - **before the data-dependent selection** every tap (decorated PFN input, PFN layer 0, PFN, canvas, every backbone |
| conv, the deblocks, the neck concat, shared conv, heatmap head, the local-max heatmap, the input-independent key |
| position embedding) at PCC >= 0.9999 against the same tensor of the deployed graph (exposed as an extra ORT output); |
| - **the top-500 selection set-based**: >= 99 % of the set shared, and every ORT proposal with heat >= 0.1 (the |
| "strong proposals" of PLAN.md 2.5) in the reference's set (two fp32 implementations may swap low-score tail entries, |
| S:transfusion:221); |
| - **after it teacher-forced**: the reference run with ORT's own top-k indices reproduces every downstream tap |
| (query features / positions / heat, decoder layers, prediction heads, the three outputs) at PCC >= 0.9999, |
| position by position; |
| - **end to end set-based**: the same detections after Autoware's post-processing (same label, centre within 1 cm, |
| score within 1e-4) and the same counts at every stage; |
| - the plain ORT session (no extra outputs) reproduces the research's bitwise output goldens on their voxel tensors, |
| and the reference reproduces the research's expected detections of the shipped PandaSet sample. |
| |
| Cases: the shipped samples (PandaSet 019 f40, test.pcd), the Autoware sample-rosbag frame 150 (62,527 pillars: above |
| Autoware's 60,000, kept whole by N_MAX = 65,536; research data, license unstated) and the same frame through the |
| ``autoware_compat`` shuffle emulation. Runs in the research venv (onnxruntime + torch, CPU only; threads <= 4):: |
| |
| cd bundles/transfusion-p150 && OMP_NUM_THREADS=4 PYTHONPATH=code \\ |
| /home/ubuntu/experiments/tt-models/tools/research-venv/bin/python -m pytest -q -p no:cacheprovider \\ |
| code/tt_transfusion/tests/test_reference_cpu.py |
| |
| Skipped where onnxruntime or the weights are absent. |
| """ |
| from __future__ import annotations |
|
|
| import json |
|
|
| import numpy as np |
| import pytest |
|
|
| pytest.importorskip("onnxruntime") |
| pytest.importorskip("onnx") |
| torch = pytest.importorskip("torch") |
|
|
| from tt_transfusion.reference.model import SELECTION_TAPS |
| from tt_transfusion.reference.ort import ALL_TAPS |
| from tt_transfusion.reference.pipeline import TransFusionReference |
| from tt_transfusion.tests import _data |
| from tt_transfusion.ttaw.golden import TapRegistry |
| from tt_transfusion.ttaw.metrics import match_detections, pcc |
|
|
| PCC_MIN = 0.9999 |
| THREADS = 4 |
| STRONG = 0.1 |
| CASES = { |
| "pandaset_019_f40": ("sample", "pandaset_019_f40.npz", None), |
| "test_pcd": ("sample", "test_pcd.npz", None), |
| "rosbag_0150": ("research", "frame_0150", None), |
| "rosbag_0150_compat": ("research", "frame_0150", {"mode": "autoware_compat", "offset": 12345}), |
| } |
|
|
|
|
| @pytest.fixture(scope="module") |
| def references(): |
| w = _data.weights_dir() |
| torch.set_num_threads(THREADS) |
| return (TransFusionReference(w, backend="torch", threads=THREADS), |
| TransFusionReference(w, backend="ort", threads=THREADS, intermediates=True)) |
|
|
|
|
| def _source(case): |
| kind, name, sampling = CASES[case] |
| src = _data.sample_path(name) if kind == "sample" else _data.research_cloud(name) |
| return src, sampling |
|
|
|
|
| _RUNS = {} |
|
|
|
|
| def _runs(references, case): |
| if case not in _RUNS: |
| rt, ro = references |
| src, sampling = _source(case) |
| tt, to = TapRegistry(), TapRegistry() |
| a = rt.run(src, sampling=sampling, taps=tt) |
| b = ro.run(src, sampling=sampling, taps=to) |
| _RUNS.clear() |
| _RUNS[case] = (a, b, tt, to) |
| return _RUNS[case] |
|
|
|
|
| @pytest.mark.parametrize("case", list(CASES)) |
| def test_taps_before_the_selection_match_onnxruntime(references, case): |
| a, b, tt, to = _runs(references, case) |
| assert np.array_equal(a.prepared.features, b.prepared.features) |
| names = [n for n in ALL_TAPS if n not in SELECTION_TAPS] |
| assert set(names) <= set(to.names()) and set(names) <= set(tt.names()) |
| worst = {} |
| for name in names: |
| x, y = tt[name], to[name] |
| assert x.size == y.size, name |
| worst[name] = pcc(x, y.reshape(x.shape)) |
| print(case, json.dumps({k: round(v, 10) for k, v in sorted(worst.items(), key=lambda kv: kv[1])[:5]})) |
| bad = {k: v for k, v in worst.items() if not v >= PCC_MIN} |
| assert not bad, bad |
|
|
|
|
| @pytest.mark.parametrize("case", list(CASES)) |
| def test_topk_selection_set_based(references, case): |
| a, b, tt, to = _runs(references, case) |
| mine, ref = tt["head.topk_indices"].ravel(), to["head.topk_indices"].ravel() |
| k = ref.size |
| shared = len(set(mine.tolist()) & set(ref.tolist())) |
| strong = ref[to["head.topk_values"].ravel() >= STRONG] |
| print(case, f"top-{k} shared {shared}, strong {len(strong)}, positional {np.mean(mine == ref):.4f}") |
| assert shared >= 0.99 * k |
| assert set(strong.tolist()) <= set(mine.tolist()) |
| np.testing.assert_allclose(np.sort(tt["head.topk_values"].ravel()), np.sort(to["head.topk_values"].ravel()), |
| atol=1e-4) |
|
|
|
|
| @pytest.mark.parametrize("case", list(CASES)) |
| def test_teacher_forced_decoder_matches_onnxruntime(references, case): |
| rt, _ = references |
| _, b, _, to = _runs(references, case) |
| tf = TapRegistry(include=SELECTION_TAPS) |
| rt.network(b.prepared, tf, topk_indices=to["head.topk_indices"]) |
| assert np.array_equal(tf["head.topk_indices"], to["head.topk_indices"]) |
| worst = {} |
| for name in SELECTION_TAPS: |
| if name == "head.topk_indices": |
| continue |
| x, y = tf[name], to[name] |
| assert x.size == y.size, name |
| worst[name] = pcc(x, y.reshape(x.shape)) |
| print(case, json.dumps({k: round(v, 10) for k, v in sorted(worst.items(), key=lambda kv: kv[1])[:5]})) |
| bad = {k: v for k, v in worst.items() if not v >= PCC_MIN} |
| assert not bad, bad |
| assert np.array_equal(tf["query.pos"], to["query.pos"]) |
|
|
|
|
| @pytest.mark.parametrize("case", list(CASES)) |
| def test_detections_match_onnxruntime_set_based(references, case): |
| a, b, _, _ = _runs(references, case) |
| assert a.counts == b.counts |
| da, db = a.output, b.output |
| m = match_detections(da.boxes, da.label_ids, da.scores, db.boxes, db.label_ids, db.scores, max_dist=0.01) |
| assert m.recall == 1.0 and m.precision == 1.0, m.to_dict() |
| assert max(m.score_errors, default=0.0) < 1e-4 |
| assert len(da) == len(db) and len(da) > 0 |
|
|
|
|
| def test_compat_voxels_reproduce_the_research_shuffle_emulation(references): |
| """``autoware_compat`` with offset 12345 = research ``run_ort_reference.py --autoware-shuffle`` (seed 0) on frame |
| 150: the voxel tensors of ``ort_sample_0150_awshuffle/io.npz`` bit for bit, and ORT's outputs on them.""" |
| a, b, _, _ = _runs(references, "rosbag_0150_compat") |
| gold = np.load(_data.research_dir() / "ort_sample_0150_awshuffle" / "io.npz") |
| mine = a.prepared.onnx_inputs() |
| for k in ("voxels", "num_points", "coors"): |
| assert np.array_equal(mine[k], gold[k]), k |
| assert a.prepared.sampling["voxelizer_points"] == 2_000_000 and a.prepared.skipped is None |
|
|
|
|
| def test_plain_onnxruntime_reproduces_the_research_output_goldens(references): |
| """The plain session (no extra outputs) on the research voxel tensors gives the research's bitwise ORT goldens |
| (S:transfusion:221: public_data goldens and ``ort_sample_0250``).""" |
| from tt_transfusion.reference.ort import OrtTransFusion |
|
|
| rt, _ = references |
| plain = OrtTransFusion(_data.weights_dir(), threads=THREADS, weights=rt.weights) |
| r = _data.research_dir() |
| cases = [(_data.pandaset_sample(), r / "public_data" / "goldens" / "ps019_f40" / "io.npz"), |
| (_data.research_cloud("frame_0250"), r / "ort_sample_0250" / "io.npz")] |
| for src, golden in cases: |
| gold = np.load(golden) |
| prep = rt.pre(src) |
| mine = prep.onnx_inputs() |
| for k in ("voxels", "num_points", "coors"): |
| assert np.array_equal(mine[k], gold[k]), (golden, k) |
| out = plain.run(mine) |
| for k in ("cls_score0", "bbox_pred0", "dir_cls_pred0"): |
| assert np.array_equal(out[k], gold[k]), (golden, k) |
|
|
|
|
| def test_pandaset_sample_reproduces_the_research_detections(references): |
| """research/transfusion/public_data/shippable/pandaset_019_f40/expected_detections_cpu_fp32.json (ORT + the |
| research post-processing, values rounded to 4-5 decimals): the same objects, labels after the remapper, scores and |
| boxes.""" |
| from tt_transfusion.ttaw.decode import LABEL_IDS |
|
|
| a, _, _, _ = _runs(references, "pandaset_019_f40") |
| gold = json.loads((_data.research_dir() / "public_data" / "shippable" / "pandaset_019_f40" / |
| "expected_detections_cpu_fp32.json").read_text())["objects"] |
| out = a.output |
| assert len(out) == len(gold) == 80 |
| ref = np.array([[o[k] for k in ("x", "y", "z", "length", "width", "height", "yaw")] for o in gold], np.float64) |
| labels = np.array([LABEL_IDS[o["label_remapped"]] for o in gold]) |
| m = match_detections(out.boxes, out.label_ids, out.scores, ref, labels, np.array([o["score"] for o in gold]), |
| max_dist=0.01) |
| assert m.recall == 1.0 and m.precision == 1.0, m.to_dict() |
| assert max(m.score_errors, default=0.0) < 1e-4 |
| for i, j in m.pairs: |
| np.testing.assert_allclose(out.boxes[i, 2:6], ref[j, 2:6], atol=2e-4) |
|
|
