File size: 9,074 Bytes
ddaad99 | 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 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 | # 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"]
|