| |
| """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 |
| from tt_bevdet.reference.config import HEAD_NAMES, default_config |
| from tt_bevdet.reference.weights import ONNX_SHA256, BEVDetWeights, module_names |
| from tt_bevdet.tests import paths |
| from tt_bevdet.ttaw.metrics import pcc |
|
|
| 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")) |
| 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 |
|
|
|
|
| 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)) |
| np.testing.assert_array_equal(got[0, 3], adj[0, 3]) |
| 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 |
|
|
|
|
| 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 |
|
|
