| |
| """C25: the CenterPoint dense head on the device (PLAN.md section 1.2 row C25; CenterPoint, PointPainting). |
| |
| mmdet3d ``CenterHead`` (single task, as Autoware's exports): ``shared = relu(conv3x3(concat[d_0..d_n]))``, then per |
| output head ``out_h = conv1x1(relu(conv1x1(shared)))``. Two exact rewrites (each differs from the literal graph |
| only by summation order; probe P14 and research/centerpoint/SPEC.md 8.1): |
| |
| - **K-split shared conv** (default; C17 ``KSplitConv``): ``conv(concat[d_i], W) == sum_i conv(d_i, W_i)``, the bias |
| on the first partial conv and the ReLU fused into the last add. The concat (CenterPoint base: 177 MB) never |
| exists; P14: 3.6 ms against 14.8 ms for the full-K conv, and more accurate. ``accumulate_dtype="float32"`` keeps |
| the partials in fp32; |
| - **merged sibling heads**: the six 1x1 -> ReLU -> 1x1 heads become two matmuls on NHWC rows (:class:`..second. |
| RowLinear`): ``hidden = relu(x @ Wh + bh)`` with ``Wh = [Wh_1 | ... | Wh_6]`` (``64 -> 6*64``), then ``out = |
| hidden @ Wo + bo`` with the block-diagonal ``Wo`` (``6*64 -> sum C_h``, zero-padded to a multiple of 32 columns so |
| the output tensor is tile-exact). |
| |
| The output is ONE ``[1, 1, H*W, C_pad]`` tensor (fp32 by default: the logits feed host thresholds), i.e. one |
| readback; :meth:`CenterHead.unpack` splits the host rows into ``{name: (C, H, W)}``. Module names for the precision |
| policy: ``<prefix>.shared``, ``<prefix>.hidden``, ``<prefix>.out``. |
| """ |
| from __future__ import annotations |
|
|
| from typing import Any, Dict, List, Optional, Sequence, Tuple |
|
|
| import numpy as np |
|
|
| from ..ops.conv import Conv2d, FeatureMap, KSplitConv |
| from ..precision import PrecisionPolicy |
| from .second import ConvSpec, RowLinear, _precision, _to_dram, free_maps |
|
|
| __all__ = ["CenterHead", "HeadSpec", "merge_heads", "centerhead_numpy"] |
|
|
| HeadSpec = Tuple[str, ConvSpec, ConvSpec] |
|
|
|
|
| def merge_heads(heads: Sequence[HeadSpec], *, pad_to: int = 32 |
| ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, List[Tuple[str, int, int]]]: |
| """``(Wh [Cin, sum hidden], bh, Wo [sum hidden, C_pad], bo [C_pad], slices)`` of the merged heads; ``slices`` |
| lists ``(name, start, stop)`` of each head's output columns; columns past ``sum C_h`` are zero.""" |
| if not heads: |
| raise ValueError("no heads") |
| cin = heads[0][1].in_channels |
| wh, bh, rows, cols = [], [], 0, 0 |
| for name, hidden, final in heads: |
| if hidden.kind != "conv" or final.kind != "conv" or hidden.kernel != 1 or final.kernel != 1: |
| raise ValueError(f"head {name!r}: merged heads need 1x1 convs") |
| if hidden.in_channels != cin or final.in_channels != hidden.out_channels: |
| raise ValueError(f"head {name!r}: channel mismatch") |
| if not hidden.relu or final.relu: |
| raise ValueError(f"head {name!r}: expected a ReLU after the hidden conv only") |
| wh.append(np.asarray(hidden.weight, np.float32)[:, :, 0, 0].T) |
| bh.append(np.zeros(hidden.out_channels, np.float32) if hidden.bias is None |
| else np.asarray(hidden.bias, np.float32)) |
| rows += hidden.out_channels |
| cols += final.out_channels |
| cpad = -(-cols // pad_to) * pad_to |
| wo = np.zeros((rows, cpad), np.float32) |
| bo = np.zeros(cpad, np.float32) |
| slices, r, c = [], 0, 0 |
| for name, hidden, final in heads: |
| co = final.out_channels |
| wo[r:r + hidden.out_channels, c:c + co] = np.asarray(final.weight, np.float32)[:, :, 0, 0].T |
| if final.bias is not None: |
| bo[c:c + co] = np.asarray(final.bias, np.float32) |
| slices.append((name, c, c + co)) |
| r += hidden.out_channels |
| c += co |
| return (np.ascontiguousarray(np.concatenate(wh, axis=1)), np.concatenate(bh).astype(np.float32), wo, bo, |
| slices) |
|
|
|
|
| class CenterHead: |
| """Shared conv (K-split over ``split``) + merged heads; see the module docstring. |
| |
| ``shared``: the 3x3 :class:`ConvSpec` over the concatenated inputs (ReLU after); ``split``: channels of each |
| input, in concat order; ``heads``: ``(name, hidden, final)`` in output order. ``k_split=False`` runs one conv |
| on a device concat instead (A/B only). ``output_dtype``: of the packed head output (``"float32"``). The shared |
| conv's output dtype is its precision's ``activations`` (bf16 by default); ``accumulate_dtype`` (default: the same) |
| is the dtype of the K-split partial sums.""" |
|
|
| def __init__(self, shared: ConvSpec, split: Sequence[int], heads: Sequence[HeadSpec], *, |
| policy: Optional[PrecisionPolicy] = None, prefix: str = "head", output_dtype: str = "float32", |
| k_split: bool = True, accumulate_dtype: Optional[str] = None): |
| if shared.kind != "conv" or shared.stride != 1 or not shared.relu: |
| raise ValueError("the shared conv must be a stride-1 conv followed by a ReLU") |
| self.prefix = prefix |
| self.split = [int(s) for s in split] |
| if sum(self.split) != shared.in_channels: |
| raise ValueError(f"split {self.split} does not cover the {shared.in_channels} shared-conv inputs") |
| self.k_split = bool(k_split) and len(self.split) > 1 |
| pad = shared.kernel // 2 if shared.padding is None else int(shared.padding) |
| prec = _precision(policy, f"{prefix}.shared") |
| if self.k_split: |
| self.shared = KSplitConv(shared.weight, shared.bias, self.split, activation="relu", padding=pad, |
| precision=prec, accumulate_dtype=accumulate_dtype or prec.activations, |
| output_dtype=prec.activations, name=f"{prefix}.shared") |
| else: |
| self.shared = Conv2d(shared.weight, shared.bias, padding=pad, activation="relu", precision=prec, |
| output_dtype=prec.activations, name=f"{prefix}.shared") |
| wh, bh, wo, bo, self.slices = merge_heads(heads) |
| self.out_channels = self.slices[-1][2] |
| self.padded_channels = int(wo.shape[1]) |
| hidden_prec = _precision(policy, f"{prefix}.hidden") |
| self.hidden = RowLinear(wh, bh, activation="relu", precision=hidden_prec, |
| output_dtype=hidden_prec.activations, name=f"{prefix}.hidden") |
| self.out = RowLinear(wo, bo, activation=None, precision=_precision(policy, f"{prefix}.out"), |
| output_dtype=output_dtype, name=f"{prefix}.out") |
|
|
| def shared_forward(self, inputs: Sequence[FeatureMap]) -> FeatureMap: |
| """``relu(conv3x3(concat(inputs)))`` (inputs are kept), DRAM interleaved.""" |
| import ttnn |
|
|
| if [x.channels for x in inputs] != self.split: |
| raise ValueError(f"{self.prefix}: inputs have {[x.channels for x in inputs]} channels, " |
| f"expected {self.split}") |
| if self.k_split: |
| return _to_dram(self.shared(inputs)) |
| cat = inputs[0].with_tensor(ttnn.concat([x.tensor for x in inputs], dim=-1), sum(self.split)) |
| out = _to_dram(self.shared(cat)) |
| free_maps(cat) |
| return out |
|
|
| def heads_forward(self, shared: FeatureMap) -> FeatureMap: |
| """Merged heads on the shared map -> ``[1, 1, H*W, C_pad]`` (``channels`` = the padded width).""" |
| hidden = self.hidden(shared.tensor) |
| out = self.out(hidden) |
| free_maps(hidden) |
| return shared.with_tensor(out, self.padded_channels) |
|
|
| def __call__(self, inputs: Sequence[FeatureMap]) -> FeatureMap: |
| shared = self.shared_forward(inputs) |
| out = self.heads_forward(shared) |
| free_maps(shared) |
| return out |
|
|
| def unpack(self, rows: Any, height: int, width: int) -> Dict[str, np.ndarray]: |
| """Host rows ``[H*W, >= C]`` (a read-back of the head output) -> ``{name: (C_h, H, W)}`` float32.""" |
| a = np.asarray(rows, dtype=np.float32) |
| a = a.reshape(-1, a.shape[-1])[: height * width] |
| return {name: np.ascontiguousarray(a[:, s:e].T.reshape(e - s, height, width)) for name, s, e in self.slices} |
|
|
| def release(self) -> None: |
| self.shared.release() |
| self.hidden.release() |
| self.out.release() |
|
|
| def describe(self) -> Dict[str, Any]: |
| return {"shared": self.shared.describe(), "k_split": self.k_split, "hidden": self.hidden.describe(), |
| "out": self.out.describe(), "slices": [list(s) for s in self.slices]} |
|
|
|
|
| def centerhead_numpy(shared: ConvSpec, heads: Sequence[HeadSpec], inputs_nchw: Sequence[Any]) -> Dict[str, np.ndarray]: |
| """fp32 oracle (torch, the literal graph: concat, conv, six heads) -> ``{"shared": NCHW, name: NCHW}``.""" |
| import torch |
|
|
| from .second import _apply_torch |
|
|
| with torch.no_grad(): |
| x = torch.cat([torch.as_tensor(np.asarray(t, dtype=np.float32)) for t in inputs_nchw], dim=1) |
| s = _apply_torch(shared, x) |
| out = {"shared": s.numpy().copy()} |
| for name, hidden, final in heads: |
| out[name] = _apply_torch(final, _apply_torch(hidden, s)).numpy().copy() |
| return out |
|
|
