code/tt_bevformer/reference/weights.py

# SPDX-License-Identifier: Apache-2.0
"""Weights of BEVFormer-small, read the way the TT port reads them.

The source is the user-decided mirror (PLAN.md D13) ``changh95/bevformer-small-weights`` @ 0e6d9d9c: the official
``bevformer_small_epoch_24.pth`` ``state_dict`` as safetensors, every tensor bit-identical to the checkpoint Autoware's
(unobtainable) ONNX was exported from (research/bevformer/SPEC.md section 1.1). It is read as DATA with the vendored
``ttaw.weights.load_state_dict`` (safetensors; a ``.pth`` goes through ``torch.load(weights_only=True)``): nothing is
executed. There is no deployed ONNX graph to read convolution attributes from, so strides and paddings come from the
architecture config (``reference.config.backbone_conv_specs``: caffe-style ResNet-101).

BatchNorm folding happens HERE, exactly where the port folds it: every backbone BatchNorm (frozen, eps 1e-5) into the
convolution before it, in float64 with one final rounding to float32 (``ttaw.weights.fold_bn_conv``); the DCNv2
main convolutions have no bias, so the folded BatchNorm becomes their bias. :meth:`BEVFormerWeights.conv` returns a
``ttaw.models.resnet.ConvParams`` (the C27 provider contract), so the CPU reference and the device ResNet consume the
same folded tensors. Nothing else is folded: the neck convs carry their own biases, LayerNorms stay LayerNorms.

:meth:`BEVFormerWeights.check_keys` is the strict-load check: the state dict must hold exactly the 1,001 tensors this
port consumes (897 float32 parameters and buffers + 104 int64 ``num_batches_tracked``), no more, no fewer.
"""
from __future__ import annotations

from pathlib import Path
from typing import Dict, List, Optional, Set, Union

import numpy as np

from ..ttaw.models.resnet import ConvParams
from ..ttaw.weights import StateDict, file_sha256, fold_bn_conv, load_state_dict
from . import config as C

__all__ = ["WEIGHTS_FILE", "WEIGHTS_SHA256", "SOURCE_PTH_SHA256", "NUM_TENSORS", "BEVFormerWeights", "find_weights",
           "expected_keys"]

WEIGHTS_FILE = "bevformer_small_epoch_24.safetensors"
# sha256 of the mirror file (HF LFS sha256 == the local conversion, research/bevformer/scripts/make_weights_mirror.py)
WEIGHTS_SHA256 = "51ba31289d85df5b90da32126bea21f284b03f08521323f6765459f1663c95d7"
# sha256 of the original checkpoint (github.com/zhiqi-li/storage v1.0; byte-identical to DerryHub's copy)
SOURCE_PTH_SHA256 = "4ebc5810201ca1e609c29452f07c08bc4d9c9ebfda5d7f8ee1034dd446274bc5"
NUM_TENSORS = 1001

PathLike = Union[str, Path]


def find_weights(path: PathLike) -> Optional[Path]:
    """The weights file inside a directory (``WEIGHTS_FILE``, else the official ``.pth``) or the file itself."""
    p = Path(path)
    if p.is_file():
        return p
    for name in (WEIGHTS_FILE, "bevformer_small_epoch_24.pth"):
        if (p / name).is_file():
            return p / name
    return None


def _bn_keys(m: str) -> List[str]:
    return [f"{m}.{k}" for k in ("weight", "bias", "running_mean", "running_var", "num_batches_tracked")]


def _linear(m: str) -> List[str]:
    return [f"{m}.weight", f"{m}.bias"]


def expected_keys() -> Set[str]:
    """Every tensor name of the checkpoint this port consumes (derived from the architecture, not from the file)."""
    keys: List[str] = []
    for module, kind, _, _, bn in C.backbone_conv_specs():
        keys.append(f"{module}.weight")
        if kind == "offset":
            keys.append(f"{module}.bias")
        if bn:
            keys += _bn_keys(bn)
    keys += _linear("img_neck.lateral_convs.0.conv") + _linear("img_neck.fpn_convs.0.conv")
    h = "pts_bbox_head"
    t = f"{h}.transformer"
    for i in range(C.ENC_LAYERS):
        e = f"{t}.encoder.layers.{i}"
        for name in ("sampling_offsets", "attention_weights", "value_proj", "output_proj"):
            keys += _linear(f"{e}.attentions.0.{name}")
        for name in ("sampling_offsets", "attention_weights", "value_proj"):
            keys += _linear(f"{e}.attentions.1.deformable_attention.{name}")
        keys += _linear(f"{e}.attentions.1.output_proj")
        keys += _linear(f"{e}.ffns.0.layers.0.0") + _linear(f"{e}.ffns.0.layers.1")
        for n in range(3):
            keys += _linear(f"{e}.norms.{n}")
    for i in range(C.DEC_LAYERS):
        d = f"{t}.decoder.layers.{i}"
        keys += [f"{d}.attentions.0.attn.in_proj_weight", f"{d}.attentions.0.attn.in_proj_bias"]
        keys += _linear(f"{d}.attentions.0.attn.out_proj")
        for name in ("sampling_offsets", "attention_weights", "value_proj", "output_proj"):
            keys += _linear(f"{d}.attentions.1.{name}")
        keys += _linear(f"{d}.ffns.0.layers.0.0") + _linear(f"{d}.ffns.0.layers.1")
        for n in range(3):
            keys += _linear(f"{d}.norms.{n}")
    keys += [f"{t}.level_embeds", f"{t}.cams_embeds"] + _linear(f"{t}.reference_points")
    keys += _linear(f"{t}.can_bus_mlp.0") + _linear(f"{t}.can_bus_mlp.2") + _linear(f"{t}.can_bus_mlp.norm")
    for i in range(C.DEC_LAYERS):
        for j in (0, 1, 3, 4, 6):          # Linear, LN, (ReLU), Linear, LN, (ReLU), Linear
            keys += _linear(f"{h}.cls_branches.{i}.{j}")
        for j in (0, 2, 4):                # Linear, (ReLU), Linear, (ReLU), Linear
            keys += _linear(f"{h}.reg_branches.{i}.{j}")
    keys += [f"{h}.positional_encoding.row_embed.weight", f"{h}.positional_encoding.col_embed.weight",
             f"{h}.bev_embedding.weight", f"{h}.query_embedding.weight", f"{h}.code_weights"]
    return set(keys)


class BEVFormerWeights:
    """The checkpoint as numpy, with the port's BatchNorm folding. ``BEVFormerWeights(path)`` accepts the mirror's
    ``.safetensors`` (or a directory holding it) and the official ``.pth``; ``verify_sha256=True`` checks the file
    against the pinned digest (slow: 238 MB)."""

    def __init__(self, path: PathLike, *, verify_sha256: bool = False, check: bool = True):
        p = find_weights(path)
        if p is None:
            raise FileNotFoundError(f"no {WEIGHTS_FILE} (or bevformer_small_epoch_24.pth) at {path}")
        self.path = p
        if p.suffix == ".pth":
            self.sd: StateDict = load_state_dict(p, key="state_dict")
        else:
            self.sd = load_state_dict(p)
        self._sha256: Optional[str] = None
        if verify_sha256:
            want = SOURCE_PTH_SHA256 if p.suffix == ".pth" else WEIGHTS_SHA256
            if self.sha256 != want:
                raise ValueError(f"{p} has sha256 {self.sha256}, expected {want}")
        if check:
            self.check_keys()
        self._conv_cache: Dict[str, ConvParams] = {}

    @property
    def sha256(self) -> str:
        if self._sha256 is None:
            self._sha256 = file_sha256(self.path)
        return self._sha256

    def check_keys(self) -> None:
        """Strict coverage: the file holds exactly :func:`expected_keys` (1,001 tensors)."""
        have = set(self.sd)
        want = expected_keys()
        missing, extra = sorted(want - have), sorted(have - want)
        if missing or extra or len(have) != NUM_TENSORS:
            raise ValueError(f"{self.path}: state dict mismatch (missing {missing[:8]}, unexpected {extra[:8]}, "
                             f"{len(have)} tensors, expected {NUM_TENSORS})")

    def array(self, name: str) -> np.ndarray:
        """A float32 copy of one tensor."""
        return np.array(self.sd[name], dtype=np.float32)

    def conv(self, module: str) -> ConvParams:
        """One backbone convolution as the C27 provider returns it: BatchNorm folded in float64 (one rounding to
        float32), stride / padding from the caffe-style ResNet-101 config. DCN ``conv2`` modules return the folded
        main weight [C, C, 3, 3] + bias; ``<module>.conv_offset`` returns the offset conv [27, C, 3, 3] + bias."""
        hit = self._conv_cache.get(module)
        if hit is not None:
            return hit
        specs = {m: (k, s, p, bn) for m, k, s, p, bn in C.backbone_conv_specs()}
        if module not in specs:
            raise KeyError(f"{module!r} is not a backbone convolution")
        kind, stride, pad, bn = specs[module]
        w = np.asarray(self.sd[f"{module}.weight"])
        if kind == "offset":
            wf, bf = w.astype(np.float32), np.asarray(self.sd[f"{module}.bias"], np.float32)
        else:
            wf, bf = fold_bn_conv(w, None, self.sd.bn(bn, C.BN_EPS))
        out = ConvParams(np.ascontiguousarray(wf, np.float32), np.ascontiguousarray(bf, np.float32),
                         stride=(stride, stride), padding=(pad, pad, pad, pad))
        self._conv_cache[module] = out
        return out

    def provider(self):
        """``module -> ConvParams``: the C27 ``ttaw.models.resnet`` provider of the backbone (DCN modules included:
        a DCN factory consumes ``conv(m)`` and ``conv(m + '.conv_offset')``)."""
        return self.conv

    def dcn_modules(self) -> List[str]:
        """The 26 DCNv2 ``conv2`` modules (layer3: 23, layer4: 3), in graph order."""
        return [m for m, kind, _, _, _ in C.backbone_conv_specs() if kind == "dcn"]