File size: 6,187 Bytes
7cc5047 | 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 | # SPDX-License-Identifier: Apache-2.0
"""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 # noqa: E402
from tt_bevfusion.reference.config import BEVFusionConfig # noqa: E402
from tt_bevfusion.reference.weights import ENCODER_PLAN, load_weights # noqa: E402
from tt_bevfusion.tests import _data # noqa: E402
from tt_bevfusion.ttaw.weights import OnnxWeights # noqa: E402
@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)
# sparse KRSC: y = x_im2col @ W^T per tap, BN over Cout
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)
# dense conv (1x1 deblock as a matmul) and conv1d heads: same fold rule along the output channels
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)
# ConvTranspose k == s: per output pixel a matmul over Cin with W[:, :, i, j]; BN over Cout (axis 1 of W)
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
|