Aluode/PerceptionLabPortable

1# This code was adapted from https://github.com/lucidrains/flamingo-pytorch licensed under the MIT License.2#3# MIT License4#5# Copyright (c) 2020  The Google AI Language Team Authors, The HuggingFace Inc. team and github/lonePatient6#7# Permission is hereby granted, free of charge, to any person obtaining a copy8# of this software and associated documentation files (the "Software"), to deal9# in the Software without restriction, including without limitation the rights10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell11# copies of the Software, and to permit persons to whom the Software is12# furnished to do so, subject to the following conditions:13#14# The above copyright notice and this permission notice shall be included in all15# copies or substantial portions of the Software.16#17# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE23# SOFTWARE.24 25 26"""27 28Generic interface to various configurations of the Perceiver Resampler, that simply takes in a series of (potentially29time-indexed) contextual embeddings, and "resamples" (compresses) them down to a pre-specified number of latents! Note30that the Perceiver in general resamples based solely off the *long-range* context; there's a nice opportunity here to31prime the Perceiver Resampler with say a single layer's worth of language embeddings (the target domain), and use that32to softly "retrieve & compress" what we need --> this would be a novel contribution we should explore.33 34References:35    - DeepMind's Flamingo: https://www.deepmind.com/blog/tackling-multiple-tasks-with-a-single-visual-language-model36    - Code borrowed w/ love from: https://github.com/lucidrains/flamingo-pytorch37 38"""39 40from typing import Optional41 42import torch43import torch.nn as nn44 45from .configuration_idefics import IdeficsConfig46 47 48class IdeficsPerceiverResampler(nn.Module):49    def __init__(50        self, config: IdeficsConfig, embed_dim: int, depth: int, n_heads: int, head_dim: int, n_latents: int51    ) -> None:52        """53        Instantiates a Perceiver Resampler that operates over a sequence of embeddings (say from a ResNet or ViT or54        MAE) of a given dimension, performs `depth` blocks of cross-attention with a fixed `n_latents` inputs, then55        returns a Tensor of shape [bsz, n_latents, embed_dim]. :param embed_dim: Dimensionality of embeddings being fed56        to the Perceiver Resampler (also dimensionality of latent embeddings *returned* by the Perceiver Resampler.57        Could be e.g., VIT embed_dim, ResNet pool dim, and so on.58 59        Args:60            config (`IdeficsConfig`): config object61            embed_dim (`int`): The size of each embedding vector62            depth (`int`): Depth of the Perceiver Resampler (Transformer w/ cross attention). Should be shallow (< 3).63            n_heads (`int`): Number of heads in each Transformer block (for multi-headed self-attention).64            head_dim (`int`): Dimensionality of each head projection in the Transformer block.65            n_latents (`int`):66                Number of latent embeddings to resample ("compress") the input sequence to (usually < 128).67 68        """69        super().__init__()70        self.embed_dim, self.n_heads, self.head_dim, self.n_latents = embed_dim, n_heads, head_dim, n_latents71        self.qk_layer_norms = config.perceiver_config.qk_layer_norms_perceiver72 73        # Create Latents for Perceiver74        self.latents = nn.Parameter(torch.randn(self.n_latents, self.embed_dim), requires_grad=True)75 76        self.intermediate_dim = (77            self.embed_dim * 478            if not hasattr(config.vision_config, "embed_dim")79            else config.vision_config.embed_dim * 480        )81        # Create Transformer Blocks82        self.blocks = nn.ModuleList(83            [84                nn.ModuleList(85                    [86                        IdeficsPerceiverAttention(self.embed_dim, self.n_heads, self.head_dim, self.qk_layer_norms),87                        IdeficsMLP(self.intermediate_dim, config),88                    ]89                )90                for _ in range(depth)91            ]92        )93        self.layer_norm = nn.LayerNorm(self.embed_dim)94 95    def forward(self, context: torch.Tensor) -> torch.Tensor:96        """Resample arbitrary length context & *compress* down to self.n_latents latent embeddings"""97        # einsum.repeat(self.latents, "seq embed -> bsz seq embed", bsz=context.shape[0])98        latents = self.latents.repeat(context.shape[0], 1, 1)99 100        # Feed through Perceiver Attention blocks...101        for attn, ff in self.blocks:102            latents = attn(context, latents) + latents103            latents = ff(latents) + latents104 105        return self.layer_norm(latents)106 107 108class IdeficsPerceiverAttention(nn.Module):109    def __init__(self, embed_dim: int, n_heads: int, head_dim: int, qk_layer_norms: bool) -> None:110        """Perceiver Cross-Attention Module --> let long-form inputs be `context`, resampled embeddings be `latents`"""111        super().__init__()112        self.embed_dim, self.n_heads, self.head_dim = embed_dim, n_heads, head_dim113        self.qk_layer_norms = qk_layer_norms114        # Normalization & Scaling115        self.context_layer_norm = nn.LayerNorm(self.embed_dim)116        self.latents_layer_norm = nn.LayerNorm(self.embed_dim)117        if self.qk_layer_norms:118            self.q_layer_norm = nn.LayerNorm(self.head_dim)119            self.k_layer_norm = nn.LayerNorm(self.head_dim)120 121        self.qk_scale = self.head_dim**-0.5122 123        # Q, K, V Projection (no bias -- detail from Perceiver/Flamingo Papers).124        self.q_proj = nn.Linear(self.embed_dim, self.n_heads * self.head_dim, bias=False)125        self.k_proj = nn.Linear(self.embed_dim, self.n_heads * self.head_dim, bias=False)126        self.v_proj = nn.Linear(self.embed_dim, self.n_heads * self.head_dim, bias=False)127 128        self.output_proj = nn.Linear(self.n_heads * self.head_dim, embed_dim, bias=False)129 130    def forward(self, context: torch.Tensor, latents: torch.Tensor) -> torch.Tensor:131        """132        Runs Perceiver Self-Attention, with special (context, latents) appended along the `seq` dimension!133 134        Args:135            context (`torch.Tensor`):136                Tensor of shape `[bsz, seq, embed_dim]` representing long-form context to resample.137            latents (`torch.Tensor`):138                Tensor of shape `[bsz, n_latents, embed_dim]` representing fixed length latents to compress to.139 140        Returns:141            `torch.Tensor`: Tensor of shape `[bsz, n_latents, embed_dim]` representing attention over latents w/ cross142            from context.143        """144        context = self.context_layer_norm(context)145        latents = self.latents_layer_norm(latents)146        batch_size, seq_length, embed_dim = context.shape[:3]147 148        # Query, Key, Value Projections --> Note that in Flamingo, latents are *concatenated* with context prior to attn!149        #   Note: This results in queries w/ `seq = n_latents`, and keys, values with `seq = len(context) + n_latents`150        q = self.q_proj(latents)151        k = self.k_proj(torch.cat([context, latents], dim=-2))152        v = self.v_proj(torch.cat([context, latents], dim=-2))153 154        # Multiheaded Self-Attention w/ stable softmax (subtract per-row max -- `amax` -- before softmax call)155        #   =>> `attn` should be a 2D matrix of shape [n_latents x (context + n_latents)]156        # einsum.rearrange(x, "bsz seq (heads embed) -> bsz heads seq embed", heads=self.n_heads)157        q, k, v = [x.reshape(batch_size, x.shape[1], self.n_heads, self.head_dim).transpose(1, 2) for x in (q, k, v)]158 159        if self.qk_layer_norms:160            q = self.q_layer_norm(q)161            k = self.k_layer_norm(k)162 163        scores = torch.einsum("... i d, ... j d -> ... i j", q * self.qk_scale, k)164        stabilized_scores = scores - (scores.amax(dim=-1, keepdim=True).detach())165        attn = stabilized_scores.softmax(dim=-1)166 167        # Attend & project back to output...168        resampled = torch.einsum("... i j, ... j d -> ... i d", attn, v)169        # einsum.rearrange(resampled, "bsz heads seq embed -> bsz seq (heads embed)", heads=self.n_heads)170        return self.output_proj(resampled.transpose(1, 2).flatten(-2))171 172 173class IdeficsMLP(nn.Module):174    def __init__(self, intermediate_size, config: IdeficsConfig):175        """Simple MLP block with intermediate_size and embedding size"""176        super().__init__()177        self.embed_dim = config.vision_config.embed_dim178        self.ln = nn.LayerNorm(self.embed_dim)179        self.fc = nn.Linear(self.embed_dim, intermediate_size, bias=False)180        self.act = nn.ReLU()181        self.c_proj = nn.Linear(intermediate_size, self.embed_dim, bias=False)182 183    def forward(self, hidden_states: Optional[tuple[torch.FloatTensor]]) -> torch.FloatTensor:184        hidden_states = self.ln(hidden_states)185        hidden_states = self.fc(hidden_states)186        hidden_states = self.act(hidden_states)187        hidden_states = self.c_proj(hidden_states)188 189        return hidden_states190