Update modeling_efficientnet.py

modeling_efficientnet.py

@@ -1,52 +1,52 @@
-from torch import nn
-
-from transformers import PreTrainedModel
-from transformers.modeling_outputs import BaseModelOutputWithPoolingAndNoAttention, ImageClassifierOutputWithNoAttention
-from timm import create_model
-
-from configuration_efficientnet import EfficientNetConfig
-
-class EfficientNetModel(PreTrainedModel):
-    config_class = EfficientNetConfig
-
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.config = config
-        self.model = create_model(config.model_name, 
-                                  pretrained = config.pretrained, 
-                                  num_classes = config.num_classes, 
-                                  global_pool = config.global_pool)
-    
-    def forward(self, pixel_values):
-        last_hidden_state = self.model.forward_features(pixel_values)
-        return BaseModelOutputWithPoolingAndNoAttention(
-            last_hidden_state = last_hidden_state
-        )
-
-class EfficientNetModelForImageClassification(PreTrainedModel):
-    config_class = EfficientNetConfig
-
-    def __init__(self, config):
-        super().__init__(config)
-
-        self.config = config
-        self.model = create_model(config.model_name, 
-                                  pretrained = config.pretrained, 
-                                  num_classes = config.num_classes, 
-                                  global_pool = config.global_pool)
-    
-    def forward(self, pixel_values, labels=None):
-        logits = self.model(pixel_values)
-        loss = None
-        if labels is not None:
-            loss = nn.CrossEntropyLoss(logits, labels)
-        return ImageClassifierOutputWithNoAttention(
-            loss = loss,
-            logits = logits
-        )
-    
-__all__ = [
-    "EfficientNetModel",
-    "EfficientNetModelForImageClassification"
+from torch import nn
+
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPoolingAndNoAttention, ImageClassifierOutputWithNoAttention
+from timm import create_model
+
+from .configuration_efficientnet import EfficientNetConfig
+
+class EfficientNetModel(PreTrainedModel):
+    config_class = EfficientNetConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.config = config
+        self.model = create_model(config.model_name, 
+                                  pretrained = config.pretrained, 
+                                  num_classes = config.num_classes, 
+                                  global_pool = config.global_pool)
+    
+    def forward(self, pixel_values):
+        last_hidden_state = self.model.forward_features(pixel_values)
+        return BaseModelOutputWithPoolingAndNoAttention(
+            last_hidden_state = last_hidden_state
+        )
+
+class EfficientNetModelForImageClassification(PreTrainedModel):
+    config_class = EfficientNetConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.config = config
+        self.model = create_model(config.model_name, 
+                                  pretrained = config.pretrained, 
+                                  num_classes = config.num_classes, 
+                                  global_pool = config.global_pool)
+    
+    def forward(self, pixel_values, labels=None):
+        logits = self.model(pixel_values)
+        loss = None
+        if labels is not None:
+            loss = nn.CrossEntropyLoss(logits, labels)
+        return ImageClassifierOutputWithNoAttention(
+            loss = loss,
+            logits = logits
+        )
+    
+__all__ = [
+    "EfficientNetModel",
+    "EfficientNetModelForImageClassification"
 ]