Coverage for cafl4ds/models/vit.py: 98%
113 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-17 14:39 +0000
1"""A small, self-contained Vision Transformer encoder (the shared backbone).
3Phase 0 needs *one* tiny, CPU-friendly encoder that both SSL methods can share:
5* **MAE** needs patch-level masking and access to the visible-token latent, plus an
6 ``ids_restore`` permutation so a decoder can put the tokens back (see
7 :class:`cafl4ds.models.heads.MAEDecoder`).
8* **SimSiam** (and later BYOL/SimCLR) only needs a pooled embedding of the full image.
10So :class:`TinyViTEncoder` exposes :meth:`forward_encoder` (optionally masked, the MAE
11entry point) and :meth:`embed` (a pooled ``[B, d]`` representation, the joint-embedding /
12measurement entry point). It is deliberately hand-rolled and tiny — no external model zoo —
13so the from-scratch smoke test runs on CPU in seconds; at HPU scale a larger / pretrained
14encoder satisfying the same two methods can be dropped in via config.
16Shapes follow the standard ViT/MAE convention: images are ``[B, C, H, W]``; a patch
17sequence is ``[B, N, d]`` with ``N = (H / p) * (W / p)`` patches; the token sequence handed
18to the transformer blocks is ``[B, 1 + N, d]`` (a prepended ``cls`` token).
19"""
21from __future__ import annotations
23from typing import cast
25import torch
26from torch import nn
29def patchify(imgs: torch.Tensor, patch_size: int) -> torch.Tensor:
30 """Split a batch of images into a sequence of flattened patches.
32 Args:
33 imgs: Images ``[B, C, H, W]`` with ``H`` and ``W`` divisible by ``patch_size``.
34 patch_size: Side length ``p`` of each square patch.
36 Returns:
37 Patches ``[B, N, p * p * C]`` with ``N = (H / p) * (W / p)``, in row-major
38 (top-left to bottom-right) patch order.
39 """
40 b, c, h, w = imgs.shape
41 p = patch_size
42 nh, nw = h // p, w // p
43 x = imgs.reshape(b, c, nh, p, nw, p)
44 x = torch.einsum("bchpwq->bhwpqc", x)
45 return x.reshape(b, nh * nw, p * p * c)
48def unpatchify(patches: torch.Tensor, patch_size: int, channels: int) -> torch.Tensor:
49 """Inverse of :func:`patchify`: reassemble a patch sequence into images.
51 Args:
52 patches: Patches ``[B, N, p * p * C]`` in row-major patch order.
53 patch_size: Side length ``p`` of each square patch.
54 channels: Number of image channels ``C``.
56 Returns:
57 Images ``[B, C, H, W]`` with ``H = W = p * sqrt(N)`` (square grid assumed).
58 """
59 b, n, _ = patches.shape
60 p, c = patch_size, channels
61 g = int(round(n**0.5))
62 x = patches.reshape(b, g, g, p, p, c)
63 x = torch.einsum("bhwpqc->bchpwq", x)
64 return x.reshape(b, c, g * p, g * p)
67class Mlp(nn.Module): # type: ignore[misc] # nn.Module is Any without torch stubs (mypy hook env)
68 """A two-layer MLP with GELU activation (transformer feed-forward block)."""
70 def __init__(self, dim: int, hidden_dim: int) -> None:
71 """Build the MLP.
73 Args:
74 dim: Input and output dimensionality.
75 hidden_dim: Width of the hidden layer.
76 """
77 super().__init__()
78 self.fc1 = nn.Linear(dim, hidden_dim)
79 self.act = nn.GELU()
80 self.fc2 = nn.Linear(hidden_dim, dim)
82 def forward(self, x: torch.Tensor) -> torch.Tensor:
83 """Apply the MLP to ``x`` (``[..., dim] -> [..., dim]``)."""
84 return self.fc2(self.act(self.fc1(x)))
87class Attention(nn.Module): # type: ignore[misc] # nn.Module is Any without torch stubs (mypy hook env)
88 """Standard multi-head self-attention over a token sequence."""
90 def __init__(self, dim: int, num_heads: int) -> None:
91 """Build the attention layer.
93 Args:
94 dim: Token dimensionality; must be divisible by ``num_heads``.
95 num_heads: Number of attention heads.
97 Raises:
98 ValueError: If ``dim`` is not divisible by ``num_heads``.
99 """
100 super().__init__()
101 if dim % num_heads != 0: 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true
102 raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads}).")
103 self.num_heads = num_heads
104 self.head_dim = dim // num_heads
105 self.qkv = nn.Linear(dim, dim * 3)
106 self.proj = nn.Linear(dim, dim)
108 def forward(self, x: torch.Tensor) -> torch.Tensor:
109 """Self-attend over the token dimension (``[B, T, d] -> [B, T, d]``)."""
110 b, t, d = x.shape
111 qkv = self.qkv(x).reshape(b, t, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
112 q, k, v = qkv[0], qkv[1], qkv[2]
113 out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
114 out = out.transpose(1, 2).reshape(b, t, d)
115 return self.proj(out)
118class Block(nn.Module): # type: ignore[misc] # nn.Module is Any without torch stubs (mypy hook env)
119 """A pre-norm transformer block: attention + MLP with residual connections."""
121 def __init__(self, dim: int, num_heads: int, mlp_ratio: float) -> None:
122 """Build the transformer block.
124 Args:
125 dim: Token dimensionality.
126 num_heads: Number of attention heads.
127 mlp_ratio: Hidden-layer width of the MLP as a multiple of ``dim``.
128 """
129 super().__init__()
130 self.norm1 = nn.LayerNorm(dim)
131 self.attn = Attention(dim, num_heads)
132 self.norm2 = nn.LayerNorm(dim)
133 self.mlp = Mlp(dim, int(dim * mlp_ratio))
135 def forward(self, x: torch.Tensor) -> torch.Tensor:
136 """Apply attention then MLP, each with a residual connection."""
137 x = x + self.attn(self.norm1(x))
138 x = x + self.mlp(self.norm2(x))
139 return x
142class TinyViTEncoder(nn.Module): # type: ignore[misc] # nn.Module is Any without torch stubs (mypy hook env)
143 """A tiny ViT encoder shared by the SSL methods.
145 Exposes two entry points:
147 * :meth:`forward_encoder` — run the transformer over the (optionally masked) patch
148 sequence; the MAE path uses ``mask_ratio > 0`` and consumes the returned mask /
149 restore indices.
150 * :meth:`embed` — a pooled ``[B, embed_dim]`` representation of the full image (mean of
151 the patch tokens after the final norm); the joint-embedding and measurement path.
153 The embedding read by the health instruments and probes is deliberately the *backbone*
154 representation (:meth:`embed`), never a projector/predictor output — those are training
155 heads, not the representation under study.
156 """
158 def __init__(
159 self,
160 img_size: int = 32,
161 patch_size: int = 8,
162 in_chans: int = 3,
163 embed_dim: int = 96,
164 depth: int = 4,
165 num_heads: int = 3,
166 mlp_ratio: float = 2.0,
167 ) -> None:
168 """Build the encoder.
170 Args:
171 img_size: Input image side length (square); must be divisible by ``patch_size``.
172 patch_size: Side length of each square patch.
173 in_chans: Number of input channels.
174 embed_dim: Token / embedding dimensionality.
175 depth: Number of transformer blocks.
176 num_heads: Number of attention heads per block.
177 mlp_ratio: MLP hidden width as a multiple of ``embed_dim``.
179 Raises:
180 ValueError: If ``img_size`` is not divisible by ``patch_size``.
181 """
182 super().__init__()
183 if img_size % patch_size != 0:
184 raise ValueError(f"img_size ({img_size}) must be divisible by patch_size ({patch_size}).")
185 self.patch_size = patch_size
186 self.in_chans = in_chans
187 self.embed_dim = embed_dim
188 self.num_patches = (img_size // patch_size) ** 2
190 self.patch_embed = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
191 self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
192 self.pos_embed = nn.Parameter(torch.zeros(1, 1 + self.num_patches, embed_dim))
193 self.blocks = nn.ModuleList([Block(embed_dim, num_heads, mlp_ratio) for _ in range(depth)])
194 self.norm = nn.LayerNorm(embed_dim)
195 self._init_weights()
197 def _init_weights(self) -> None:
198 """Initialize positional/token parameters and the linear/conv layers."""
199 nn.init.trunc_normal_(self.pos_embed, std=0.02)
200 nn.init.trunc_normal_(self.cls_token, std=0.02)
201 for m in self.modules():
202 if isinstance(m, nn.Linear):
203 nn.init.trunc_normal_(m.weight, std=0.02)
204 if m.bias is not None: 204 ↛ 201line 204 didn't jump to line 201 because the condition on line 204 was always true
205 nn.init.zeros_(m.bias)
206 elif isinstance(m, nn.LayerNorm):
207 nn.init.ones_(m.weight)
208 nn.init.zeros_(m.bias)
210 def _embed_patches(self, imgs: torch.Tensor) -> torch.Tensor:
211 """Convolutional patch embedding, adding the patch positional encodings.
213 Args:
214 imgs: Images ``[B, C, H, W]``.
216 Returns:
217 Patch tokens ``[B, N, embed_dim]`` with positional embeddings added (the ``cls``
218 token and its position are handled by the callers).
219 """
220 x = self.patch_embed(imgs).flatten(2).transpose(1, 2) # [B, N, D]
221 return cast(torch.Tensor, x + self.pos_embed[:, 1:, :])
223 @staticmethod
224 def random_masking(x: torch.Tensor, mask_ratio: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
225 """Per-sample random masking of a patch sequence (MAE, He et al. 2022).
227 Args:
228 x: Patch tokens ``[B, N, d]``.
229 mask_ratio: Fraction of patches to drop, in ``[0, 1)``.
231 Returns:
232 A tuple ``(x_kept, mask, ids_restore)`` where ``x_kept`` is ``[B, N_keep, d]``
233 (the visible tokens), ``mask`` is ``[B, N]`` (1 = masked/removed, 0 = kept), and
234 ``ids_restore`` is ``[B, N]`` — the permutation that undoes the shuffle.
235 """
236 b, n, d = x.shape
237 len_keep = max(1, int(round(n * (1.0 - mask_ratio))))
238 noise = torch.rand(b, n, device=x.device)
239 ids_shuffle = torch.argsort(noise, dim=1)
240 ids_restore = torch.argsort(ids_shuffle, dim=1)
241 ids_keep = ids_shuffle[:, :len_keep]
242 x_kept = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).expand(-1, -1, d))
243 mask = torch.ones(b, n, device=x.device)
244 mask[:, :len_keep] = 0
245 mask = torch.gather(mask, dim=1, index=ids_restore)
246 return x_kept, mask, ids_restore
248 def forward_encoder(
249 self, imgs: torch.Tensor, mask_ratio: float = 0.0
250 ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
251 """Run the transformer over the (optionally masked) patch sequence.
253 Args:
254 imgs: Images ``[B, C, H, W]``.
255 mask_ratio: Fraction of patches to mask before encoding (MAE). ``0`` disables
256 masking and processes the full sequence (joint-embedding / measurement).
258 Returns:
259 A tuple ``(tokens, mask, ids_restore)``. ``tokens`` is ``[B, 1 + N_keep, d]``
260 (final-normed, ``cls`` token first). ``mask`` and ``ids_restore`` are ``[B, N]``
261 when ``mask_ratio > 0`` and ``None`` otherwise.
262 """
263 x = self._embed_patches(imgs)
264 mask: torch.Tensor | None = None
265 ids_restore: torch.Tensor | None = None
266 if mask_ratio > 0.0:
267 x, mask, ids_restore = self.random_masking(x, mask_ratio)
268 cls = (self.cls_token + self.pos_embed[:, :1, :]).expand(x.shape[0], -1, -1)
269 x = torch.cat([cls, x], dim=1)
270 for blk in self.blocks:
271 x = blk(x)
272 return self.norm(x), mask, ids_restore
274 def forward(self, imgs: torch.Tensor) -> torch.Tensor:
275 """Return the full-sequence tokens ``[B, 1 + N, d]`` (unmasked)."""
276 tokens, _, _ = self.forward_encoder(imgs, mask_ratio=0.0)
277 return tokens
279 def embed(self, imgs: torch.Tensor) -> torch.Tensor:
280 """Return a pooled ``[B, embed_dim]`` representation (mean of the patch tokens).
282 This is the representation consumed by the health instruments and probes — the
283 backbone output, never a training-head projection.
285 Args:
286 imgs: Images ``[B, C, H, W]``.
288 Returns:
289 The mean-pooled patch-token embedding ``[B, embed_dim]``.
290 """
291 # The measurement/probe callers hold their fixed eval tensors on CPU, so coerce to
292 # the backbone's device here (a no-op when both are CPU). This keeps ``embed`` — the
293 # instrument entry point — usable when the model has been moved to an accelerator
294 # (e.g. ``hpu``) without pushing device bookkeeping into the monitor/measurements.
295 imgs = imgs.to(self.pos_embed.device)
296 tokens = self.forward(imgs)
297 return tokens[:, 1:, :].mean(dim=1)