File size: 8,461 Bytes
ddaad99 | 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 | # SPDX-License-Identifier: Apache-2.0
"""The CPU reference's building blocks (``reference.ops``, ``reference.weights``, ``reference.model``) on the host:
the plugin semantics against literal float64 ports of the CUDA loops, the strict weight load and the port's BatchNorm
folding, and the port's rewrites against the deployed form (rebatched vs dense spatial cross-attention, rotation as a
gather). Needs torch (tt-env has it); the whole-network comparisons with ONNX Runtime are in test_reference_cpu.py.
TT_VISIBLE_DEVICES=none python -m pytest -q -p fake_ttnn_plugin code/tt_bevformer/tests/test_reference_host.py
"""
from __future__ import annotations
import numpy as np
import pytest
torch = pytest.importorskip("torch")
from tt_bevformer.host import geometry as hg # noqa: E402
from tt_bevformer.reference import config as C # noqa: E402
from tt_bevformer.reference import ops # noqa: E402
from tt_bevformer.tests import paths # noqa: E402
from tt_bevformer.ttaw.metrics import pcc # noqa: E402
torch.set_num_threads(4)
Q = C.NUM_QUERY_BEV
@pytest.fixture(scope="module")
def weights():
p = paths.weights_file()
if p is None:
pytest.skip("BEVFormer weights not found (set BEVFORMER_WEIGHTS_DIR)")
from tt_bevformer.reference.weights import BEVFormerWeights
return BEVFormerWeights(p)
# --------------------------------------------------------------------------------------------- plugin semantics
@pytest.mark.parametrize("case", [dict(N=2, shapes=[(7, 9)], H=2, Cc=4, Qn=5, Pg=1, P=4), # TSA-like
dict(N=3, shapes=[(5, 8)], H=2, Cc=4, Qn=6, Pg=4, P=8), # SCA-like
dict(N=1, shapes=[(6, 6), (3, 3)], H=2, Cc=4, Qn=4, Pg=1, P=2)])
def test_msda_matches_the_cuda_loop(case):
torch.manual_seed(0)
S = sum(h * w for h, w in case["shapes"])
L = len(case["shapes"])
value = torch.randn(case["N"], S, case["H"], case["Cc"])
ref = torch.rand(case["N"], case["Qn"], 1, case["Pg"] * 2) * 1.4 - 0.2 # some points outside
off = torch.randn(case["N"], case["Qn"], case["H"], L * case["P"] * 2) * 3.0
aw = torch.randn(case["N"], case["Qn"], case["H"], L * case["P"])
a = ops.msda(value, case["shapes"], ref, off, aw, q_chunk=3)
b = ops.msda_naive(value, case["shapes"], ref, off, aw)
assert float((a.double() - b).abs().max()) < 1e-5
def test_mdcn_matches_the_cuda_loop_and_conv2d():
torch.manual_seed(1)
x = torch.randn(2, 5, 7, 9)
offset = torch.randn(2, 18, 7, 9) * 2.0
mask = torch.sigmoid(torch.randn(2, 9, 7, 9))
w = torch.randn(4, 5, 3, 3)
bias = torch.randn(4)
a = ops.mdcn(x, offset, mask, w, bias)
b = ops.mdcn_naive(x, offset, mask, w) + bias.double().view(1, -1, 1, 1)
assert float((a.double() - b).abs().max()) < 2e-5
a0 = ops.mdcn(x, torch.zeros_like(offset), torch.ones_like(mask), w) # zero offset, mask 1
assert float((a0 - torch.nn.functional.conv2d(x, w, None, 1, 1)).abs().max()) < 2e-5
def test_rotate_identity_and_tie_rule():
img = torch.randn(3, C.BEV_H, C.BEV_W)
assert torch.equal(ops.rotate(img, torch.tensor([0.0]), C.ROTATE_CENTER), img)
a = ops.rotate(img, torch.tensor([0.37]), C.ROTATE_CENTER)
b = ops.rotate(img, torch.tensor([0.37]), C.ROTATE_CENTER, plugin_rounding=False) # torchvision (training)
assert 0 < int((a != b).any(0).sum()) <= 4 # ties only (SPEC 4.2)
# ---------------------------------------------------------------------------------------------------- weights
def test_strict_weight_coverage(weights):
from tt_bevformer.reference.weights import NUM_TENSORS, expected_keys
assert len(weights.sd) == NUM_TENSORS == len(expected_keys())
assert len(weights.dcn_modules()) == 26 # layer3 23 + layer4 3
assert len(C.backbone_conv_specs()) == 1 + 33 * 3 + 26 + 4 # stem, blocks, offsets, ds
assert weights.conv("img_backbone.layer2.0.conv1").stride == (2, 2) # caffe: stride on conv1
assert weights.conv("img_backbone.layer2.0.conv2").stride == (1, 1)
assert weights.conv("img_backbone.layer3.5.conv2").weight.shape == (256, 256, 3, 3)
assert weights.conv("img_backbone.layer3.5.conv2.conv_offset").weight.shape == (27, 256, 3, 3)
with pytest.raises(KeyError):
weights.conv("img_neck.lateral_convs.0.conv")
def test_bn_folding_equals_conv_then_bn(weights):
"""The folded convs equal conv + frozen BN (float64 oracle) to float32 rounding, also through a DCN."""
torch.manual_seed(2)
sd = weights.sd
for module, bn, cin in (("img_backbone.conv1", "img_backbone.bn1", 3),
("img_backbone.layer3.0.conv2", "img_backbone.layer3.0.bn2", 256)):
x = torch.randn(1, cin, 12, 20, dtype=torch.float64)
w = torch.from_numpy(np.asarray(sd[f"{module}.weight"], np.float64))
cp = weights.conv(module)
stride, pad = cp.stride[0], cp.padding[0]
g, b, m, v = (torch.from_numpy(np.asarray(sd[f"{bn}.{k}"], np.float64)) for k in
("weight", "bias", "running_mean", "running_var"))
if module.endswith("conv2"): # DCN with zero offsets
ref = ops.mdcn(x, torch.zeros(1, 18, 12, 20, dtype=torch.float64),
torch.ones(1, 9, 12, 20, dtype=torch.float64), w)
got = ops.mdcn(x.float(), torch.zeros(1, 18, 12, 20), torch.ones(1, 9, 12, 20),
torch.from_numpy(cp.weight), torch.from_numpy(cp.bias)).double()
else:
ref = torch.nn.functional.conv2d(x, w, None, stride, pad)
got = torch.nn.functional.conv2d(x.float(), torch.from_numpy(cp.weight), torch.from_numpy(cp.bias),
stride, pad).double()
ref = (ref - m.view(1, -1, 1, 1)) / torch.sqrt(v.view(1, -1, 1, 1) + C.BN_EPS) * g.view(1, -1, 1, 1) \
+ b.view(1, -1, 1, 1)
rel = float((got - ref).abs().max() / ref.abs().max())
assert rel < 1e-5, (module, rel)
# ----------------------------------------------------------------------------------------- port rewrites
@pytest.fixture(scope="module")
def net(weights):
from tt_bevformer.reference.model import BEVFormerNet
return BEVFormerNet(weights)
def test_rebatched_sca_equals_dense(net):
"""Encoder layer 0's spatial cross-attention, dense over 6 x 22,500 pairs (the TRT graph) vs rebatched to the
valid pairs with the gather-form scatter (the port), on the sample0 rig with random query / value."""
golden = paths.research_file("golden", "sample0_golden.npz")
if golden is None:
pytest.skip("research/bevformer golden not in this workspace")
from tt_bevformer.reference.model import PortGeometry, point_sampling, reference_points_3d
with np.load(golden) as g:
l2i = g["lidar2img"]
geo_np = hg.frame_geometry(l2i, np.zeros(18, np.float32))
geo = PortGeometry.from_numpy(geo_np)
torch.manual_seed(3)
query = torch.randn(Q, C.EMBED)
value = torch.randn(C.NUM_CAMS, C.FEAT_LEN, C.EMBED)
ref_cam, mask_n = point_sampling(reference_points_3d(), torch.from_numpy(l2i))
with torch.no_grad():
dense = net.sca_dense(0, query, value, ref_cam, mask_n)
dense_host = net.sca_dense(0, query, value, geo.ref_cam, torch.from_numpy(geo_np["bev_mask_n"]))
rebatched = net.sca_rebatched(0, query, value, geo)
assert float((rebatched - dense_host).abs().max()) < 1e-4 # same geometry: summation order
assert pcc(rebatched.numpy(), dense.numpy()) > 0.9999999 # host vs graph geometry
def test_rotation_gather_equals_rotate_op():
"""forward_port rotates prev_bev with the host map as a row gather; on angles without ties it is the plugin."""
torch.manual_seed(4)
prev = torch.randn(Q, C.EMBED)
table = torch.cat([prev, torch.zeros(1, C.EMBED)], 0)
for a in (0.0, 1.0, -3.0, 3.0, 12.5):
idx = torch.from_numpy(hg.rotation_index_map(a))
ours = table[idx]
theirs = ops.rotate(prev.view(C.BEV_H, C.BEV_W, C.EMBED).permute(2, 0, 1), torch.tensor([a]),
C.ROTATE_CENTER).permute(1, 2, 0).reshape(Q, C.EMBED)
assert int((ours != theirs).any(1).sum()) <= 2, a
|