Upload model

config.json

@@ -1,5 +1,4 @@
 {
-  "_name_or_path": "./efficientnet/temp",
   "architectures": [
     "EfficientNetModelForImageClassification"
   ],
@@ -7,8 +6,10 @@
     "AutoConfig": "configuration_efficientnet.EfficientNetConfig",
     "AutoModelForImageClassification": "modeling_efficientnet.EfficientNetModelForImageClassification"
   },
+  "global_pool": "avg",
   "model_name": "efficientnet_b1",
   "model_type": "efficientnet",
+  "num_classes": 1000,
   "pretrained": true,
   "torch_dtype": "float32",
   "transformers_version": "4.49.0"

configuration_efficientnet.py

@@ -1,10 +1,6 @@
 from transformers.configuration_utils import PretrainedConfig
-
-from optimum.utils.normalized_config import NormalizedVisionConfig
-from optimum.utils.input_generators import DummyVisionInputGenerator
 from optimum.exporters.onnx.model_configs import ViTOnnxConfig
-
-from typing import OrderedDict, Dict
+from typing import Dict
 
 MODEL_NAMES = [
         'efficientnet_b0',
@@ -26,6 +22,8 @@ class EfficientNetConfig(PretrainedConfig):
             self,
             model_name: str = 'efficientnet_b0',
             pretrained: bool = False,
+            num_classes: int = 1000,
+            global_pool: str = 'avg',
             **kwargs
             ):
         
@@ -34,6 +32,8 @@ class EfficientNetConfig(PretrainedConfig):
         
         self.model_name = model_name
         self.pretrained = pretrained
+        self.num_classes = num_classes
+        self.global_pool = global_pool
        
         super().__init__(**kwargs)
 

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:58a67887d3b08b775276e0ed4ced6c66aeb03b0d9ad069a98a17440cc41557e6
+oid sha256:f1813e3c9f91308823701bc60e65f1417a1bc776274096c60784f4756a5a1d11
 size 31474952

modeling_efficientnet.py

@@ -1,46 +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)
-    
-    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)
-    
-    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"
 ]