VisionLanguageGroup/MicroscopyMatching
0
1"""Transformer class."""2 3import logging4import math5from collections import OrderedDict6from pathlib import Path7from typing import Literal, Tuple8 9import torch10import torch.nn.functional as F11 12import yaml13from torch import nn14 15import sys, os16 17from .utils import blockwise_causal_norm18 19logger = logging.getLogger(__name__)20 21 22def _pos_embed_fourier1d_init(23 cutoff: float = 256, n: int = 32, cutoff_start: float = 124):25 return (26 torch.exp(torch.linspace(-math.log(cutoff_start), -math.log(cutoff), n))27 .unsqueeze(0)28 .unsqueeze(0)29 )30 31 32def _rope_pos_embed_fourier1d_init(cutoff: float = 128, n: int = 32):33 # Maximum initial frequency is 134 return torch.exp(torch.linspace(0, -math.log(cutoff), n)).unsqueeze(0).unsqueeze(0)35 36 37def _rotate_half(x: torch.Tensor) -> torch.Tensor:38 """Rotate pairs of scalars as 2d vectors by pi/2."""39 x = x.unflatten(-1, (-1, 2))40 x1, x2 = x.unbind(dim=-1)41 return torch.stack((-x2, x1), dim=-1).flatten(start_dim=-2)42 43 44class RotaryPositionalEncoding(nn.Module):45 def __init__(self, cutoffs: Tuple[float] = (256,), n_pos: Tuple[int] = (32,)):46 super().__init__()47 assert len(cutoffs) == len(n_pos)48 if not all(n % 2 == 0 for n in n_pos):49 raise ValueError("n_pos must be even")50 51 self._n_dim = len(cutoffs)52 self.freqs = nn.ParameterList([53 nn.Parameter(_rope_pos_embed_fourier1d_init(cutoff, n // 2))54 for cutoff, n in zip(cutoffs, n_pos)55 ])56 57 def get_co_si(self, coords: torch.Tensor):58 _B, _N, D = coords.shape59 assert D == len(self.freqs)60 co = torch.cat(61 tuple(62 torch.cos(0.5 * math.pi * x.unsqueeze(-1) * freq) / math.sqrt(len(freq))63 for x, freq in zip(coords.moveaxis(-1, 0), self.freqs)64 ),65 axis=-1,66 )67 si = torch.cat(68 tuple(69 torch.sin(0.5 * math.pi * x.unsqueeze(-1) * freq) / math.sqrt(len(freq))70 for x, freq in zip(coords.moveaxis(-1, 0), self.freqs)71 ),72 axis=-1,73 )74 return co, si75 76 def forward(self, q: torch.Tensor, k: torch.Tensor, coords: torch.Tensor):77 _B, _N, D = coords.shape78 _B, _H, _N, _C = q.shape79 80 if D != self._n_dim:81 raise ValueError(f"coords must have {self._n_dim} dimensions, got {D}")82 83 co, si = self.get_co_si(coords)84 co = co.unsqueeze(1).repeat_interleave(2, dim=-1)85 si = si.unsqueeze(1).repeat_interleave(2, dim=-1)86 q2 = q * co + _rotate_half(q) * si87 k2 = k * co + _rotate_half(k) * si88 return q2, k289 90 91class FeedForward(nn.Module):92 def __init__(self, d_model, expand: float = 2, bias: bool = True):93 super().__init__()94 self.fc1 = nn.Linear(d_model, int(d_model * expand))95 self.fc2 = nn.Linear(int(d_model * expand), d_model, bias=bias)96 self.act = nn.GELU()97 98 def forward(self, x):99 return self.fc2(self.act(self.fc1(x)))100 101 102class PositionalEncoding(nn.Module):103 def __init__(104 self,105 cutoffs: Tuple[float] = (256,),106 n_pos: Tuple[int] = (32,),107 cutoffs_start=None,108 ):109 super().__init__()110 if cutoffs_start is None:111 cutoffs_start = (1,) * len(cutoffs)112 113 assert len(cutoffs) == len(n_pos)114 self.freqs = nn.ParameterList([115 nn.Parameter(_pos_embed_fourier1d_init(cutoff, n // 2))116 for cutoff, n, cutoff_start in zip(cutoffs, n_pos, cutoffs_start)117 ])118 119 def forward(self, coords: torch.Tensor):120 _B, _N, D = coords.shape121 assert D == len(self.freqs)122 embed = torch.cat(123 tuple(124 torch.cat(125 (126 torch.sin(0.5 * math.pi * x.unsqueeze(-1) * freq),127 torch.cos(0.5 * math.pi * x.unsqueeze(-1) * freq),128 ),129 axis=-1,130 )131 / math.sqrt(len(freq))132 for x, freq in zip(coords.moveaxis(-1, 0), self.freqs)133 ),134 axis=-1,135 )136 return embed137 138 139def _bin_init_exp(cutoff: float, n: int):140 return torch.exp(torch.linspace(0, math.log(cutoff + 1), n))141 142 143def _bin_init_linear(cutoff: float, n: int):144 return torch.linspace(-cutoff, cutoff, n)145 146 147class RelativePositionalBias(nn.Module):148 def __init__(149 self,150 n_head: int,151 cutoff_spatial: float,152 cutoff_temporal: float,153 n_spatial: int = 32,154 n_temporal: int = 16,155 ):156 super().__init__()157 self._spatial_bins = _bin_init_exp(cutoff_spatial, n_spatial)158 self._temporal_bins = _bin_init_linear(cutoff_temporal, 2 * n_temporal + 1)159 self.register_buffer("spatial_bins", self._spatial_bins)160 self.register_buffer("temporal_bins", self._temporal_bins)161 self.n_spatial = n_spatial162 self.n_head = n_head163 self.bias = nn.Parameter(164 -0.5 + torch.rand((2 * n_temporal + 1) * n_spatial, n_head)165 )166 167 def forward(self, coords: torch.Tensor):168 _B, _N, _D = coords.shape169 t = coords[..., 0]170 yx = coords[..., 1:]171 temporal_dist = t.unsqueeze(-1) - t.unsqueeze(-2)172 spatial_dist = torch.cdist(yx, yx)173 174 spatial_idx = torch.bucketize(spatial_dist, self.spatial_bins)175 torch.clamp_(spatial_idx, max=len(self.spatial_bins) - 1)176 temporal_idx = torch.bucketize(temporal_dist, self.temporal_bins)177 torch.clamp_(temporal_idx, max=len(self.temporal_bins) - 1)178 179 idx = spatial_idx.flatten() + temporal_idx.flatten() * self.n_spatial180 bias = self.bias.index_select(0, idx).view((*spatial_idx.shape, self.n_head))181 bias = bias.transpose(-1, 1)182 return bias183 184 185class RelativePositionalAttention(nn.Module):186 def __init__(187 self,188 coord_dim: int,189 embed_dim: int,190 n_head: int,191 cutoff_spatial: float = 256,192 cutoff_temporal: float = 16,193 n_spatial: int = 32,194 n_temporal: int = 16,195 dropout: float = 0.0,196 mode: Literal["bias", "rope", "none"] = "bias",197 attn_dist_mode: str = "v0",198 ):199 super().__init__()200 201 if not embed_dim % (2 * n_head) == 0:202 raise ValueError(203 f"embed_dim {embed_dim} must be divisible by 2 times n_head {2 * n_head}"204 )205 206 self.q_pro = nn.Linear(embed_dim, embed_dim, bias=True)207 self.k_pro = nn.Linear(embed_dim, embed_dim, bias=True)208 self.v_pro = nn.Linear(embed_dim, embed_dim, bias=True)209 self.proj = nn.Linear(embed_dim, embed_dim)210 self.dropout = dropout211 self.n_head = n_head212 self.embed_dim = embed_dim213 self.cutoff_spatial = cutoff_spatial214 self.attn_dist_mode = attn_dist_mode215 216 if mode == "bias" or mode is True:217 self.pos_bias = RelativePositionalBias(218 n_head=n_head,219 cutoff_spatial=cutoff_spatial,220 cutoff_temporal=cutoff_temporal,221 n_spatial=n_spatial,222 n_temporal=n_temporal,223 )224 elif mode == "rope":225 n_split = 2 * (embed_dim // (2 * (coord_dim + 1) * n_head))226 self.rot_pos_enc = RotaryPositionalEncoding(227 cutoffs=((cutoff_temporal,) + (cutoff_spatial,) * coord_dim),228 n_pos=(embed_dim // n_head - coord_dim * n_split,)229 + (n_split,) * coord_dim,230 )231 elif mode == "none":232 pass233 elif mode is None or mode is False:234 logger.warning(235 "attn_positional_bias is not set (None or False), no positional bias."236 )237 else:238 raise ValueError(f"Unknown mode {mode}")239 240 self._mode = mode241 242 def forward(243 self,244 query: torch.Tensor,245 key: torch.Tensor,246 value: torch.Tensor,247 coords: torch.Tensor,248 padding_mask: torch.Tensor = None,249 ):250 B, N, D = query.size()251 q = self.q_pro(query)252 k = self.k_pro(key)253 v = self.v_pro(value)254 k = k.view(B, N, self.n_head, D // self.n_head).transpose(1, 2)255 q = q.view(B, N, self.n_head, D // self.n_head).transpose(1, 2)256 v = v.view(B, N, self.n_head, D // self.n_head).transpose(1, 2)257 258 attn_mask = torch.zeros(259 (B, self.n_head, N, N), device=query.device, dtype=q.dtype260 )261 attn_ignore_val = -1e3262 263 yx = coords[..., 1:]264 spatial_dist = torch.cdist(yx, yx)265 spatial_mask = (spatial_dist > self.cutoff_spatial).unsqueeze(1)266 attn_mask.masked_fill_(spatial_mask, attn_ignore_val)267 268 if coords is not None:269 if self._mode == "bias":270 attn_mask = attn_mask + self.pos_bias(coords)271 elif self._mode == "rope":272 q, k = self.rot_pos_enc(q, k, coords)273 274 if self.attn_dist_mode == "v0":275 dist = torch.cdist(coords, coords, p=2)276 attn_mask += torch.exp(-0.1 * dist.unsqueeze(1))277 elif self.attn_dist_mode == "v1":278 attn_mask += torch.exp(279 -5 * spatial_dist.unsqueeze(1) / self.cutoff_spatial280 )281 else:282 raise ValueError(f"Unknown attn_dist_mode {self.attn_dist_mode}")283 284 if padding_mask is not None:285 ignore_mask = torch.logical_or(286 padding_mask.unsqueeze(1), padding_mask.unsqueeze(2)287 ).unsqueeze(1)288 attn_mask.masked_fill_(ignore_mask, attn_ignore_val)289 290 y = F.scaled_dot_product_attention(291 q, k, v, attn_mask=attn_mask, dropout_p=self.dropout if self.training else 0292 )293 y = y.transpose(1, 2).contiguous().view(B, N, D)294 y = self.proj(y)295 return y296 297 298class EncoderLayer(nn.Module):299 def __init__(300 self,301 coord_dim: int = 2,302 d_model=256,303 num_heads=4,304 dropout=0.1,305 cutoff_spatial: int = 256,306 window: int = 16,307 positional_bias: Literal["bias", "rope", "none"] = "bias",308 positional_bias_n_spatial: int = 32,309 attn_dist_mode: str = "v0",310 ):311 super().__init__()312 self.positional_bias = positional_bias313 self.attn = RelativePositionalAttention(314 coord_dim,315 d_model,316 num_heads,317 cutoff_spatial=cutoff_spatial,318 n_spatial=positional_bias_n_spatial,319 cutoff_temporal=window,320 n_temporal=window,321 dropout=dropout,322 mode=positional_bias,323 attn_dist_mode=attn_dist_mode,324 )325 self.mlp = FeedForward(d_model)326 self.norm1 = nn.LayerNorm(d_model)327 self.norm2 = nn.LayerNorm(d_model)328 329 def forward(330 self,331 x: torch.Tensor,332 coords: torch.Tensor,333 padding_mask: torch.Tensor = None,334 ):335 x = self.norm1(x)336 337 # setting coords to None disables positional bias338 a = self.attn(339 x,340 x,341 x,342 coords=coords if self.positional_bias else None,343 padding_mask=padding_mask,344 )345 346 x = x + a347 x = x + self.mlp(self.norm2(x))348 349 return x350 351 352class DecoderLayer(nn.Module):353 def __init__(354 self,355 coord_dim: int = 2,356 d_model=256,357 num_heads=4,358 dropout=0.1,359 window: int = 16,360 cutoff_spatial: int = 256,361 positional_bias: Literal["bias", "rope", "none"] = "bias",362 positional_bias_n_spatial: int = 32,363 attn_dist_mode: str = "v0",364 ):365 super().__init__()366 self.positional_bias = positional_bias367 self.attn = RelativePositionalAttention(368 coord_dim,369 d_model,370 num_heads,371 cutoff_spatial=cutoff_spatial,372 n_spatial=positional_bias_n_spatial,373 cutoff_temporal=window,374 n_temporal=window,375 dropout=dropout,376 mode=positional_bias,377 attn_dist_mode=attn_dist_mode,378 )379 380 self.mlp = FeedForward(d_model)381 self.norm1 = nn.LayerNorm(d_model)382 self.norm2 = nn.LayerNorm(d_model)383 self.norm3 = nn.LayerNorm(d_model)384 385 def forward(386 self,387 x: torch.Tensor,388 y: torch.Tensor,389 coords: torch.Tensor,390 padding_mask: torch.Tensor = None,391 ):392 x = self.norm1(x)393 y = self.norm2(y)394 # cross attention395 # setting coords to None disables positional bias396 a = self.attn(397 x,398 y,399 y,400 coords=coords if self.positional_bias else None,401 padding_mask=padding_mask,402 )403 404 x = x + a405 x = x + self.mlp(self.norm3(x))406 407 return x408 409 410 411class TrackingTransformer(torch.nn.Module):412 def __init__(413 self,414 coord_dim: int = 3,415 feat_dim: int = 0,416 d_model: int = 128,417 nhead: int = 4,418 num_encoder_layers: int = 4,419 num_decoder_layers: int = 4,420 dropout: float = 0.1,421 pos_embed_per_dim: int = 32,422 feat_embed_per_dim: int = 1,423 window: int = 6,424 spatial_pos_cutoff: int = 256,425 attn_positional_bias: Literal["bias", "rope", "none"] = "rope",426 attn_positional_bias_n_spatial: int = 16,427 causal_norm: Literal[428 "none", "linear", "softmax", "quiet_softmax"429 ] = "quiet_softmax",430 attn_dist_mode: str = "v0",431 ):432 super().__init__()433 434 self.config = dict(435 coord_dim=coord_dim,436 feat_dim=feat_dim,437 pos_embed_per_dim=pos_embed_per_dim,438 d_model=d_model,439 nhead=nhead,440 num_encoder_layers=num_encoder_layers,441 num_decoder_layers=num_decoder_layers,442 window=window,443 dropout=dropout,444 attn_positional_bias=attn_positional_bias,445 attn_positional_bias_n_spatial=attn_positional_bias_n_spatial,446 spatial_pos_cutoff=spatial_pos_cutoff,447 feat_embed_per_dim=feat_embed_per_dim,448 causal_norm=causal_norm,449 attn_dist_mode=attn_dist_mode,450 )451 452 # TODO remove, alredy present in self.config453 # self.window = window454 # self.feat_dim = feat_dim455 # self.coord_dim = coord_dim456 457 self.proj = nn.Linear(458 (1 + coord_dim) * pos_embed_per_dim + feat_dim * feat_embed_per_dim, d_model459 )460 self.norm = nn.LayerNorm(d_model)461 462 self.encoder = nn.ModuleList([463 EncoderLayer(464 coord_dim,465 d_model,466 nhead,467 dropout,468 window=window,469 cutoff_spatial=spatial_pos_cutoff,470 positional_bias=attn_positional_bias,471 positional_bias_n_spatial=attn_positional_bias_n_spatial,472 attn_dist_mode=attn_dist_mode,473 )474 for _ in range(num_encoder_layers)475 ])476 self.decoder = nn.ModuleList([477 DecoderLayer(478 coord_dim,479 d_model,480 nhead,481 dropout,482 window=window,483 cutoff_spatial=spatial_pos_cutoff,484 positional_bias=attn_positional_bias,485 positional_bias_n_spatial=attn_positional_bias_n_spatial,486 attn_dist_mode=attn_dist_mode,487 )488 for _ in range(num_decoder_layers)489 ])490 491 self.head_x = FeedForward(d_model)492 self.head_y = FeedForward(d_model)493 494 if feat_embed_per_dim > 1:495 self.feat_embed = PositionalEncoding(496 cutoffs=(1000,) * feat_dim,497 n_pos=(feat_embed_per_dim,) * feat_dim,498 cutoffs_start=(0.01,) * feat_dim,499 )500 else:501 self.feat_embed = nn.Identity()502 503 self.pos_embed = PositionalEncoding(504 cutoffs=(window,) + (spatial_pos_cutoff,) * coord_dim,505 n_pos=(pos_embed_per_dim,) * (1 + coord_dim),506 )507 508 # self.pos_embed = NoPositionalEncoding(d=pos_embed_per_dim * (1 + coord_dim))509 510 # @profile511 def forward(self, coords, features=None, padding_mask=None, attn_feat=None):512 assert coords.ndim == 3 and coords.shape[-1] in (3, 4)513 _B, _N, _D = coords.shape514 515 # disable padded coords (such that it doesnt affect minimum)516 if padding_mask is not None:517 coords = coords.clone()518 coords[padding_mask] = coords.max()519 520 # remove temporal offset521 min_time = coords[:, :, :1].min(dim=1, keepdims=True).values522 coords = coords - min_time523 524 pos = self.pos_embed(coords)525 526 if features is None or features.numel() == 0:527 features = pos528 else:529 features = self.feat_embed(features)530 features = torch.cat((pos, features), axis=-1)531 532 features = self.proj(features)533 if attn_feat is not None:534 # add attention embedding535 features = features + attn_feat536 537 features = self.norm(features)538 539 x = features540 541 # encoder542 for enc in self.encoder:543 x = enc(x, coords=coords, padding_mask=padding_mask)544 545 y = features546 # decoder w cross attention547 for dec in self.decoder:548 y = dec(y, x, coords=coords, padding_mask=padding_mask)549 # y = dec(y, y, coords=coords, padding_mask=padding_mask)550 551 x = self.head_x(x)552 y = self.head_y(y)553 554 # outer product is the association matrix (logits)555 A = torch.einsum("bnd,bmd->bnm", x, y)556 557 return A558 559 def normalize_output(560 self,561 A: torch.FloatTensor,562 timepoints: torch.LongTensor,563 coords: torch.FloatTensor,564 ) -> torch.FloatTensor:565 """Apply (parental) softmax, or elementwise sigmoid.566 567 Args:568 A: Tensor of shape B, N, N569 timepoints: Tensor of shape B, N570 coords: Tensor of shape B, N, (time + n_spatial)571 """572 assert A.ndim == 3573 assert timepoints.ndim == 2574 assert coords.ndim == 3575 assert coords.shape[2] == 1 + self.config["coord_dim"]576 577 # spatial distances578 dist = torch.cdist(coords[:, :, 1:], coords[:, :, 1:])579 invalid = dist > self.config["spatial_pos_cutoff"]580 581 if self.config["causal_norm"] == "none":582 # Spatially distant entries are set to zero583 A = torch.sigmoid(A)584 A[invalid] = 0585 else:586 return torch.stack([587 blockwise_causal_norm(588 _A, _t, mode=self.config["causal_norm"], mask_invalid=_m589 )590 for _A, _t, _m in zip(A, timepoints, invalid)591 ])592 return A593 594 def save(self, folder):595 folder = Path(folder)596 folder.mkdir(parents=True, exist_ok=True)597 yaml.safe_dump(self.config, open(folder / "config.yaml", "w"))598 torch.save(self.state_dict(), folder / "model.pt")599 600 @classmethod601 def from_folder(602 cls, folder, map_location=None, checkpoint_path: str = "model.pt"603 ):604 folder = Path(folder)605 606 config = yaml.load(open(folder / "config.yaml"), Loader=yaml.FullLoader)607 608 model = cls(**config)609 610 fpath = folder / checkpoint_path611 logger.info(f"Loading model state from {fpath}")612 613 state = torch.load(fpath, map_location=map_location, weights_only=True)614 # if state is a checkpoint, we have to extract state_dict615 if "state_dict" in state:616 state = state["state_dict"]617 state = OrderedDict(618 (k[6:], v) for k, v in state.items() if k.startswith("model.")619 )620 model.load_state_dict(state)621 622 return model623 624 @classmethod625 def from_cfg(626 cls, cfg_path627 ):628 629 cfg_path = Path(cfg_path)630 631 config = yaml.load(open(cfg_path), Loader=yaml.FullLoader)632 633 model = cls(**config)634 635 return model636