← Walkthroughmapa/encoder.py
Source snapshot · source SHA-256 92e8a2cb388d
1"""Frozen encoder API for per-contact feature extraction from intracranial EEG.23Main path: prepare() builds metadata; forward() encodes windows and returns taps.45This is the API wrapper, not another neural-network stage.6prepare(): contact metadata + clip length -> reusable token-layout metadata.7forward(): normalized clips + layout -> tower -> features grouped by contact.89Inputs are three normalized magnitude STFT tensors and a SensorSidecar with the array identity,10clinical contact number, and anatomical region of each contact. Outputs are keyed by encoder11block. Block 0 returns the frontend baseline features; block 12 returns the final encoder features.12"""13from __future__ import annotations1415import os16from dataclasses import dataclass17from numbers import Integral18from typing import Sequence1920import torch21from torch import Tensor, nn2223from mapa.models.geometry import build_l1_geometry24from mapa.models.pack_r4 import BAND_STRIDES, build_r4_grid25from mapa.models.sidecar import SensorSidecar26from mapa.models.stem import clip_band27from mapa.models.tower import MaeEncoder2829N_REGIONS = 75 # DKT K=74 plus one reserved id for contacts outside every region30FPS = 32.0 # the uniform frame clock every band is resampled onto31DEPTH = 12 # encoder blocks; 3/6/9/12 carry deep-supervision norms323334@dataclass(frozen=True)35class SessionGrid:36 """Per-session token layout. Clip-independent, so build it once and reuse it."""3738 grid: object39 region_packed: Tensor40 contact_order: Tensor # (n,) indices into the sidecar row axis, in grid order41 region_of_contact: Tensor # (n,) DKT id per contact, grid order42 k_full: int43 n_time: int444546class MapaEncoder(nn.Module):47 """The frozen encoder that ships as the representation.4849 Pretraining is masked autoencoding. The decoder that reconstructs the removed patches is50 discarded when pretraining ends, so this is the whole released model. Loading reads the51 architecture off the checkpoint rather than trusting arguments, so a mismatched shell52 fails loud.53 """5455 def __init__(self, tower: MaeEncoder, *, space_rope: bool, region_embed: bool) -> None:56 super().__init__()57 self.tower = tower58 self.space_rope = bool(space_rope)59 self.region_embed = bool(region_embed)6061 @property62 def d_model(self) -> int:63 return int(self.tower.d_model)6465 @classmethod66 def from_checkpoint(67 cls,68 path: str,69 *,70 device: str | torch.device = "cpu",71 space_rope: bool | None = None,72 ) -> MapaEncoder:73 """Load a released checkpoint.7475 A released file records whether the relative positional encoding is active, because76 that setting cannot be recovered from the weights: disabling it zeroes a buffer that77 is not persisted, so the state-dict layout alone does not identify that ablation.78 Pass ``space_rope`` only to override what the file records, or to read a raw training79 checkpoint, which records nothing.80 """81 if not os.path.exists(path):82 raise FileNotFoundError(83 f"no checkpoint at {path!r}. Pass an absolute path, or set MAPA_WEIGHTS to the "84 "directory holding the released checkpoints and load through torch.hub."85 )86 raw = torch.load(path, map_location="cpu", weights_only=True)87 sub, recorded = _tower_state(raw)88 if space_rope is None:89 if recorded is None:90 raise ValueError(91 f"{path} does not record whether the relative positional encoding is "92 "active, so it cannot be inferred. Pass space_rope explicitly."93 )94 space_rope = recorded9596 peek = [v.shape[0] for kk, v in sub.items() if kk.endswith("region_embed.embed.weight")]97 if peek and int(peek[0]) != N_REGIONS:98 raise ValueError(f"checkpoint region table {peek[0]} != expected {N_REGIONS}")99 dkey = "encoder.blocks.0.norm1.weight"100 if dkey not in sub:101 raise RuntimeError(f"{path} has no '{dkey}'; cannot infer encoder width")102103 region_embed = bool(peek)104 deep_sup = any(kk.startswith("encoder.norms_block.") for kk in sub)105 d_model = int(sub[dkey].shape[0])106107 shell = MaeEncoder(108 n_regions=N_REGIONS, deep_sup=deep_sup, region_embed=region_embed,109 space_rope=space_rope, d_model=d_model,110 )111 missing, unexpected = shell.load_state_dict(sub, strict=False)112 bad = [m for m in missing if "num_batches_tracked" not in m]113 if bad or unexpected:114 raise RuntimeError(115 f"state_dict mismatch: missing={bad[:8]} unexpected={unexpected[:8]}"116 )117 shell.eval().to(device)118 for p in shell.parameters():119 p.requires_grad_(False)120 return cls(shell, space_rope=space_rope, region_embed=region_embed)121122 def prepare(self, sidecar: SensorSidecar, *, n_time: int) -> SessionGrid:123 """Build the token layout for one session. ``n_time`` is the window length in frames124 on the 32 Hz clock, so a one second window is 32."""125 if isinstance(n_time, bool) or not isinstance(n_time, Integral) or n_time <= 0:126 raise ValueError("n_time must be a positive integer on the 32 Hz frame clock")127 device = next(self.tower.parameters()).device128 # Metadata only: group contacts, then expand contact/band/time per token.129 geom = build_l1_geometry(sidecar).to(device)130 grid = build_r4_grid(geom, n_time=int(n_time))131 region_id = sidecar.region_id.to(device)132 contact_order = grid.contact.reshape(-1, grid.k_full)[:, 0]133 return SessionGrid(134 grid=grid,135 region_packed=region_id[grid.contact],136 contact_order=contact_order.cpu(),137 region_of_contact=region_id[contact_order].cpu(),138 k_full=int(grid.k_full),139 n_time=int(n_time),140 )141142 @torch.no_grad()143 def forward(144 self,145 bands: Sequence[Tensor],146 session: SessionGrid,147 *,148 taps: Sequence[int] = (DEPTH,),149 ) -> dict[int, Tensor]:150 """Encode windows and return selected layer features in grid contact order.151152 Taps 1..12 have shape (B, contacts, K, d); all requested layers share153 one encoder pass. Tap 0 flattens each contact's subsampled bins into154 (B, contacts, 1, sum(F_b*T_b)); it skips the learned stem and encoder.155 """156 taps = tuple(taps)157 if not taps or any(158 isinstance(t, bool) or not isinstance(t, Integral) or not 0 <= t <= DEPTH159 for t in taps160 ):161 raise ValueError(f"taps must contain integer block numbers from 0 to {DEPTH}")162 _validate_bands(bands, session)163 device = next(self.tower.parameters()).device164 if session.grid.contact.device != device or bands[0].device != device:165 raise ValueError(166 "encoder, session and bands must be on the same device; move the encoder first, "167 "then call prepare() and move the bands to that device"168 )169 out: dict[int, Tensor] = {}170 gpu_taps = tuple(t for t in taps if t != 0)171 dtype = next(self.tower.parameters()).dtype172 if gpu_taps and bands[0].dtype != dtype and (173 torch.float64 in (bands[0].dtype, dtype) or not torch.is_autocast_enabled(device.type)174 ):175 raise ValueError(176 f"bands use {bands[0].dtype}, encoder uses {dtype}; convert bands with "177 f"[band.to(dtype={dtype}) for band in bands]"178 )179 if 0 in taps:180 out[0] = stem_input(bands, session)181 if gpu_taps:182 # Use requested raw layer taps (width d), not the concatenated 4d output _z.183 _z, tapped = self.tower.forward(184 list(bands), session.grid, session.region_packed, tap_blocks=gpu_taps185 )186 n = session.contact_order.shape[0]187 for t in gpu_taps:188 x = tapped[t]189 out[t] = x.reshape(x.shape[0], n, session.k_full, x.shape[-1])190 return {t: out[t] for t in taps}191192193# The region embedding was called ``parcel_embed`` before the rename, and checkpoints194# written then carry that key. Both spellings load.195_LEGACY_KEYS = (("parcel_embed", "region_embed"),)196197198def _canonical_keys(sd: dict) -> dict:199 out = {}200 for k, v in sd.items():201 for old, new in _LEGACY_KEYS:202 if old in k:203 k = k.replace(old, new)204 out[k] = v205 return out206207208def _tower_state(raw) -> tuple[dict, bool | None]:209 """Pull the encoder weights out of either checkpoint layout.210211 A released file is ``{"model": ..., "meta": ...}`` and carries bare keys. A raw training212 checkpoint carries the whole autoencoder, of which the encoder sits under213 ``objective.online.``; the decoder is training only and is dropped here.214 """215 if isinstance(raw, dict) and "model" in raw and isinstance(raw["model"], dict):216 meta = raw.get("meta") or {}217 rope = meta.get("space_rope")218 return _canonical_keys(raw["model"]), (None if rope is None else bool(rope))219220 sd = raw["state_dict"] if "state_dict" in raw else raw221 pref = "objective.online."222 sub = {}223 for k, v in sd.items():224 kk = k.replace("_orig_mod.", "")225 if kk.startswith("model."):226 kk = kk[len("model."):]227 if kk.startswith(pref):228 sub[kk[len(pref):]] = v229 if not sub:230 raise RuntimeError("no encoder weights found; not a released checkpoint layout")231 return _canonical_keys(sub), None232233234def stem_input(bands: Sequence[Tensor], session: SessionGrid) -> Tensor:235 """Tap 0: the exact tensor the stem's first linear layer consumes, in grid contact order.236237 Each band is decimated by its own stride onto the token lattice, made time-major, reordered238 to the grid contact axis, and concatenated. A readout on this is the input-linear floor: the239 same clips, the same contacts, the same pooling, no encoder.240 """241 _validate_bands(bands, session)242 canon = session.contact_order243 per_band = []244 for b, (x, stride) in enumerate(zip(bands, BAND_STRIDES)):245 xd = clip_band(x[..., ::stride], b).transpose(-1, -2).contiguous()246 xd = xd[:, canon]247 per_band.append(xd.reshape(xd.shape[0], xd.shape[1], -1))248 return torch.cat(per_band, dim=-1)[:, :, None, :]249250251def _validate_bands(bands: Sequence[Tensor], session: SessionGrid) -> None:252 """Reject mismatched axes before indexing can silently discard contacts or frames."""253 if len(bands) != 3:254 raise ValueError("provide three bands in Slow, Mid, Fast order")255 n = len(session.contact_order)256 for name, bins, band in zip(("Slow", "Mid", "Fast"), (7, 6, 7), bands):257 expected = (n, bins, session.n_time)258 if not isinstance(band, Tensor) or band.ndim != 4 or tuple(band.shape[1:]) != expected:259 raise ValueError(f"{name} band must have shape (batch, {n}, {bins}, {session.n_time})")260 if not band.is_floating_point():261 raise ValueError(f"{name} band must be a floating-point tensor")262 if band.shape[0] < 1 or band.shape[0] != bands[0].shape[0]:263 raise ValueError("all bands must have the same nonempty batch axis")264 if band.device != bands[0].device or band.dtype != bands[0].dtype:265 raise ValueError("all bands must use the same device and dtype")266267268def pool_to_regions(x: Tensor, session: SessionGrid) -> Tensor:269 """Average contacts within each region: ``(B, n, ...)`` to ``(B, n_regions, F)``.270271 Regions are returned in sorted ID order, as given by ``session.region_of_contact.unique()``.272 Align the region IDs shared by two sessions before fitting a cross-subject readout.273 """274 regions = session.region_of_contact275 present = torch.unique(regions)276 blocks = [x[:, regions == p].mean(1).reshape(x.shape[0], -1) for p in present.tolist()]277 return torch.stack(blocks, dim=1)