| |
| """The CPU reference reads the weights the way the spec describes them (S:transfusion:57-208): the Autoware |
| parameters, and the module structure checked from the ONNX graph (layer counts, strides, channels, the decoder's |
| heads / scale / LayerNorm eps, the in-graph decoration constants, ``bev_pos``, parameter counts after the BN folds), |
| without onnxruntime or torch. |
| |
| TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_transfusion/tests/test_reference_structure_host.py |
| """ |
| from __future__ import annotations |
|
|
| import numpy as np |
| import pytest |
|
|
| pytest.importorskip("onnx") |
| pytest.importorskip("yaml") |
|
|
| from tt_transfusion.reference.config import (N_MAX, NODE_PARAMS, OUT_SIZE_FACTOR, POINTS_PER_VOXEL, |
| TransFusionConfig) |
| from tt_transfusion.reference.weights import PREDICTION_HEADS, load_weights |
| from tt_transfusion.tests import _data |
|
|
|
|
| @pytest.fixture(scope="module") |
| def config(): |
| return TransFusionConfig.from_weights(_data.weights_dir()) |
|
|
|
|
| @pytest.fixture(scope="module") |
| def weights(config): |
| return load_weights(_data.weights_dir(), config) |
|
|
|
|
| def test_autoware_parameters(config): |
| c = config |
| assert c.version == "v2.1" and c.class_names == ("CAR", "TRUCK", "BUS", "BICYCLE", "PEDESTRIAN") |
| assert c.voxels_num == (5000, 30000, 60000) and (c.min_voxels, c.max_voxels) == (5000, 60000) |
| assert c.point_cloud_range == (-92.16, -92.16, -3.0, 92.16, 92.16, 7.0) and c.voxel_size == (0.24, 0.24, 10.0) |
| assert (c.grid_x, c.grid_y, c.grid_z) == (768, 768, 1) |
| assert (c.feature_x, c.feature_y) == (192, 192) and OUT_SIZE_FACTOR == 4 |
| assert c.num_proposals == 500 and POINTS_PER_VOXEL == 20 and N_MAX == 65536 |
| spec = c.grid_spec |
| assert spec.order == "raster" and spec.max_pillars == N_MAX and spec.max_points_per_pillar == 20 |
| assert c.score_threshold == NODE_PARAMS["score_threshold"] == 0.1 |
| assert c.yaw_norm_thresholds.tolist() == pytest.approx([0.3, 0.3, 0.3, 0.3, 0.0]) |
| assert c.with_node(yaw_norm_thresholds=(1.5, -1, 0.3, 0.999, 1.0)).yaw_norm_thresholds.tolist() == pytest.approx( |
| [0.0, 0.0, 0.3, 0.999, 0.0]) |
| assert TransFusionConfig.sanitized_score_threshold(0.0) == 0.2 and TransFusionConfig.sanitized_score_threshold(-1) == 0.2 |
| assert TransFusionConfig.sanitized_circle_nms(0.0) == 0.5 and TransFusionConfig.sanitized_circle_nms(1.2) == 1.2 |
| rm = c.remapper |
| assert rm.num_labels == 8 and rm.allow[1, 2] and rm.allow[1, 4] and rm.allow[2, 4] and rm.allow[3, 4] |
| assert int(rm.allow.sum()) == 4 |
| assert (rm.min_area[1, 2], rm.max_area[1, 2], rm.min_area[3, 4], rm.max_area[3, 4]) == (12.1, 36.0, 36.0, 999.999) |
| assert c.decoration_scale_zyx.tolist() == [np.float32(10.0), np.float32(0.24), np.float32(0.24)] |
| assert c.decoration_offset_zyx.tolist() == [np.float32(2.0), np.float32(-92.04), np.float32(-92.04)] |
|
|
|
|
| def test_refuses_another_release(tmp_path, config): |
| import shutil |
|
|
| for f in ("transfusion_ml_package.param.yaml", "detection_class_remapper.param.yaml"): |
| shutil.copy(_data.weights_dir() / f, tmp_path / f) |
| (tmp_path / "deploy_metadata.yaml").write_text("version: v2\n") |
| with pytest.raises(ValueError, match="not supported"): |
| TransFusionConfig.from_model_dir(tmp_path) |
|
|
|
|
| def test_graph_structure_and_parameter_counts(weights): |
| w = weights |
| assert [len(b) for b in w.blocks] == [4, 6, 6] |
| assert [(b[0].in_channels, b[-1].out_channels, b[0].stride) for b in w.blocks] == [(64, 64, 2), (64, 128, 2), |
| (128, 256, 2)] |
| assert all(layer.kernel == 3 and layer.padding == 1 and layer.relu for b in w.blocks for layer in b) |
| assert [(d.kind, d.in_channels, d.out_channels, d.kernel, d.stride, d.padding) for d in w.deblocks] == [ |
| ("conv", 64, 128, 2, 2, 0), ("conv", 128, 128, 1, 1, 0), ("conv_transpose", 256, 128, 2, 2, 0)] |
| assert [d.bn_node is not None for d in w.deblocks] == [False, False, True] |
| h = w.head |
| assert (h.shared.in_channels, h.shared.out_channels, h.shared.kernel, h.shared.relu) == (384, 128, 3, False) |
| assert (h.heatmap_hidden.out_channels, h.heatmap_hidden.relu, h.heatmap_out.out_channels) == (128, True, 5) |
| assert h.nms_kernel == 3 and h.num_proposals == 500 |
| d = h.decoder |
| for a in (d.self_attn, d.cross_attn): |
| assert (a.num_heads, a.head_dim, a.q_divisor, a.scale) == (8, 16, 4.0, 0.25) |
| assert a.wq.shape == a.wk.shape == a.wv.shape == a.wo.shape == (128, 128) |
| assert [eps for _, _, eps in d.norms] == [pytest.approx(1e-5)] * 3 |
| assert d.ffn_w1.shape == (128, 256) and d.ffn_w2.shape == (256, 128) |
| assert [(n, fin.out_channels) for n, _, fin in h.prediction_heads] == list(PREDICTION_HEADS) |
| assert w.pfn.w0.shape == (11, 32) and w.pfn.w1.shape == (64, 64) and w.grid == (768, 768) |
| |
| |
| assert w.num_params() == 5_270_580 - (4 * 96 - 96) - (4 * 128 - 128) |
|
|
|
|
| def test_constants_and_tables(weights): |
| h = weights.head |
| k = np.arange(192 * 192) |
| assert h.bev_pos.shape == (36864, 2) |
| assert np.array_equal(h.bev_pos[:, 0], (k % 192).astype(np.float32) + 0.5) |
| assert np.array_equal(h.bev_pos[:, 1], (k // 192).astype(np.float32) + 0.5) |
| assert h.bev_pos[1].tolist() == [1.5, 0.5] and h.bev_pos[192].tolist() == [0.5, 1.5] |
| t = h.class_table |
| assert t.shape == (5, 128) and np.array_equal(t[3], (h.class_encoding_w[:, 3] + h.class_encoding_b).astype( |
| np.float32)) |
| p = weights.pfn |
| assert np.array_equal(p.scale_zyx, np.array([10.0, 0.24, 0.24], np.float32)) |
| assert np.array_equal(p.offset_zyx, np.array([2.0, -92.04, -92.04], np.float32)) |
|
|
|
|
| def test_onnx_tap_map_covers_the_reference_taps(weights): |
| from tt_transfusion.reference.model import EXTRA_TAPS, MODULE_TAPS |
|
|
| missing = [t for t in MODULE_TAPS + EXTRA_TAPS if t not in weights.onnx_taps] |
| assert not missing, missing |
| assert weights.onnx_taps["head.topk_indices"].endswith("TopK_output_1") |
| assert weights.onnx_taps["pfn.input"] == "/pts_voxel_encoder/Mul_1_output_0" |
|
|
