File size: 6,247 Bytes
77e3c67 | 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 | # SPDX-License-Identifier: Apache-2.0
"""The CPU reference's building blocks that need torch but neither onnxruntime nor the full network run:
weights addressed by consuming node + the deployed-graph check, the three plugins against the research numpy
implementations, and the TT port's exact rewrites (``reference.port_form``) against the deployed form.
TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_bevdet/tests/test_reference_host.py
Needs the pinned ONNX (``tests/paths.py``: ``$BEVDET_WEIGHTS_DIR``, the HF cache or the workspace download); the
weight tests skip without it. The end-to-end reference against ONNX Runtime is ``test_reference_cpu.py``.
"""
from __future__ import annotations
import numpy as np
import pytest
torch = pytest.importorskip("torch")
from tt_bevdet.reference import plugins, port_form # noqa: E402
from tt_bevdet.reference.config import HEAD_NAMES, default_config # noqa: E402
from tt_bevdet.reference.weights import ONNX_SHA256, BEVDetWeights, module_names # noqa: E402
from tt_bevdet.tests import paths # noqa: E402
from tt_bevdet.ttaw.metrics import pcc # noqa: E402
CFG = default_config()
DN = "img_view_transformer.depth_net."
@pytest.fixture(scope="module")
def weights():
path = paths.weights_onnx()
if path is None:
pytest.skip("bevdet_one_lt_d.onnx not found (set BEVDET_WEIGHTS_DIR)")
pytest.importorskip("onnx")
return BEVDetWeights(path)
@pytest.fixture(scope="module")
def net(weights):
from tt_bevdet.reference.model import BEVDetNet
torch.set_num_threads(2)
return BEVDetNet(weights, CFG)
def test_weights_by_consuming_node(weights):
assert weights.sha256 == ONNX_SHA256
names = module_names()
n_conv = sum(len(v) for k, v in names.items() if k != "depth_net_linear")
assert n_conv == 107 == len(weights.onnx.nodes("Conv")) # SPEC 4.2: Conv 107
c = weights.conv("img_backbone.conv1")
assert c.node == "Conv_1" and c.weight.shape == (64, 3, 7, 7) and c.stride == (2, 2) and c.padding == (3, 3)
assert [weights.conv(DN + f"depth_conv.3.aspp{i}.atrous_conv").dilation for i in (2, 3, 4)] == \
[(6, 6), (12, 12), (18, 18)]
assert weights.conv("img_bev_encoder_backbone.layers.0.0.downsample").stride == (2, 2)
w, b = weights.linear(DN + "context_mlp.fc1")
assert w.shape == (512, 27) and b.shape == (512,)
bn = weights.depth_bn()
assert bn.channels == 27 and bn.eps == pytest.approx(1e-5)
n_params = sum(c.weight.size + (0 if c.bias is None else c.bias.size) for _, c in weights.iter_convs())
n_params += sum(a.size for m in names["depth_net_linear"] for a in weights.linear(m)) + 4 * 27
assert n_params == 76_187_718 # SPEC 4.3
def test_align_gather_bevpool_vs_research():
bc = paths.import_research()
if bc is None:
pytest.skip("research reference not present")
rng = np.random.default_rng(2)
adj = rng.normal(size=(1, 8, 4, 20, 24)).astype(np.float32)
trs = np.zeros((1, 8, 6), np.float32)
for n in range(8):
a = rng.uniform(-0.3, 0.3)
trs[0, n] = [np.cos(a), -np.sin(a), rng.uniform(-6, 6), np.sin(a), np.cos(a), rng.uniform(-6, 6)]
trs[0, 3] = [1, 0, 0, 0, 1, 0]
got = plugins.align_bev(adj, trs).numpy()
np.testing.assert_array_equal(got, bc.align_bev(adj, trs)) # bit for bit
np.testing.assert_array_equal(got[0, 3], adj[0, 3]) # identity = copy
curr = rng.normal(size=(1, 4, 20, 24)).astype(np.float32)
for flag in (0, 1):
np.testing.assert_array_equal(plugins.gather_bev(got, curr, flag).numpy(), bc.gather_bev(got, curr, flag))
def test_camera_gates_host_fold_vs_graph(weights, net):
rng = np.random.default_rng(4)
cp = np.zeros((1, 6, 27), np.float32)
cp[0, :, :4] = rng.uniform([1000, 1000, 780, 430], [1300, 1300, 830, 470], size=(6, 4))
cp[0, :, 4:15] = [0.44, 0, 0, 0, 0.44, -140, 1, 0, 0, 1, 1]
cp[0, :, 15:] = rng.normal(size=(6, 12))
g_ctx, g_depth = port_form.camera_gates(weights, cp)
t_ctx, t_depth = net.camera_gates(torch.from_numpy(cp))
assert np.abs(g_ctx - t_ctx.numpy()).max() < 1e-5 and np.abs(g_depth - t_depth.numpy()).max() < 1e-5
w, b = weights.linear(DN + "depth_mlp.fc1")
bn = weights.depth_bn()
wf, bf = port_form.fold_bn_into_linear_input(w, b, bn)
x = cp[0].astype(np.float64)
np.testing.assert_allclose(x @ wf.T + bf, bn.apply(x, axis=1) @ w.astype(np.float64).T + b, rtol=1e-9, atol=1e-9)
def test_depthnet_port_form_matches_deployed(weights, net):
"""1x3 ASPP branch 4, GAP branch as a per-frame bias, depth logits padded to 128 with -1e4, host gates."""
rng = np.random.default_rng(6)
x = torch.from_numpy((rng.normal(size=(2, 512, 16, 44)) * 0.5).astype(np.float32))
gates = (rng.uniform(0, 1, (2, 512)).astype(np.float32), rng.uniform(0, 1, (2, 512)).astype(np.float32))
with torch.no_grad():
d_ref, f_ref = net.depth_net(x, gates=tuple(torch.from_numpy(g) for g in gates))
d_port, f_port = port_form.forward_depth_net_port(net, x, gates, weights, CFG)
np.testing.assert_array_equal(f_port.numpy(), f_ref.numpy())
assert d_port.shape == d_ref.shape and np.abs(d_port.numpy() - d_ref.numpy()).max() < 2e-6
assert pcc(d_port.numpy(), d_ref.numpy()) > 0.99999999
pad = port_form.pad_depth_logits(weights)
assert pad.weight.shape[0] == 128 and pad.n_valid == 118 and np.all(pad.bias[118:] == -1e4)
assert float(np.exp(np.float32(-1e4) - np.float32(-50.0))) == 0.0 # the pad columns vanish in fp32
def test_merged_head_matches_deployed(weights, net):
rng = np.random.default_rng(8)
x = torch.from_numpy(rng.normal(size=(1, 256, 32, 32)).astype(np.float32))
with torch.no_grad():
ref = net.center_head(x)
got = port_form.forward_center_head_port(net, x, weights)
m = port_form.merge_center_head(weights)
assert m.w1.shape == (384, 64, 3, 3) and m.w2.shape == (20, 384, 3, 3) and [s[0] for s in m.splits] == \
list(HEAD_NAMES)
for k in HEAD_NAMES:
assert np.abs(got[k].numpy() - ref[k].numpy()).max() < 5e-5, k
assert pcc(got[k].numpy(), ref[k].numpy()) > 0.9999999, k
|