File size: 9,946 Bytes
68a3286 | 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 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 | # SPDX-License-Identifier: Apache-2.0
"""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 # noqa: E402
from tt_transfusion.reference.ort import ALL_TAPS # noqa: E402
from tt_transfusion.reference.pipeline import TransFusionReference # noqa: E402
from tt_transfusion.tests import _data # noqa: E402
from tt_transfusion.ttaw.golden import TapRegistry # noqa: E402
from tt_transfusion.ttaw.metrics import match_detections, pcc # noqa: E402
PCC_MIN = 0.9999 # PLAN.md 0.3 item 1
THREADS = 4
STRONG = 0.1 # PLAN.md 2.5: overlap of strong proposals (heat >= 0.1) 100 %
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() # every tap (module and per layer)
a = rt.run(src, sampling=sampling, taps=tt)
b = ro.run(src, sampling=sampling, taps=to)
_RUNS.clear() # keep one case in memory (~1.5 GB of taps)
_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) # identical decorated 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"]) # exact arithmetic 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)
|