File size: 8,531 Bytes
ac1609a | 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 | # SPDX-License-Identifier: Apache-2.0
"""The CPU fp32 reference (``tt_streampetr.reference``) against ONNX Runtime on the three deployed ONNX graphs
(PLAN.md 0.3 item 1; ORT runs them as-is, S:6).
Gates:
- before any data-dependent selection: PCC >= 0.9999 (``img_feats``, ``pos_embed`` / ``cone``, every head output
computed before the top-256: cls, bbox, ``outs_dec``, the decoder inputs);
- after it (the top-256 selection, the post-update memory rows, the detections): set-based or teacher-forced -- the
top-256 overlap, the post memory with ORT's own selection (recovered from its ``post_memory_embedding``), and the
detection F1 of the end-to-end runs (same label, 1 m).
Runs in the research venv (onnxruntime); skipped where onnxruntime, torch or the weights are absent::
cd bundles/streampetr-p150 && OMP_NUM_THREADS=3 PYTHONPATH=$PWD/code \\
/home/ubuntu/experiments/tt-models/tools/research-venv/bin/python -m pytest -q -s -p no:cacheprovider \\
code/tt_streampetr/tests/test_reference_cpu.py
Inputs: the staged PandaSet sample (frames 40-43 through the full node pipeline) and the research
teacher-forcing records of the public sequences (``research/streampetr/public_data``; nuScenes is internal only).
"""
from __future__ import annotations
import os
import numpy as np
import pytest
ort = pytest.importorskip("onnxruntime")
torch = pytest.importorskip("torch")
pytest.importorskip("onnx")
from tt_streampetr.host import inputs as hin # noqa: E402
from tt_streampetr.host import postprocess as hp # noqa: E402
from tt_streampetr.host.agreement import detection_f1, recover_topk_indices, topk_overlap # noqa: E402
from tt_streampetr.host.temporal import StreamState # noqa: E402
from tt_streampetr.reference.config import default_config # noqa: E402
from tt_streampetr.tests import paths # noqa: E402
from tt_streampetr.ttaw.metrics import pcc # noqa: E402
REF_PCC = 0.9999 # PLAN.md 0.3 item 1
CFG = default_config()
THREADS = int(os.environ.get("STREAMPETR_TEST_THREADS", "3")) # workspace rule: <= 4 CPU threads
PRE_SELECTION = ("all_cls_scores", "all_bbox_preds", "outs_dec", "reference_points", "tgt", "query_pos",
"temp_memory", "temp_pos", "query_pos_in")
POST_MEMORY = ("post_memory_embedding", "post_memory_reference_point", "post_memory_egopose", "post_memory_velo")
FRAMES = (40, 41, 42, 43)
@pytest.fixture(scope="module")
def weights():
w = paths.weights_dir()
if w is None:
pytest.skip("camera_streampetr v1.0 weights not found (set STREAMPETR_WEIGHTS_DIR)")
return w
@pytest.fixture(scope="module")
def oracle(weights):
from tt_streampetr.reference.ort import OrtOracle
return OrtOracle(weights, threads=THREADS)
@pytest.fixture(scope="module")
def reference(weights):
from tt_streampetr.reference.pipeline import StreamPETRReference
return StreamPETRReference(weights, CFG, threads=THREADS, preprocess_workers=2)
@pytest.fixture(scope="module")
def sample_run(reference, oracle):
"""The staged PandaSet sample (019 frames 40-43, one stream from a cold start) through the reference pipeline
and, with the same host inputs and the same host memory loop, through ORT."""
from tt_streampetr.ttaw.golden import TapRegistry
paths.need_pandaset()
reference.streams.clear()
st_ort = StreamState(CFG)
runs = []
for f in FRAMES:
kw = hin.load_sample(paths.pandaset_sample(f))
taps = TapRegistry(include=("img_feats",))
det, outs, plan, img = reference(**kw, taps=taps, return_outputs=True)
frame = hin.prepare_frame(**kw, cfg=CFG)
p_ort = st_ort.begin(frame.anchor_stamp, frame.stream.T_world_from_ego, frame.camera_stamps)
pe_o = oracle.position_embedding(**frame.rig.pe_inputs())
feats_o = oracle.backbone(img)
feeds = {"x": feats_o, "pos_embed": pe_o["pos_embed"], "cone": pe_o["cone"], "data_ego_pose": p_ort.E,
"data_ego_pose_inv": p_ort.Einv, **p_ort.pre}
outs_o = oracle.head(feeds)
st_ort.end(p_ort, outs_o)
runs.append({"frame": f, "img": img, "rig": frame.rig, "plan": plan, "outs": outs, "det": det,
"feats_o": feats_o, "pe_o": pe_o, "feeds_o": feeds, "outs_o": outs_o, "plan_o": p_ort})
return runs
def test_backbone_vs_ort(sample_run):
for r in sample_run[:2]:
p = pcc(r["outs"]["img_feats"], r["feats_o"])
print(f" frame {r['frame']} img_feats pcc={p:.12f} max_abs="
f"{np.abs(r['outs']['img_feats'] - r['feats_o']).max():.2e} (|x| <= {np.abs(r['feats_o']).max():.0f})")
assert p >= REF_PCC
def test_position_embedding_vs_ort(reference, sample_run):
r = sample_run[0]
pos, cone = reference.position_embedding(r["rig"])
p = pcc(pos, r["pe_o"]["pos_embed"])
print(f" pos_embed pcc={p:.12f} max_abs={np.abs(pos - r['pe_o']['pos_embed']).max():.2e}; cone max_abs="
f"{np.abs(cone - r['pe_o']['cone']).max():.2e}")
assert p >= REF_PCC and pcc(cone, r["pe_o"]["cone"]) >= REF_PCC
def test_head_vs_ort_on_the_same_inputs(reference, sample_run):
"""The head with ORT's own inputs (ORT backbone / PE / memory) at every frame of the sample."""
for r in sample_run:
fo = r["feeds_o"]
pre = {k: fo[k] for k in fo if k.startswith("pre_")}
out = reference.head(fo["x"], fo["pos_embed"], fo["cone"], fo["data_ego_pose"], fo["data_ego_pose_inv"], pre)
for k in PRE_SELECTION:
p = pcc(out[k], r["outs_o"][k])
assert p >= REF_PCC, (r["frame"], k, p)
idx_o = recover_topk_indices(r["outs_o"]["post_memory_embedding"], r["outs_o"]["outs_dec"])
ov = topk_overlap(out["topk_idx"].reshape(-1), idx_o)
tf = reference.head(fo["x"], fo["pos_embed"], fo["cone"], fo["data_ego_pose"], fo["data_ego_pose_inv"], pre,
topk_idx=idx_o)
worst = min(pcc(tf[k], r["outs_o"][k]) for k in POST_MEMORY)
print(f" frame {r['frame']}: cls pcc={pcc(out['all_cls_scores'], r['outs_o']['all_cls_scores']):.10f} "
f"top-256 overlap={ov:.4f} teacher-forced post-memory min pcc={worst:.10f}")
assert ov >= 0.99 and worst >= REF_PCC
assert np.array_equal(tf["post_memory_timestamp"], r["outs_o"]["post_memory_timestamp"])
def test_end_to_end_detections_vs_ort(sample_run):
"""Free-running over the 4 frames (each side with its own memory): the node's objects agree (set-based)."""
for r in sample_run:
ours = hp.autoware_postprocess(r["outs"]["all_cls_scores"], r["outs"]["all_bbox_preds"], CFG)["objects"]
theirs = hp.autoware_postprocess(r["outs_o"]["all_cls_scores"], r["outs_o"]["all_bbox_preds"], CFG)["objects"]
m = detection_f1(ours, theirs, dist=1.0)
print(f" frame {r['frame']}: {len(ours)} vs {len(theirs)} objects, F1={m['f1']:.4f} "
f"centre err max {m['centre_err_max']:.4f} m, score err max {m['score_err_max']:.2e}")
assert m["f1"] >= 0.99 and m["centre_err_max"] < 0.05
assert r["det"].to_dict()["num_detections"] == len(ours)
def _tf_records():
out = []
for run, f in (("ns0103", 1), ("ns0916", 0)):
g = paths.public_golden(run)
if g is not None and (g / f"tf_{f:04d}.npz").is_file():
out.append((run, f))
return out
@pytest.mark.parametrize("run,frame", _tf_records(), ids=[f"{r}-{f}" for r, f in _tf_records()])
def test_head_vs_ort_on_nuscenes_records(reference, oracle, run, frame):
"""A second rig (1600x900 nuScenes cameras, internal validation data): PE and head vs ORT."""
g = paths.public_golden(run)
rec, pe = np.load(g / f"tf_{frame:04d}.npz"), np.load(g / "pe_inputs_outputs.npz")
pe_o = oracle.position_embedding(pe["img_metas_pad"], pe["intrinsics"], pe["img2lidar"])
pos, cone = reference.net.pe(pe["img_metas_pad"], pe["intrinsics"], pe["img2lidar"])
assert pcc(pos.detach().numpy(), pe_o["pos_embed"]) >= REF_PCC
pre = {k: rec[k] for k in rec.files if k.startswith("pre_")}
feeds = {"x": rec["img_feats"], "pos_embed": pe_o["pos_embed"], "cone": pe_o["cone"],
"data_ego_pose": rec["data_ego_pose"], "data_ego_pose_inv": rec["data_ego_pose_inv"], **pre}
outs_o = oracle.head(feeds)
out = reference.head(rec["img_feats"], pe_o["pos_embed"], pe_o["cone"], rec["data_ego_pose"],
rec["data_ego_pose_inv"], pre)
for k in PRE_SELECTION:
assert pcc(out[k], outs_o[k]) >= REF_PCC, k
|