Update modeling_internimage.py

modeling_internimage.py

@@ -800,23 +800,31 @@ class InternImage(nn.Module):
             'pooler_output': x if self.num_classes > 0 else None
         }
 
-    def forward(self, x):
+    def forward(self, 
+            pixel_values, 
+            output_attentions=None,
+            output_hidden_states=None,
+            return_dict=None):
         if self.use_clip_projector:  # for InternImage-H/G
-            outputs = self.forward_clip_projector(x)
+            outputs = self.forward_clip_projector(pixel_values)
         else:  # for InternImage-T/S/B/L/XL
-            outputs = self.forward_features(x)
+            outputs = self.forward_features(pixel_values)
 
-        hidden_states = outputs['hidden_states']
-        pooler_output = outputs['pooler_output']
+        hidden_states = outputs['hidden_states'] if output_hidden_states is not None else None
+        pooler_output = outputs['pooler_output'] if output_attentions is not None else None
+        last_hidden_state = outputs['hidden_states'][-1] if output_hidden_states is not None else None
 
         if self.num_classes > 0:
-            logits = self.head(pooler_output)
+            logits = self.head(outputs['pooler_output'])
         else:
             logits = None
 
+        if not return_dict:
+            return tuple(v for v in [logits, hidden_states, pooler_output, last_hidden_state] if v is not None)
+
         return BackboneOutput(
             hidden_states=hidden_states,
-            last_hidden_state=hidden_states[-1],
+            last_hidden_state=last_hidden_state,
             pooler_output=pooler_output,
             logits=logits
         )
@@ -853,8 +861,17 @@ class InternImageModel(PreTrainedModel):
             remove_center=config.remove_center,  # for InternImage-H/G
         )
 
-    def forward(self, pixel_values):
-        return self.model.forward_features(pixel_values)
+    def forward(self, 
+                pixel_values, 
+                output_attentions=None, 
+                output_hidden_states=None,
+                return_dict=None):
+        
+        return self.model.forward_features(
+            pixel_values,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict)
 
 
 class InternImageModelForImageClassification(PreTrainedModel):
@@ -862,6 +879,7 @@ class InternImageModelForImageClassification(PreTrainedModel):
 
     def __init__(self, config):
         super().__init__(config)
+        self.config = config
         self.model = InternImage(
             core_op=config.core_op,
             channels=config.channels,
@@ -888,12 +906,34 @@ class InternImageModelForImageClassification(PreTrainedModel):
             remove_center=config.remove_center,  # for InternImage-H/G
         )
 
-    def forward(self, pixel_values, labels=None):
-        outputs = self.model.forward(pixel_values)
-
+    def forward(self, 
+                pixel_values, 
+                labels=None, 
+                output_attentions=None, 
+                output_hidden_states=None, 
+                return_dict=None):
+        
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model.forward(
+            pixel_values, 
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict)
+
+        loss = None
         if labels is not None:
-            logits = outputs['logits']
+            logits = outputs.logits if return_dict else outputs[0]
             loss = F.cross_entropy(logits, labels)
-            outputs['loss'] = loss
 
-        return outputs
+        if not return_dict:
+            output = (outputs[0],) + outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return BackboneOutput(
+            loss = loss, 
+            logits = outputs.logits,
+            hidden_states = outputs.hidden_states,
+            last_hidden_state = outputs.last_hidden_state,
+            pooler_output = outputs.pooler_output
+        )