Move transformer impl to flava folder (#439)

Summary:
Only used by flava so moving it to that folder to make space for the other transformer

Pull Request resolved: #439

Test Plan:
pytest tests/
sanity check
torchrun --nproc_per_node=1 -m flava.native.train config=flava/native/configs/pretrain_debug.yaml

Reviewed By: ebsmothers

Differential Revision: D47839480

Pulled By: ankitade

fbshipit-source-id: ed635ae192baa16ee1a244b2e8e59bad37a80ad4

Author: ankitade

examples/flava/native/train.py

@@ -53,8 +53,10 @@
 from torch.utils.tensorboard import SummaryWriter
 from torchmultimodal.models.flava.image_encoder import ImageTransformer
 from torchmultimodal.models.flava.text_encoder import BERTTextEncoder
-from torchmultimodal.models.flava.transformer import FLAVATransformerWithoutEmbeddings
-from torchmultimodal.modules.layers.transformer import TransformerEncoderLayer
+from torchmultimodal.models.flava.transformer import (
+    FLAVATransformerWithoutEmbeddings,
+    TransformerEncoderLayer,
+)
 from torchmultimodal.modules.losses.flava import FLAVAPretrainingLossOutput
 
 

examples/flava/requirements.txt

@@ -7,3 +7,4 @@ omegaconf==2.1.2
 hydra-core==1.1.2
 transformers==4.30.0
 pycocotools==2.0.4
+tensorboard

tests/models/flava/test_image_encoder.py

@@ -10,7 +10,7 @@
 from tests.test_utils import assert_expected, set_rng_seed
 from torch import nn
 from torchmultimodal.models.flava.image_encoder import ImageEmbeddings, ImageTransformer
-from torchmultimodal.modules.layers.transformer import TransformerEncoder
+from torchmultimodal.models.flava.transformer import TransformerEncoder
 
 
 @pytest.fixture(autouse=True)

tests/models/flava/test_text_encoder.py

@@ -11,10 +11,12 @@
 import torch
 from tests.test_utils import assert_expected, set_rng_seed
 from torch import nn
-from torchmultimodal.models.flava.transformer import init_transformer_weights
+from torchmultimodal.models.flava.transformer import (
+    init_transformer_weights,
+    TransformerEncoder,
+)
 from torchmultimodal.modules.encoders.bert_text_encoder import BERTTextEncoder
 from torchmultimodal.modules.layers.text_embedding import BERTTextEmbeddings
-from torchmultimodal.modules.layers.transformer import TransformerEncoder
 
 
 @pytest.fixture(autouse=True)

tests/modules/layers/test_transformer.py renamed to tests/models/flava/test_transformer.py

@@ -9,7 +9,7 @@
 import torch
 from tests.test_utils import assert_expected, set_rng_seed
 from torch import nn
-from torchmultimodal.modules.layers.transformer import (
+from torchmultimodal.models.flava.transformer import (
     TransformerEncoder,
     TransformerEncoderLayer,
 )

torchmultimodal/models/flava/image_encoder.py

@@ -11,12 +11,12 @@
 
 import torch
 from torch import nn, Tensor
-from torchmultimodal.models.flava.transformer import init_transformer_weights
-from torchmultimodal.modules.layers.normalizations import Fp32LayerNorm
-from torchmultimodal.modules.layers.transformer import (
+from torchmultimodal.models.flava.transformer import (
+    init_transformer_weights,
     TransformerEncoder,
-    TransformerOutput,
 )
+from torchmultimodal.modules.layers.normalizations import Fp32LayerNorm
+from torchmultimodal.modules.layers.transformer import TransformerOutput
 from torchmultimodal.modules.losses.flava import Pooler
 
 

torchmultimodal/models/flava/model.py

@@ -17,13 +17,13 @@
 from torch import nn, Tensor
 from torchmultimodal.models.flava.image_encoder import flava_image_encoder
 from torchmultimodal.models.flava.text_encoder import flava_text_encoder
-from torchmultimodal.models.flava.transformer import FLAVATransformerWithoutEmbeddings
-from torchmultimodal.modules.layers.mlp import MLP
-from torchmultimodal.modules.layers.normalizations import Fp32LayerNorm
-from torchmultimodal.modules.layers.transformer import (
+from torchmultimodal.models.flava.transformer import (
+    FLAVATransformerWithoutEmbeddings,
     TransformerEncoder,
-    TransformerOutput,
 )
+from torchmultimodal.modules.layers.mlp import MLP
+from torchmultimodal.modules.layers.normalizations import Fp32LayerNorm
+from torchmultimodal.modules.layers.transformer import TransformerOutput
 from torchmultimodal.modules.losses.flava import (
     FLAVAPretrainingLoss,
     FLAVAPretrainingLossOutput,

torchmultimodal/models/flava/text_encoder.py

@@ -8,11 +8,13 @@
 from typing import Callable
 
 from torch import nn
-from torchmultimodal.models.flava.transformer import init_transformer_weights
+from torchmultimodal.models.flava.transformer import (
+    init_transformer_weights,
+    TransformerEncoder,
+)
 from torchmultimodal.modules.encoders.bert_text_encoder import BERTTextEncoder
 from torchmultimodal.modules.layers.normalizations import Fp32LayerNorm
 from torchmultimodal.modules.layers.text_embedding import BERTTextEmbeddings
-from torchmultimodal.modules.layers.transformer import TransformerEncoder
 from torchmultimodal.modules.losses.flava import Pooler
 
 

torchmultimodal/models/flava/transformer.py

@@ -5,10 +5,13 @@
 # LICENSE file in the root directory of this source tree.
 
 from functools import partial
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Optional, Tuple, Union
 
 import torch
 from torch import nn, Tensor
+from torchmultimodal.modules.layers.attention import MultiHeadAttention, SelfAttention
+from torchmultimodal.modules.layers.mlp import MLP
+from torchmultimodal.modules.layers.normalizations import Fp32LayerNorm
 
 from torchmultimodal.modules.layers.transformer import TransformerOutput
 
@@ -75,6 +78,223 @@ def forward(
         )
 
 
+class TransformerEncoderLayer(nn.Module):
+    """Transformer encoder layer is made up of multihead self-attention and feedforward blocks,
+    based on the architecture in "Attention Is All You Need" (Vaswani et al. 2017). Similar to
+    ``nn.TransformerEncoderLayer``, but uses a custom ``MultiHeadAttention`` that supports
+    n-dimensional inputs (including sequences, images, video) and head-masking.
+
+    Attributes:
+        d_model (int): size of hidden dimension of input
+        n_head (int): number of attention heads
+        dim_feedforward (int): size of hidden dimension of feedforward network
+        dropout (float): dropout probability for all dropouts. Defaults to 0.
+        activation (Callable): activation function in feedforward network. Defaults to ``nn.ReLU``.
+        layer_norm_eps (float): the eps value in layer norms. Default is 1e-12.
+        norm_first (bool): if True, layer norm is done prior to each of self-attention, cross-attention,
+            and feedforward. Otherwise, layer norm is done after.
+
+    Args:
+        hidden_states (Tensor): input tensor of shape [b, d1, ..., dn, c] to calculate self-attention on.
+        attention_mask (Tensor, optional): mask to be applied to self-attention inputs, ``hidden_states``. See
+            ``MultiHeadAttention`` for shape requirements.
+        head_mask (Tensor, optional): mask to be applied to self-attention inputs after softmax and dropout,
+            before matrix multiplication with values. See ``MultiHeadAttention`` for shape requirements.
+        return_attn_weights (bool, optional): return attention probabilities in addition to attention output.
+            Defaults to False.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_head: int,
+        dim_feedforward: int,
+        dropout: float = 0.0,
+        activation: Callable[..., nn.Module] = nn.ReLU,
+        layer_norm_eps: float = 1e-12,
+        norm_first: bool = False,
+    ) -> None:
+        super().__init__()
+        # attention block
+        self.attention = MultiHeadAttention(
+            dim_q=d_model,
+            dim_kv=d_model,
+            n_head=n_head,
+            attn_module=SelfAttention(dropout),
+        )
+        self.attention_dropout = nn.Dropout(dropout)
+        # feedforward block
+        self.feedforward = MLP(
+            d_model, d_model, dim_feedforward, dropout=dropout, activation=activation
+        )
+        self.feedforward_dropout = nn.Dropout(dropout)
+        # layernorms
+        self.attention_layernorm = Fp32LayerNorm(d_model, eps=layer_norm_eps)
+        self.feedforward_layernorm = Fp32LayerNorm(d_model, eps=layer_norm_eps)
+        self.norm_first = norm_first
+
+    def _attention_block(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        head_mask: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        output, attn_weights = self.attention(
+            hidden_states,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            return_attn_weights=True,
+        )
+        output = self.attention_dropout(output)
+        return output, attn_weights
+
+    def _feedforward_block(self, hidden_states: Tensor) -> Tensor:
+        h = self.feedforward(hidden_states)
+        h = self.feedforward_dropout(h)
+        return h
+
+    def _forward_prenorm(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        head_mask: Optional[Tensor] = None,
+        return_attn_weights: bool = False,
+    ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+        x = hidden_states
+        inputs = self.attention_layernorm(x)
+        attn_output, attn_weights = self._attention_block(
+            inputs,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+        )
+        attn_residual = attn_output + x
+        ff_residual = attn_residual + self._feedforward_block(
+            self.feedforward_layernorm(attn_residual)
+        )
+        if return_attn_weights:
+            return ff_residual, attn_weights
+        else:
+            return ff_residual
+
+    def _forward_postnorm(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        head_mask: Optional[Tensor] = None,
+        return_attn_weights: bool = False,
+    ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+        x = hidden_states
+        attn_output, attn_weights = self._attention_block(
+            x,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+        )
+        attn_residual = attn_output + x
+        attn_residual = self.attention_layernorm(attn_residual)
+        ff_residual = attn_residual + self._feedforward_block(attn_residual)
+        outputs = self.feedforward_layernorm(ff_residual)
+        if return_attn_weights:
+            return outputs, attn_weights
+        else:
+            return outputs
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        head_mask: Optional[Tensor] = None,
+        return_attn_weights: bool = False,
+    ) -> Union[Tensor, Tuple[Tensor, Tensor]]:
+        if self.norm_first:
+            return self._forward_prenorm(
+                hidden_states,
+                attention_mask,
+                head_mask,
+                return_attn_weights,
+            )
+        else:
+            return self._forward_postnorm(
+                hidden_states,
+                attention_mask,
+                head_mask,
+                return_attn_weights,
+            )
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(
+        self,
+        n_layer: int,
+        d_model: int,
+        n_head: int,
+        dim_feedforward: int,
+        dropout: float = 0.0,
+        activation: Callable[..., nn.Module] = nn.ReLU,
+        layer_norm_eps: float = 1e-12,
+        norm_first: bool = False,
+        final_layer_norm_eps: Optional[float] = None,
+    ):
+        super().__init__()
+        self.layer = nn.ModuleList(
+            [
+                TransformerEncoderLayer(
+                    d_model,
+                    n_head,
+                    dim_feedforward,
+                    dropout,
+                    activation,
+                    layer_norm_eps,
+                    norm_first,
+                )
+                for _ in range(n_layer)
+            ]
+        )
+        self.final_layer_norm = None
+        if final_layer_norm_eps:
+            self.final_layer_norm = Fp32LayerNorm(d_model, eps=final_layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: Tensor,
+        attention_mask: Optional[Tensor] = None,
+        head_mask: Optional[Tensor] = None,
+        return_attn_weights: bool = False,
+        return_hidden_states: bool = False,
+    ) -> TransformerOutput:
+
+        all_hidden_states = [] if return_hidden_states else None
+        all_self_attentions = [] if return_attn_weights else None
+
+        for layer_module in self.layer:
+            if return_hidden_states:
+                all_hidden_states.append(hidden_states)
+
+            layer_outputs = layer_module(
+                hidden_states,
+                attention_mask=attention_mask,
+                head_mask=head_mask,
+                return_attn_weights=return_attn_weights,
+            )
+
+            if return_attn_weights:
+                hidden_states = layer_outputs[0]
+                all_self_attentions.append(layer_outputs[1])
+            else:
+                hidden_states = layer_outputs
+
+        if return_hidden_states:
+            all_hidden_states.append(hidden_states)
+
+        if self.final_layer_norm is not None:
+            hidden_states = self.final_layer_norm(hidden_states)
+
+        return TransformerOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+        )
+
+
 def init_transformer_weights(module: nn.Module, initializer_range: float) -> None:
     """Initialize the weights"""
     if isinstance(module, (nn.Linear, nn.Conv2d)):

torchmultimodal/modules/encoders/bert_text_encoder.py

@@ -8,11 +8,9 @@
 
 import torch
 from torch import nn, Tensor
+from torchmultimodal.models.flava.transformer import TransformerEncoder
 from torchmultimodal.modules.layers.text_embedding import BERTTextEmbeddings
-from torchmultimodal.modules.layers.transformer import (
-    TransformerEncoder,
-    TransformerOutput,
-)
+from torchmultimodal.modules.layers.transformer import TransformerOutput
 from torchmultimodal.utils.attention import get_extended_attention_mask