# 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; ``.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"]