Upload model

config.json

@@ -1,9 +1,13 @@
 {
-  "_attn_implementation_autoset": true,
+  "architectures": [
+    "MobileNetV3Model"
+  ],
   "auto_map": {
-    "AutoConfig": "config.MobileNetV3Config"
+    "AutoConfig": "config.MobileNetV3Config",
+    "AutoModelForImageClassification": "model.MobileNetV3Model"
   },
   "model_type": "mobilenetv3",
   "num_classes": 6,
+  "torch_dtype": "float32",
   "transformers_version": "4.46.3"
 }

model.py

@@ -0,0 +1,27 @@
+from transformers import PreTrainedModel
+from torchvision.models import mobilenet_v3_large, MobileNet_V3_Large_Weights
+from torch import nn
+import torch.nn.functional as F
+from .config import MobileNetV3Config
+
+class MobileNetV3Model(PreTrainedModel):
+    config_class = MobileNetV3Config
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = mobilenet_v3_large(weights=MobileNet_V3_Large_Weights.DEFAULT)
+        self.model.classifier = nn.Sequential(
+            nn.Linear(960, 1280),
+            nn.Hardswish(),
+            nn.Dropout(p=0.2, inplace=True),
+            nn.Linear(1280, config.num_classes),
+        )
+        
+    def forward(self, tensor, labels=None):
+        logits = self.model(tensor)
+        
+        if labels is not None:
+            loss = F.cross_entropy(logits, labels)
+            return {"loss": loss, "logits": logits}
+        
+        return {"logits": logits}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f56288776e34197e9911489341fbbcaa1a74119c47e77a0e92cddff7c01ba17d
+size 16968952