Create modeling_efficientnetv25.py

modeling_efficientnetv25.py

@@ -0,0 +1,38 @@
+from transformers import PreTrainedModel
+from .configuration_efficientnetv25 import EfficientNetV25Config
+import torch, sys, os
+from huggingface_hub import hf_hub_download
+
+class EfficientNetV25ForImageClassification(PreTrainedModel):
+    config_class = EfficientNetV25Config
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        repo_id = '/'.join(config.url.split('/')[3:5])
+        file_name = config.url.split('/')[-1]
+        path = f"./models/{file_name}"
+        if not os.path.exists(path):
+            hf_hub_download(repo_id=repo_id, filename=file_name, local_dir="./models")
+        
+        self.model = torch.load(path)
+        self.input_size = config.input_size
+        shape = [2] + self.input_size
+        example_inputs = torch.randn(shape)
+        example_inputs = (example_inputs - example_inputs.min()) / (example_inputs.max() - example_inputs.min())
+
+        self.num_classes = config.num_classes
+        if self.num_classes != 1000:
+            self.model.classifier = torch.nn.Linear(in_features=1984, out_features=self.num_classes, bias=True)
+        
+        traced_model = torch.jit.trace(self.model, example_inputs)
+        traced_model.save(file_name)
+
+        self.model = torch.jit.load(file_name)
+
+    def forward(self, tensor, labels=None):
+        logits = self.model(tensor)
+        if labels is not None:
+            loss = torch.nn.cross_entropy(logits, labels)
+            return {"loss": loss, "logits": logits}
+        return {"logits": logits}