Update modeling_intern_vit.py

modeling_intern_vit.py

@@ -219,12 +219,11 @@ class InternAttention(nn.Module):
 
         attn = ((q * self.scale) @ k.transpose(-2, -1))
         attn = attn.softmax(dim=-1)
-        attn = self.attn_drop(attn)
 
-        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = (self.attn_drop(attn) @ v).transpose(1, 2).reshape(B, N, C)
         x = self.proj(x)
         x = self.proj_drop(x)
-        return x
+        return x, attn
 
     def _flash_attn(self, x, key_padding_mask=None, need_weights=False):
         qkv = self.qkv(x)
@@ -243,8 +242,11 @@ class InternAttention(nn.Module):
         outs = self.proj_drop(outs)
         return outs
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        x = self._naive_attn(hidden_states) if not self.use_flash_attn else self._flash_attn(hidden_states)
+    def forward(self, hidden_states: torch.Tensor, output_attentions: Optional[bool] = False,) -> torch.Tensor:
+        if not self.use_flash_attn:
+            x = self._naive_attn(hidden_states)
+        else:
+            x = self._flash_attn(hidden_states)
         return x
 
 
@@ -283,23 +285,37 @@ class InternVisionEncoderLayer(nn.Module):
     def forward(
             self,
             hidden_states: torch.Tensor,
+            output_attentions: Optional[bool] = False,
     ) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor], Optional[Tuple[torch.FloatTensor]]]:
         """
         Args:
             hidden_states (`Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]`): input to the layer of shape `(batch, seq_len, embed_dim)`
         """
-        hidden_states = hidden_states + self.drop_path1(self.attn(self.norm1(hidden_states).to(hidden_states.dtype)) * self.ls1)
+        residual = hidden_states
 
-        hidden_states = hidden_states + self.drop_path2(self.mlp(self.norm2(hidden_states).to(hidden_states.dtype)) * self.ls2)
+        hidden_states = self.norm1(hidden_states).to(hidden_states.dtype)
+        hidden_states, attn_weights = self.attn(
+            hidden_states=hidden_states,
+            output_attentions=output_attentions,
+        )
+        hidden_states = residual + self.drop_path1(hidden_states * self.ls1)
 
-        return hidden_states
+        residual = hidden_states
+        hidden_states = self.norm2(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + self.drop_path2(hidden_states * self.ls2)
 
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (attn_weights,)
+
+        return outputs
 
 class InternVisionEncoder(nn.Module):
     """
     Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
     [`InternEncoderLayer`].
-
     Args:
         config (`InternConfig`):
             The corresponding vision configuration for the `InternEncoder`.
@@ -318,6 +334,7 @@ class InternVisionEncoder(nn.Module):
             self,
             inputs_embeds,
             output_hidden_states: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
             return_dict: Optional[bool] = None,
     ) -> Union[Tuple, BaseModelOutput]:
         r"""
@@ -336,6 +353,8 @@ class InternVisionEncoder(nn.Module):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
         encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
         hidden_states = inputs_embeds
 
         for idx, encoder_layer in enumerate(self.layers):
@@ -348,8 +367,12 @@ class InternVisionEncoder(nn.Module):
             else:
                 layer_outputs = encoder_layer(
                     hidden_states,
+                    output_attentions=output_attentions,
                 )
-            hidden_states = layer_outputs
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
 
         if output_hidden_states:
             encoder_states = encoder_states + (hidden_states,)
@@ -357,7 +380,7 @@ class InternVisionEncoder(nn.Module):
         if not return_dict:
             return tuple(v for v in [hidden_states, encoder_states] if v is not None)
         return BaseModelOutput(
-            last_hidden_state=hidden_states, hidden_states=encoder_states
+            last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
         )
 
 
@@ -393,6 +416,7 @@ class InternVisionModel(PreTrainedModel):
             self,
             pixel_values: Optional[torch.FloatTensor] = None,
             output_hidden_states: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
             return_dict: Optional[bool] = None,
             pixel_embeds: Optional[torch.FloatTensor] = None,
     ) -> Union[Tuple, BaseModelOutputWithPooling]:
@@ -414,6 +438,7 @@ class InternVisionModel(PreTrainedModel):
         encoder_outputs = self.encoder(
             inputs_embeds=hidden_states,
             output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
             return_dict=return_dict,
         )
         last_hidden_state = encoder_outputs.last_hidden_state