| |
| """Structure of the weights the reference (and later the device graph) reads from ``bevfusion_lidar.onnx``: every |
| wiring assertion of ``reference/weights.py`` holds (it raises otherwise), the layer plan matches the host rulebook |
| plan, every BatchNorm fold is exact (float64 folded vs unfolded on random inputs, per layer kind), the parameter |
| count after the folds equals an independent count over the graph, constants (``bev_pos``, pooled classes, K, heads) |
| are the SPEC's. numpy + onnx only (no torch, no onnxruntime): runs in tt-env with the fake ttnn and in the research |
| venv; skipped without the weights. |
| """ |
| from __future__ import annotations |
|
|
| import numpy as np |
| import pytest |
|
|
| pytest.importorskip("onnx") |
| pytest.importorskip("yaml") |
|
|
| from tt_bevfusion.host.preprocess import STRIDED_LAYERS, SUBM_LAYERS |
| from tt_bevfusion.reference.config import BEVFusionConfig |
| from tt_bevfusion.reference.weights import ENCODER_PLAN, load_weights |
| from tt_bevfusion.tests import _data |
| from tt_bevfusion.ttaw.weights import OnnxWeights |
|
|
|
|
| @pytest.fixture(scope="module") |
| def loaded(): |
| d = _data.weights_dir() |
| cfg = BEVFusionConfig.from_model_dir(d) |
| return cfg, load_weights(d, cfg), OnnxWeights(d / _data.ONNX) |
|
|
|
|
| def test_config_from_the_weights(loaded): |
| cfg, _, _ = loaded |
| assert cfg.version == "v2.0" and cfg.class_names == ("CAR", "TRUCK", "BUS", "BICYCLE", "PEDESTRIAN") |
| assert cfg.grid == (1440, 1440, 40) and cfg.num_proposals == 500 and cfg.max_voxels == 256000 |
| assert cfg.max_points_per_voxel == 10 and not cfg.use_intensity and cfg.out_size_factor == 8 |
| assert cfg.remapper.apply(np.array([1, 1, 2]), np.array([5.0, 10.0, 10.0]), np.array([3.0, 4.0, 4.0])).tolist() \ |
| == [2, 4, 4] |
|
|
|
|
| def test_layer_plan(loaded): |
| _, w, _ = loaded |
| assert len(w.encoder) == 21 and [x.name for x in w.encoder] == [p[0] for p in ENCODER_PLAN] |
| maps = {x.map for x in w.encoder} |
| assert maps == set(SUBM_LAYERS) | set(STRIDED_LAYERS) |
| chans = [(x.in_channels, x.out_channels, x.kernel_volume) for x in w.encoder] |
| assert chans[0] == (4, 16, 27) and chans[-1] == (128, 128, 3) |
| assert [x.name for x in w.encoder if x.residual] == [f"sparse.layer{s}.{b}" for s in range(1, 5) for b in range(2)] |
| assert [len(b) for b in w.blocks] == [6, 6] and w.blocks[1][0].stride == 2 |
| assert [(d.kind, d.kernel, d.stride) for d in w.deblocks] == [("conv", 1, 1), ("conv_transpose", 2, 2)] |
| h = w.head |
| assert h.pooled_classes == (0, 1, 2, 3) and h.num_proposals == 500 and h.nms_kernel == 3 |
| assert h.decoder.self_attn.num_heads == 8 and h.decoder.self_attn.head_dim == 16 |
| assert h.decoder.cross_attn.q_divisor == 4.0 |
| assert h.bev_pos[181].tolist() == [1.5, 1.5] and h.bev_pos[-1].tolist() == [179.5, 179.5] |
|
|
|
|
| def test_parameter_count_after_the_folds(loaded): |
| """ONNX float32 weights (8,308,692, S:bevfusion:218) minus what each BatchNorm fold removes: its 4 C statistics |
| become a C bias (-3 C) when the folded conv had no bias, or merge into the existing bias (-4 C).""" |
| _, w, o = loaded |
| import onnx |
| from onnx import numpy_helper |
|
|
| inits = onnx.load(str(_data.weights_dir() / _data.ONNX)).graph.initializer |
| total = sum(numpy_helper.to_array(t).size for t in inits if t.data_type == onnx.TensorProto.FLOAT) |
| removed = 0 |
| for bn in o.nodes("BatchNormalization"): |
| prod = o.producer(bn.inputs[0]) |
| c = o.array(bn.inputs[1]).size |
| has_bias = prod.op_type in ("Conv", "ConvTranspose") and len(prod.inputs) > 2 and prod.inputs[2] |
| removed += 4 * c if has_bias else 3 * c |
| assert total == 8_308_692 |
| assert w.num_params() == total - removed |
|
|
|
|
| def _bn_apply(y, o, bn_name, axis): |
| bn = o.batchnorm(bn_name) |
| shape = [1] * y.ndim |
| shape[axis] = -1 |
| return (y - bn.mean.reshape(shape)) / np.sqrt(bn.var.reshape(shape) + bn.eps) * bn.gamma.reshape(shape) \ |
| + bn.beta.reshape(shape) |
|
|
|
|
| def test_bn_folds_are_exact(loaded): |
| _, w, o = loaded |
| rng = np.random.default_rng(0) |
| |
| for layer in (w.encoder[0], w.encoder[6], w.encoder[-1]): |
| node = o.node(layer.onnx_node) |
| W = np.asarray(o.array(node.inputs[1]), np.float64) |
| x = rng.normal(size=(7, layer.kernel_volume, layer.in_channels)) |
| y = np.einsum("nkc,okc->no", x, W.reshape(W.shape[0], -1, W.shape[-1])) |
| ref = _bn_apply(y, o, layer.bn_node, 1) |
| got = np.einsum("nkc,okc->no", x, layer.weight.astype(np.float64).reshape(W.shape[0], -1, W.shape[-1])) \ |
| + layer.bias |
| np.testing.assert_allclose(got, ref, rtol=2e-6, atol=2e-5) |
| |
| d0 = w.deblocks[0] |
| W = np.asarray(o.conv(d0.onnx_node).weight, np.float64)[:, :, 0, 0] |
| x = rng.normal(size=(5, W.shape[1])) |
| np.testing.assert_allclose(x @ d0.weight[:, :, 0, 0].astype(np.float64).T + d0.bias, |
| _bn_apply(x @ W.T, o, d0.bn_node, 1), rtol=2e-6, atol=2e-5) |
| |
| d1 = w.deblocks[1] |
| W = np.asarray(o.conv(d1.onnx_node).weight, np.float64) |
| x = rng.normal(size=(5, W.shape[0])) |
| np.testing.assert_allclose(x @ d1.weight[:, :, 1, 0].astype(np.float64) + d1.bias, |
| _bn_apply(x @ W[:, :, 1, 0], o, d1.bn_node, 1), rtol=2e-6, atol=2e-5) |
| name, hid, _ = w.head.prediction_heads[0] |
| c = o.conv(hid.onnx_node) |
| x = rng.normal(size=(5, 128)) |
| y = x @ np.asarray(c.weight, np.float64)[:, :, 0].T + (0 if c.bias is None else c.bias) |
| np.testing.assert_allclose(x @ hid.weight[:, :, 0].astype(np.float64).T + hid.bias, |
| _bn_apply(y, o, hid.bn_node, 1), rtol=2e-6, atol=2e-5) |
|
|
|
|
| def test_onnx_taps_name_real_tensors(loaded): |
| _, w, o = loaded |
| produced = {t for node in o.nodes() for t in node.outputs} |
| missing = {k: v for k, v in w.onnx_taps.items() if v not in produced} |
| assert not missing, missing |
|
|
