File size: 6,829 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 | # SPDX-License-Identifier: Apache-2.0
"""The TT-friendly forms of ``reference/rewrites.py`` equal the literal modules of ``reference/model.py`` with the real
weights (CPU, torch fp32): what the device graph computes is proven here before device time is spent. Inputs are
seeded random tensors of the real channel counts (small spatial sizes), plus ties for the local max.
TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_transfusion/tests/test_reference_rewrites_host.py
"""
from __future__ import annotations
import numpy as np
import pytest
pytest.importorskip("onnx")
torch = pytest.importorskip("torch")
F = torch.nn.functional
from tt_transfusion.reference import rewrites as rw # noqa: E402
from tt_transfusion.reference.model import TransFusionTorch # noqa: E402
from tt_transfusion.reference.weights import load_weights # noqa: E402
from tt_transfusion.tests import _data # noqa: E402
RNG = np.random.default_rng(20261008)
@pytest.fixture(scope="module")
def net():
torch.set_num_threads(2)
return TransFusionTorch(load_weights(_data.weights_dir()), threads=2)
def _t(a):
return torch.from_numpy(np.array(a, dtype=np.float32, copy=True))
def _pillars(n=300, slots=20, feats=11):
"""Random decorated pillars: 1..20 points (a third of them full), padded slots zero."""
num = RNG.integers(1, slots + 1, n)
num[: n // 3] = slots
f = RNG.normal(0.0, 1.0, (n, slots, feats)).astype(np.float32)
f[:, :, 3] = RNG.integers(0, 256, (n, slots)) # intensity scale
f *= (np.arange(slots)[None, :] < num[:, None])[:, :, None]
return f, num
def test_pfn_tile_form_and_32_slot_pillars(net):
f, num = _pillars()
p = net.weights.pfn
with torch.no_grad():
ref = net.pfn(f).numpy()
f32 = rw.pad_pillar_slots(f) # (n, 32, 16): slots 20-31 repeat slot 0
w0 = np.zeros((rw.PFN_IN_PADDED, p.w0.shape[1]), np.float32)
w0[: p.w0.shape[0]] = p.w0
got = rw.pfn_tile_form(_t(f32), _t(w0), _t(p.b0), _t(p.w1a), _t(p.w1b), _t(p.b1)).numpy()
# negative control: zero rows in slots 20-31 add relu(b0) to the maxima of the full pillars
zero = np.zeros_like(f32)
zero[:, :20, :11] = f
bad = rw.pfn_tile_form(_t(zero), _t(w0), _t(p.b0), _t(p.w1a), _t(p.w1b), _t(p.b1)).numpy()
assert f32.shape == (300, 32, 16) and np.array_equal(f32[:, 20:, :11], np.repeat(f[:, :1], 12, axis=1))
np.testing.assert_allclose(got, ref, rtol=1e-5, atol=2e-5)
full = num == 20
assert np.abs(bad[full] - ref[full]).max() > 1e-3 # the trap the duplicated slots avoid
np.testing.assert_allclose(bad[~full], ref[~full], rtol=1e-5, atol=2e-5)
def test_k_split_shared_conv(net):
x = [_t(np.maximum(RNG.normal(0, 1, (1, 128, 24, 24)), 0)) for _ in range(3)]
shared = net.weights.head.shared
with torch.no_grad():
ref = net.layer("head.shared", torch.cat(x, dim=1))
got = rw.k_split_torch(x, shared)
np.testing.assert_allclose(got.numpy(), ref.numpy(), rtol=1e-4, atol=1e-4)
def test_deblock_rewrites(net):
d0, _, d2 = net.weights.deblocks
x0 = np.maximum(RNG.normal(0, 1, (1, 64, 16, 16)), 0).astype(np.float32)
x2 = np.maximum(RNG.normal(0, 1, (1, 256, 8, 8)), 0).astype(np.float32)
with torch.no_grad():
ref0 = net.layer("neck.deblock0", _t(x0)).numpy()[0].transpose(1, 2, 0) # (8, 8, 128)
ref2 = net.layer("neck.deblock2", _t(x2)).numpy()[0].transpose(1, 2, 0) # (16, 16, 128)
wl, bl = rw.deblock0_as_linear(d0)
s2d = rw.space_to_depth_nhwc(x0[0].transpose(1, 2, 0), 2) # (8, 8, 256)
got0 = np.maximum(s2d.reshape(64, 256) @ wl + bl, 0).reshape(8, 8, 128)
np.testing.assert_allclose(got0, ref0, rtol=1e-4, atol=1e-4)
wl2, bl2 = rw.deconv_as_linear(d2)
y = x2[0].transpose(1, 2, 0).reshape(64, 256) @ wl2 + bl2
got2 = np.maximum(rw.depth_to_space_nhwc(y, 8, 8, 2, 128), 0)
np.testing.assert_allclose(got2, ref2, rtol=1e-4, atol=1e-4)
def test_masked_local_max_chain_is_exact(net):
"""C22's chain with TransFusion's masks equals MaxPool p0 + zero border + equality bit for bit, with ties."""
heat = np.round(RNG.uniform(0, 1, (1, 5, 192, 192)) * 64) / 64 # many exact ties
heat[0, 2, 50:60, 70:80] = 0.5 # a flat plateau
heat = _t(heat)
lp = torch.zeros_like(heat)
lp[:, :, 1:-1, 1:-1] = F.max_pool2d(heat, 3, 1, 0)
literal = heat * (heat == lp).to(heat.dtype)
m_pass, m_eq = rw.local_max_masks(192, 192, 5)
got = rw.local_max_masked_torch(heat, m_pass, m_eq)
assert torch.equal(got, literal)
assert float(got[0, :, 0].abs().max()) == 0.0 and float(got[0, :, :, -1].abs().max()) == 0.0
def test_query_init_tables_are_exact(net):
h = net.weights.head
pos = RNG.integers(0, 36864, 500)
assert np.array_equal(rw.query_pos_from_index(pos, 192), h.bev_pos[pos])
onehot = np.zeros((1, 5, 500), np.float32)
cls = RNG.integers(0, 5, 500)
onehot[0, cls, np.arange(500)] = 1.0
with torch.no_grad():
conv = F.conv1d(_t(onehot), _t(h.class_encoding_w[:, :, None]), _t(h.class_encoding_b))[0].T.numpy()
np.testing.assert_array_equal(h.class_table[cls], conv)
def test_cross_attention_constant_kv_tables_and_padded_heads(net):
d = net.weights.head.decoder
kpe = net.key_pos_embed().numpy() # (36864, 128)
lidar = RNG.normal(0, 1, (36864, 128)).astype(np.float32)
x = RNG.normal(0, 1, (500, 128)).astype(np.float32)
with torch.no_grad():
ref = net._mha("cross_attn", _t(x), _t(lidar + kpe)).numpy()
ck, cv = rw.cross_attention_kv_tables(d.cross_attn, kpe)
got = rw.padded_head_attention(_t(x), _t(lidar), _t(lidar), d.cross_attn, k_const=ck, v_const=cv).numpy()
ref_self = net._mha("self_attn", _t(x), _t(x)).numpy()
got_self = rw.padded_head_attention(_t(x), _t(x), _t(x), d.self_attn).numpy()
np.testing.assert_allclose(got, ref, rtol=1e-4, atol=1e-4)
np.testing.assert_allclose(got_self, ref_self, rtol=1e-4, atol=1e-4)
def test_merged_prediction_heads(net):
x = RNG.normal(0, 1, (500, 128)).astype(np.float32)
merged = rw.merged_prediction_head_weights(net.weights.head)
wh, bh, wo, bo, slices = merged
assert wh.shape == (128, 384) and wo.shape == (384, 32) and slices[-1] == ("heatmap", 10, 15)
with torch.no_grad():
got = rw.merged_heads_torch(_t(x), merged)
y = _t(x).T[None]
for name, hid, fin in net.weights.head.prediction_heads:
ref = net.layer(fin.name, net.layer(hid.name, y))
np.testing.assert_allclose(got[name].numpy(), ref.numpy(), rtol=1e-4, atol=1e-4)
|