File size: 6,908 Bytes
fe51a63 | 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 | # SPDX-License-Identifier: Apache-2.0
"""The TT device dataflow (``reference/rewrites.py``) equals the literal graph (``reference/model.py``) on the CPU
(torch + the weights; no device, no onnxruntime):
- the whole network on the staged PandaSet sample (QT128, the faster variant; OT128 on the same frame in
``test_reference_cpu.py``) and on a small synthetic frame with tight capacities: every module tap PCC >= 0.9999999
and ``pred_probs`` within 1e-4, argmax identical -- the plan's sorted / padded points, zero-sentinel gathers,
last-voxel-wins pixel tables, log-step segmented max, K-split convs and point layers and interpolation-matrix
up-sampling change nothing but the summation order;
- padding is harmless: two capacities give the same outputs on the real rows (to the last bits of BLAS blocking);
- the log-step segmented max equals ``scatter_reduce(amax)`` for any rounds >= the frame's need, and fails with one
round fewer when a segment is longer;
- the align-corners interpolation matrices equal ``F.interpolate``; the emulation modes round as documented.
TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_frnet/tests/test_reference_rewrites_host.py
"""
from __future__ import annotations
import numpy as np
import pytest
torch = pytest.importorskip("torch")
pytest.importorskip("onnx")
import torch.nn.functional as F # noqa: E402
from tt_frnet.host.frame import FRNetPreprocessor # noqa: E402
from tt_frnet.host.preprocess import autoware_preprocess, build_plan, segmax_rounds # noqa: E402
from tt_frnet.reference.config import FRNetConfig # noqa: E402
from tt_frnet.reference.model import FRNetTorch, scatter_max # noqa: E402
from tt_frnet.reference.rewrites import (Emulation, FRNetPlanTorch, plan_taps_to_literal, round_mode, # noqa: E402
segmax_logstep)
from tt_frnet.reference.weights import MODULE_TAPS, load_weights # noqa: E402
from tt_frnet.tests import _data # noqa: E402
from tt_frnet.ttaw.golden import TapRegistry # noqa: E402
from tt_frnet.ttaw.metrics import pcc # noqa: E402
from tt_frnet.ttaw.ops.upsample import interp_matrix # noqa: E402
THREADS = 4
PATTERNS = list(MODULE_TAPS) + ["backbone.segmax*", "decoder.up*"]
@pytest.fixture(scope="module")
def qt128():
torch.set_num_threads(THREADS)
w = _data.weights_dir()
return load_weights(w, "qt128"), FRNetConfig.from_model_dir(w, "qt128")
def _compare(weights, cfg, inputs, n_cap=None, m_cap=None):
plan = build_plan(inputs, cfg, pre_norm=(weights.pre_norm_scale, weights.pre_norm_shift), n_cap=n_cap,
m_cap=m_cap)
tl, tr = TapRegistry(include=PATTERNS), TapRegistry(include=PATTERNS)
lit = FRNetTorch(weights)(inputs.points, inputs.coors, inputs.voxel_coors, inputs.inverse_map, taps=tl)
out = FRNetPlanTorch(weights)(plan, taps=tr)
worst = min((pcc(plan_taps_to_literal(n, tr[n], plan), tl[n]), n) for n in tl.names())
return plan, lit, out, worst, tr
def test_plan_equals_literal_on_the_sample(qt128):
weights, cfg = qt128
_data.need_pandaset()
prep = FRNetPreprocessor(cfg, plan=False)(_data.SAMPLE, _data.SAMPLE_CALIB)
plan, lit, out, worst, _ = _compare(weights, cfg, prep.inputs)
assert plan.n_cap == cfg.n_cap == 120000 and plan.m_cap == 60000
assert worst[0] >= 0.9999999, worst
got = plan.unsort(out)
assert np.abs(got - lit).max() < 1e-4
np.testing.assert_array_equal(got.argmax(1), lit.argmax(1))
assert np.allclose(out[plan.n:].sum(1), 1.0, atol=1e-5) # padded rows: finite, ignored
def test_plan_on_a_small_frame_and_padding_is_harmless(qt128):
weights, cfg = qt128
rng = np.random.default_rng(7)
n = 6000
az = rng.uniform(-np.pi, np.pi, n)
el = np.radians(rng.uniform(-50, 50, n))
r = rng.uniform(2, 60, n)
pts = np.c_[r * np.cos(el) * np.cos(az), r * np.cos(el) * np.sin(az), r * np.sin(el),
rng.uniform(0, 255, n)].astype(np.float32)
pts[:500, :3] = pts[0, :3] + rng.normal(0, 0.01, (500, 3)) # one dense cell: a long segment
inputs = autoware_preprocess(pts, cfg, crop_bounds=(0.0,) * 6)
p1, lit, o1, worst, t1 = _compare(weights, cfg, inputs, n_cap=inputs.num_points + 1, m_cap=inputs.num_voxels + 1)
assert p1.max_segment_length >= 400 and segmax_rounds(p1.max_segment_length) >= 9
assert worst[0] >= 0.9999999, worst
assert np.abs(p1.unsort(o1) - lit).max() < 1e-4
p2 = build_plan(inputs, cfg, pre_norm=(weights.pre_norm_scale, weights.pre_norm_shift),
n_cap=inputs.num_points + 4096, m_cap=inputs.num_voxels + 777)
o2 = FRNetPlanTorch(weights)(p2)
# padded rows never reach a real row: only BLAS blocking of the taller matrices may move the last bits
np.testing.assert_allclose(p2.unsort(o2), p1.unsort(o1), rtol=0, atol=1e-6)
def test_segmax_logstep_equals_scatter_reduce(qt128):
_, cfg = qt128
rng = np.random.default_rng(3)
pts = np.c_[rng.normal(0, 15, (3000, 3)), rng.uniform(0, 255, 3000)].astype(np.float32)
pts[:300, :3] = pts[0, :3]
inputs = autoware_preprocess(pts, cfg, crop_bounds=(0.0,) * 6)
plan = build_plan(inputs, cfg, n_cap=inputs.num_points + 64, m_cap=inputs.num_voxels + 9)
x = torch.randn(plan.n_cap, 5)
longest = int(np.argmax(np.diff(plan.seg_offsets)))
x[int(plan.seg_offsets[longest])] = 100.0 # the longest segment's max sits on its FIRST row
ref = scatter_max(x[: plan.n], torch.from_numpy(plan.seg_id[: plan.n]), plan.m)
need = segmax_rounds(plan.max_segment_length)
for rounds in (need, need + 2):
assert torch.equal(segmax_logstep(x, plan, rounds)[: plan.m], ref)
assert not torch.equal(segmax_logstep(x, plan, need - 1)[: plan.m], ref)
@pytest.mark.parametrize("hw,s", [((128, 1024), 2), ((128, 256), 8), ((16, 128), 4)])
def test_interp_matrices_equal_align_corners_interpolate(hw, s):
H, W = hw
x = torch.randn(1, 3, H // s, W // s, dtype=torch.float64)
ah = torch.from_numpy(interp_matrix(H // s, H, mode="linear", align_corners=True))
aw = torch.from_numpy(interp_matrix(W // s, W, mode="linear", align_corners=True))
y = torch.einsum("hi,bcij,wj->bchw", ah, x, aw)
ref = F.interpolate(x, size=(H, W), mode="bilinear", align_corners=True)
assert torch.allclose(y, ref, atol=1e-12)
def test_rounding_modes():
x = torch.tensor([1.0, 1.0 + 2 ** -8, 1.0 + 2 ** -7 + 2 ** -9, -3.14159265, 1e-30])
assert torch.equal(round_mode(x, "fp32"), x)
assert torch.equal(round_mode(x, "bf16"), x.to(torch.bfloat16).float())
t = round_mode(x, "tf32")
assert (t.view(torch.int32) & 0x1FFF).eq(0).all() and (t.abs() <= x.abs()).all()
e = Emulation.parse("conv_act=bf16,point_act=tf32")
assert e.label() == "conv_act=bf16,point_act=tf32" and not e.exact and Emulation().exact
with pytest.raises(ValueError):
Emulation.parse("point_act=fp16")
|