| |
| |
| from collections.abc import Callable |
|
|
| from einops import rearrange, repeat |
| from einops.layers.torch import Rearrange |
| import torch |
| from torch import nn |
| from transformers import PreTrainedModel |
| from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS |
| from transformers.utils import logging |
|
|
| from .configuration_metom import MetomConfig |
|
|
| logger = logging.get_logger(__name__) |
|
|
|
|
| def size_pair(t): |
| return t if isinstance(t, tuple) else (t, t) |
|
|
|
|
| def metom_eager_attention_forward( |
| module: nn.Module, |
| query: torch.Tensor, |
| key: torch.Tensor, |
| value: torch.Tensor, |
| attention_mask: torch.Tensor | None, |
| scaling: float | None = None, |
| dropout: float = 0.0, |
| **kwargs, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| if scaling is None: |
| scaling = query.size(-1) ** -0.5 |
|
|
| attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling |
| if attention_mask is not None: |
| attention_mask = attention_mask[:, :, :, : key.shape[-2]] |
| attn_weights = attn_weights + attention_mask |
|
|
| attn_weights = nn.functional.softmax(attn_weights, dim=-1) |
| attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) |
| attn_output = torch.matmul(attn_weights, value) |
| attn_output = attn_output.transpose(1, 2).contiguous() |
| return attn_output, attn_weights |
|
|
|
|
| class MetomFeedForward(nn.Module): |
| def __init__(self, dim, hidden_dim, dropout): |
| super().__init__() |
| self.net = nn.Sequential( |
| nn.LayerNorm(dim), |
| nn.Linear(dim, hidden_dim), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(hidden_dim, dim), |
| nn.Dropout(dropout), |
| ) |
|
|
| def forward(self, x): |
| return self.net(x) |
|
|
|
|
| class MetomAttention(nn.Module): |
| def __init__(self, config: MetomConfig): |
| super().__init__() |
| inner_dim = config.dim_head * config.heads |
| project_out = not (config.heads == 1 and config.dim_head == config.dim) |
|
|
| self.config = config |
| self.heads = config.heads |
| self.scale = config.dim_head ** -0.5 |
| self.norm = nn.LayerNorm(config.dim) |
| self.dropout = nn.Dropout(config.dropout) |
| self.to_qkv = nn.Linear(config.dim, inner_dim * 3, bias=False) |
| self.to_out = nn.Sequential( |
| nn.Linear(inner_dim, config.dim), |
| nn.Dropout(config.dropout), |
| ) if project_out else nn.Identity() |
|
|
| def forward(self, x: torch.Tensor, **kwargs): |
| x = self.norm(x) |
| qkv = self.to_qkv(x).chunk(3, dim=-1) |
| q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.heads), qkv) |
|
|
| attn_implementation = self.config._attn_implementation or "eager" |
| if attn_implementation == "flex_attention": |
| if self.training and self.dropout.p > 0: |
| logger.warning_once( |
| "`flex_attention` does not support attention dropout during training. Falling back to `eager`." |
| ) |
| attn_implementation = "eager" |
| attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface( |
| attn_implementation, |
| metom_eager_attention_forward, |
| ) |
| out, _ = attention_interface( |
| self, |
| q, |
| k, |
| v, |
| None, |
| is_causal=False, |
| scaling=self.scale, |
| dropout=0.0 if not self.training else self.dropout.p, |
| **kwargs, |
| ) |
| out = rearrange(out, "b n h d -> b n (h d)") |
| return self.to_out(out) |
|
|
|
|
| class MetomTransformer(nn.Module): |
| def __init__(self, config: MetomConfig): |
| super().__init__() |
| self.norm = nn.LayerNorm(config.dim) |
| self.layers = nn.ModuleList([]) |
| for _ in range(config.depth): |
| self.layers.append( |
| nn.ModuleList( |
| [ |
| MetomAttention(config), |
| MetomFeedForward(config.dim, config.mlp_dim, dropout=config.dropout), |
| ] |
| ) |
| ) |
|
|
| def forward(self, x: torch.Tensor, **kwargs): |
| for attn, ff in self.layers: |
| x = attn(x, **kwargs) + x |
| x = ff(x) + x |
| return self.norm(x) |
|
|
|
|
| class MetomModel(PreTrainedModel): |
| config_class = MetomConfig |
| main_input_name = "pixel_values" |
| _supports_attention_backend = True |
| _supports_flash_attn = True |
| _supports_flash_attn_2 = True |
| _supports_sdpa = True |
| _supports_flex_attn = True |
|
|
| def __init__(self, config: MetomConfig): |
| super().__init__(config) |
| image_height, image_width = size_pair(config.image_size) |
| patch_height, patch_width = size_pair(config.patch_size) |
| assert image_height % patch_height == 0 and image_width % patch_width == 0, ( |
| "Image dimensions must be divisible by the patch size." |
| ) |
|
|
| num_patches = (image_height // patch_height) * (image_width // patch_width) |
| patch_dim = config.channels * patch_height * patch_width |
| assert config.pool in {"cls", "mean"}, "pool type must be either cls (cls token) or mean (mean pooling)" |
| assert len(config.labels) > 0, "labels must be composed of at least one label" |
|
|
| self.to_patch_embedding = nn.Sequential( |
| Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=patch_height, p2=patch_width), |
| nn.LayerNorm(patch_dim), |
| nn.Linear(patch_dim, config.dim), |
| nn.LayerNorm(config.dim), |
| ) |
| self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, config.dim)) |
| self.cls_token = nn.Parameter(torch.randn(1, 1, config.dim)) |
| self.dropout = nn.Dropout(config.emb_dropout) |
| self.transformer = MetomTransformer(config) |
| self.pool = config.pool |
| self.to_latent = nn.Identity() |
| self.mlp_head = nn.Linear(config.dim, len(config.labels)) |
| self.labels = config.labels |
| self.post_init() |
|
|
| def forward(self, pixel_values: torch.Tensor, **kwargs): |
| x = self.to_patch_embedding(pixel_values) |
| b, n, _ = x.shape |
| cls_tokens = repeat(self.cls_token, "1 1 d -> b 1 d", b=b) |
| x = torch.cat((cls_tokens, x), dim=1) |
| x += self.pos_embedding[:, : (n + 1)] |
| x = self.dropout(x) |
| x = self.transformer(x, **kwargs) |
| x = x.mean(dim=1) if self.pool == "mean" else x[:, 0] |
| x = self.to_latent(x) |
| return self.mlp_head(x) |
|
|
| def get_predictions(self, pixel_values: torch.Tensor) -> list[str]: |
| logits = self(pixel_values=pixel_values) |
| indices = torch.argmax(logits, dim=-1) |
| return [self.labels[i] for i in indices] |
|
|
| def get_topk_labels( |
| self, pixel_values: torch.Tensor, k: int = 5, return_probs: bool = False |
| ) -> list[list[str]] | list[list[tuple[str, float]]]: |
| assert 0 < k <= len(self.labels), "k must be a positive integer less than or equal to the number of labels" |
| logits = self(pixel_values=pixel_values) |
| probs = torch.softmax(logits, dim=-1) |
| topk_probs, topk_indices = torch.topk(probs, k, dim=-1) |
| topk_labels = [[self.labels[i] for i in ti] for ti in topk_indices] |
| if return_probs: |
| return [ |
| [(label, prob.item()) for label, prob in zip(labels, probs)] |
| for labels, probs in zip(topk_labels, topk_probs) |
| ] |
| return topk_labels |
|
|
