Add BetterTransformer support for FlavaModel

docs/source/bettertransformer/overview.mdx

@@ -33,6 +33,7 @@ The list of supported model below:
 - [DeiT](https://arxiv.org/abs/2012.12877)
 - [Electra](https://arxiv.org/abs/2003.10555)
 - [Ernie](https://arxiv.org/abs/1904.09223)
+- [Flava](https://arxiv.org/abs/2112.04482)
 - [FSMT](https://arxiv.org/abs/1907.06616)
 - [HuBERT](https://arxiv.org/pdf/2106.07447.pdf)
 - [LayoutLM](https://arxiv.org/abs/1912.13318)

optimum/bettertransformer/models/__init__.py

@@ -18,6 +18,7 @@
     BartEncoderLayerBetterTransformer,
     BertLayerBetterTransformer,
     DistilBertLayerBetterTransformer,
+    FlavaLayerBetterTransformer,
     FSMTEncoderLayerBetterTransformer,
     MBartEncoderLayerBetterTransformer,
     ViltLayerBetterTransformer,
@@ -74,6 +75,7 @@
     # FSMTModel:
     "EncoderLayer": FSMTEncoderLayerBetterTransformer,
     "ViltLayer": ViltLayerBetterTransformer,
+    "FlavaLayer": FlavaLayerBetterTransformer,
 }
 
 

optimum/bettertransformer/models/encoder_models.py

@@ -1068,3 +1068,99 @@ def forward(self, hidden_states, attention_mask, position_bias=None, *_, **__):
         if hidden_states.is_nested and self.is_last_layer:
             hidden_states = hidden_states.to_padded_tensor(0.0)
         return (hidden_states, attention_mask)
+
+
+class FlavaLayerBetterTransformer(BetterTransformerBaseLayer):
+    def __init__(self, flava_layer, config):
+        r"""
+        A simple conversion of the FlavaLayer to its `BetterTransformer` implementation.
+
+        Args:
+            flava_layer (`torch.nn.Module`):
+                The original `FlavaLayer` where the weights needs to be retrieved.
+        """
+        super().__init__(config)
+        # In_proj layer
+        self.in_proj_weight = nn.Parameter(
+            torch.cat(
+                [
+                    flava_layer.attention.attention.query.weight,
+                    flava_layer.attention.attention.key.weight,
+                    flava_layer.attention.attention.value.weight,
+                ]
+            )
+        )
+        self.in_proj_bias = nn.Parameter(
+            torch.cat(
+                [
+                    flava_layer.attention.attention.query.bias,
+                    flava_layer.attention.attention.key.bias,
+                    flava_layer.attention.attention.value.bias,
+                ]
+            )
+        )
+
+        # Out proj layer
+        self.out_proj_weight = flava_layer.attention.output.dense.weight
+        self.out_proj_bias = flava_layer.attention.output.dense.bias
+
+        # Linear layer 1
+        self.linear1_weight = flava_layer.intermediate.dense.weight
+        self.linear1_bias = flava_layer.intermediate.dense.bias
+
+        # Linear layer 2
+        self.linear2_weight = flava_layer.output.dense.weight
+        self.linear2_bias = flava_layer.output.dense.bias
+
+        # Layer norm 1
+        self.norm1_eps = flava_layer.layernorm_before.eps
+        self.norm1_weight = flava_layer.layernorm_before.weight
+        self.norm1_bias = flava_layer.layernorm_before.bias
+
+        # Layer norm 2
+        self.norm2_eps = flava_layer.layernorm_after.eps
+        self.norm2_weight = flava_layer.layernorm_after.weight
+        self.norm2_bias = flava_layer.layernorm_after.bias
+
+        # Model hyper parameters
+        self.num_heads = flava_layer.attention.attention.num_attention_heads
+        self.embed_dim = int(flava_layer.attention.attention.attention_head_size * self.num_heads)
+
+        # Last step: set the last layer to `False` -> this will be set to `True` when converting the model
+        self.is_last_layer = False
+        self.norm_first = True
+
+        self.validate_bettertransformer()
+
+    def forward(self, hidden_states, *_, **__):
+        r"""
+        This is just a wrapper around the forward function proposed in:
+        https://github.com/huggingface/transformers/pull/19553
+        """
+        super().forward_checker()
+        attention_mask = None
+
+        hidden_states = torch._transformer_encoder_layer_fwd(
+            hidden_states,
+            self.embed_dim,
+            self.num_heads,
+            self.in_proj_weight,
+            self.in_proj_bias,
+            self.out_proj_weight,
+            self.out_proj_bias,
+            self.use_gelu,
+            self.norm_first,
+            self.norm1_eps,
+            self.norm1_weight,
+            self.norm1_bias,
+            self.norm2_weight,
+            self.norm2_bias,
+            self.linear1_weight,
+            self.linear1_bias,
+            self.linear2_weight,
+            self.linear2_bias,
+            attention_mask,
+        )
+        if hidden_states.is_nested and self.is_last_layer:
+            hidden_states = hidden_states.to_padded_tensor(0.0)
+        return (hidden_states,)

tests/bettertransformer/test_bettertransformer_vision.py

@@ -32,6 +32,7 @@
 
 ALL_VISION_TEXT_MODELS_TO_TEST = [
     "hf-internal-testing/tiny-vilt-random-vqa",
+    "hf-internal-testing/tiny-random-FlavaModel",
 ]
 
 
@@ -66,3 +67,20 @@ def prepare_inputs_for_class(self, model_id=None):
         processor = AutoProcessor.from_pretrained(model_id)
         inputs = processor(image, text, return_tensors="pt")
         return inputs
+
+
+class BetterTransformersFlavaTest(BetterTransformersTestMixin, unittest.TestCase):
+    r"""
+    Testing suite for Vision and Text Models - tests all the tests defined in `BetterTransformersTestMixin`
+    """
+    all_models_to_test = ALL_VISION_TEXT_MODELS_TO_TEST
+
+    def prepare_inputs_for_class(self, model_id=None):
+        url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        image = Image.open(requests.get(url, stream=True).raw)
+        text = "How many cats are there?"
+
+        # Model takes image and text as input
+        processor = AutoProcessor.from_pretrained(model_id)
+        inputs = processor(image, text, return_tensors="pt")
+        return inputs