File size: 9,078 Bytes
4d9b003 | 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 | # SPDX-License-Identifier: Apache-2.0
"""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] # (output name, hidden 1x1 + ReLU, final 1x1)
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
|