Upload model

config.json

@@ -1,11 +1,10 @@
 {
-  "_name_or_path": "./efficientnet/temp",
   "architectures": [
-    "EfficientNetModelForImageClassification"
+    "EfficientNetModel"
   ],
   "auto_map": {
     "AutoConfig": "configuration_efficientnet.EfficientNetConfig",
-    "AutoModelForImageClassification": "modeling_efficientnet.EfficientNetModelForImageClassification"
+    "AutoModel": "modeling_efficientnet.EfficientNetModel"
   },
   "global_pool": "avg",
   "model_name": "efficientnet_b1",

configuration_efficientnet.py

@@ -1,6 +1,7 @@
+from typing import Dict
+
 from transformers.configuration_utils import PretrainedConfig
 from optimum.exporters.onnx.model_configs import ViTOnnxConfig
-from typing import Dict
 
 MODEL_NAMES = [
         'efficientnet_b0',
@@ -15,18 +16,18 @@ MODEL_NAMES = [
         'efficientnet_l2'
     ]
 
+
 class EfficientNetConfig(PretrainedConfig):
     model_type = 'efficientnet'
 
     def __init__(
-            self,
-            model_name: str = 'efficientnet_b0',
-            pretrained: bool = False,
-            num_classes: int = 1000,
-            global_pool: str = 'avg',
-            **kwargs
-            ):
-        
+        self,
+        model_name: str = 'efficientnet_b0',
+        pretrained: bool = False,
+        num_classes: int = 1000,
+        global_pool: str = 'avg',
+        **kwargs,
+    ): 
         if model_name not in MODEL_NAMES:
             raise ValueError(f'`model_name` must be one of these: {MODEL_NAMES}, but got {model_name}')
         
@@ -34,9 +35,9 @@ class EfficientNetConfig(PretrainedConfig):
         self.pretrained = pretrained
         self.num_classes = num_classes
         self.global_pool = global_pool
-       
         super().__init__(**kwargs)
 
+
 class EfficientNetOnnxConfig(ViTOnnxConfig):
     @property
     def outputs(self) -> Dict[str, Dict[int, str]]:
@@ -47,6 +48,7 @@ class EfficientNetOnnxConfig(ViTOnnxConfig):
         
         return common_outputs
     
+
 __all__ = [
     'MODEL_NAMES',
     'EfficientNetConfig',

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,52 +1,126 @@
-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, Tensor, tensor
+from typing import Union, List, Optional
+
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPoolingAndNoAttention, 
+    ImageClassifierOutputWithNoAttention
+)
+from timm import create_model
+
+from .configuration_efficientnet import EfficientNetConfig
+
+
+class EfficientNetModel(PreTrainedModel):
+    """
+    EfficientNet model wrapper using Hugging Face's PreTrainedModel.
+
+    This class initializes an EfficientNet model from `timm` library 
+    and defines a forward method that extracts feature representations.
+
+    Attributes
+    ----------
+        config:  
+            Configuration object containing model parameters.
+        model:   
+            Instantiated EfficientNet model.
+    """
+    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: Tensor) -> BaseModelOutputWithPoolingAndNoAttention:
+        """
+        Parameters
+        ----------
+            pixel_values : torch.Tensor
+                Input tensor representing image pixel values.
+
+        Returns
+        -------
+            BaseModelOutputWithPoolingAndNoAttention
+                Object containing the `last_hidden_state` and `pooled_output`.
+        """
+        last_hidden_state = self.model.forward_features(pixel_values)
+        pooler_output = self.model.forward_head(last_hidden_state, pre_logits=True)
+        
+        return BaseModelOutputWithPoolingAndNoAttention(
+            last_hidden_state = last_hidden_state,
+            pooler_output=pooler_output
+        )
+
+
+class EfficientNetModelForImageClassification(PreTrainedModel):
+    """
+    EfficientNet model wrapper using Hugging Face's PreTrainedModel.
+
+    This class initializes an EfficientNet model from `timm` library 
+    and defines a forward method that return logits.
+
+    It supports training when labels are provided
+
+    Attributes
+    ----------
+        config :  
+            Configuration object containing model parameters.
+        model :   
+            Instantiated EfficientNet model.
+    """
+    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: Tensor, 
+        labels: Optional[Union[List[int], Tensor]] = None
+    ) -> ImageClassifierOutputWithNoAttention:
+        """
+        Parameters
+        ----------
+            pixel_values : torch.Tensor
+                Input tensor representing image pixel values.
+            labels : Optional[Union[List[int], torch.Tensor]]
+                Ground truth labels for training and computing loss. 
+                List of integers/tensor representing class IDs.
+
+        Returns
+        -------
+            ImageClassifierOutputWithNoAttention
+                Object containing `logits` and `loss`.
+        """
+        self.model.training = False if labels is None else True
+
+        logits = self.model(pixel_values)
+
+        loss = None
+        if self.model.training:
+            labels = tensor(labels)
+            ce_loss = nn.CrossEntropyLoss()
+            loss = ce_loss(logits, labels)
+
+        return ImageClassifierOutputWithNoAttention(
+            loss = loss,
+            logits = logits,
+        )
+    
+__all__ = [
+    "EfficientNetModel",
+    "EfficientNetModelForImageClassification"
 ]