code/tt_frnet/tests/test_reference_structure_host.py

# SPDX-License-Identifier: Apache-2.0
"""The deployed ONNX as ``reference/weights.py`` reads it (onnx + the weights; no torch, no device): the structure
recovered from the graph (S:frnet:250-254: 53 convs = 47 x 3x3 s1 + 3 x 3x3 s2 + 3 x 1x1 s2, 17 point layers,
BasicBlocks {3, 4, 6, 3}, strides 1 / 2 / 4 / 8, the frustum grid of both variants), the folds (float64, one
rounding: ``x @ W' + b'`` == BatchNorm(x @ W) for every point layer; ``pre_norm`` == ``x * scale + shift``), the
parameter count (10,016,971 values in the initializers), the pinned file hashes and every module tap's ONNX tensor.

    TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_frnet/tests/test_reference_structure_host.py
"""
from __future__ import annotations

import numpy as np
import pytest

pytest.importorskip("onnx")

from tt_frnet.api import FRNet  # noqa: E402
from tt_frnet.reference.config import BLOCKS, ONNX_SHA256, STRIDES, VARIANTS, FRNetConfig  # noqa: E402
from tt_frnet.reference.weights import MODULE_TAPS, OPT_IN_TAPS, load_weights  # noqa: E402
from tt_frnet.tests import _data  # noqa: E402
from tt_frnet.ttaw.weights import OnnxWeights  # noqa: E402

GRID = {"ot128": (128, 1024), "qt128": (128, 256)}


@pytest.fixture(scope="module", params=VARIANTS)
def loaded(request):
    w = _data.weights_dir()
    return request.param, load_weights(w, request.param), FRNetConfig.from_model_dir(w, request.param), w


def test_structure(loaded):
    variant, W, cfg, _ = loaded
    assert W.frustum_hw == GRID[variant] == (cfg.frustum_h, cfg.frustum_w) and W.strides == STRIDES
    assert tuple(len(s.blocks) for s in W.stages) == BLOCKS
    convs = W.convs
    kinds = sorted((c.kernel, c.stride) for c in convs)
    assert len(convs) == 53 and kinds.count((3, 1)) == 47 and kinds.count((3, 2)) == 3 and kinds.count((1, 2)) == 3
    assert sum(c.activation == "hardswish" for c in convs) == 3 + 1 + 16 + 4 + 4 + 2
    assert sum(c.activation == "sigmoid" for c in convs) == 4
    assert [(lin.in_features, lin.out_features) for lin in W.linears] == [
        (8, 64), (64, 128), (128, 256), (256, 256), (256, 16), (384, 128), (256, 128), (256, 128), (256, 128),
        (256, 128), (640, 256), (256, 128), (128, 128), (128, 256), (256, 128), (128, 64), (64, 27)]
    assert [lin.relu for lin in W.linears] == [True] * 3 + [False, True] + [True] * 11 + [False]
    assert W.raw_parameters == 10_016_971 and W.num_classes == 27
    assert W.onnx_sha256 == ONNX_SHA256[variant]
    assert set(MODULE_TAPS) <= set(W.onnx_taps) and len(W.onnx_taps) == 67
    assert {t for t in W.onnx_taps if t.startswith("backbone.segmax")} == {f"backbone.segmax{i}" for i in range(5)}
    assert len(OPT_IN_TAPS) == 3


def test_folds_are_exact_in_float64(loaded):
    variant, W, _, wdir = loaded
    onnx = OnnxWeights(wdir / f"frnet_{variant}.onnx")
    rng = np.random.default_rng(0)
    for lin in W.linears:
        x = rng.normal(size=(64, lin.in_features))
        ours = x @ lin.weight.astype(np.float64) + lin.bias
        if lin.bn_node is not None:
            raw = x @ onnx.matmul_weight(lin.onnx_node).astype(np.float64)
            ref = onnx.batchnorm(lin.bn_node).apply(raw, axis=1)
        else:
            g = onnx.gemm(lin.onnx_node)
            ref = x @ np.asarray(g.weight, np.float64).T + g.bias
        np.testing.assert_allclose(ours, ref, rtol=2e-6, atol=2e-6 * np.abs(ref).max(), err_msg=lin.name)
    bn = onnx.batchnorm(onnx.find_node("/model/voxel_encoder/pre_norm/*", "BatchNormalization").name)
    x = rng.normal(size=(64, 8)) * 30
    np.testing.assert_allclose(x * W.pre_norm_scale + W.pre_norm_shift, bn.apply(x, axis=1), rtol=1e-6, atol=1e-5)
    assert bn.eps == pytest.approx(1e-3)


def test_both_variants_share_classes_and_differ_in_weights():
    w = _data.weights_dir()
    a, b = (FRNetConfig.from_model_dir(w, v) for v in VARIANTS)
    assert a.class_names == b.class_names == tuple(FRNet.LABELS) and a.palette == b.palette
    assert (a.frustum_w, a.interp_w, a.num_points[2]) == (1024, 4096, 160000)
    assert (b.frustum_w, b.interp_w, b.num_points[2]) == (256, 1024, 120000)
    assert a.num_unique == b.num_unique == (3000, 30000, 60000) and a.version == "v2.0"
    wa, wb = (load_weights(w, v) for v in VARIANTS)
    assert not np.array_equal(wa.stem[0].weight, wb.stem[0].weight)
    assert [c.name for c in wa.convs] == [c.name for c in wb.convs]


def test_a_modified_file_is_refused(tmp_path):
    w = _data.weights_dir()
    for name in ("frnet_ot128.onnx", "ml_package_frnet_ot128.param.yaml", "deploy_metadata.yaml"):
        (tmp_path / name).write_bytes((w / name).read_bytes())
    with open(tmp_path / "frnet_ot128.onnx", "ab") as f:
        f.write(b"\0")
    with pytest.raises(ValueError, match="sha256"):
        load_weights(tmp_path, "ot128")
    (tmp_path / "deploy_metadata.yaml").write_text("version: v1.0\n")
    with pytest.raises(ValueError, match="v1.0"):
        FRNetConfig.from_model_dir(tmp_path, "ot128")