Upload model

ConvNextBase.py

@@ -0,0 +1,43 @@
+from transformers import PretrainedConfig, PreTrainedModel
+from typing import List
+from torchvision.models import convnext_base, ConvNeXt_Base_Weights
+
+
+class ConvNextBaseConfig(PretrainedConfig):
+    model_type = "ConvNext"
+
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+
+class ConvNextBaseModel(PreTrainedModel):
+    config_class = ConvNextBaseConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = convnext_base()
+
+    def forward(self, tensor):
+        return self.model(tensor)
+
+
+class ConvNextBaseModelForImageClassification(PreTrainedModel):
+    config_class = ConvNextBaseConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = convnext_base()
+
+    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}
+
+ConvNextBaseConfig.register_for_auto_class()
+ConvNextBaseModel.register_for_auto_class("AutoModel")
+ConvNextBaseModelForImageClassification.register_for_auto_class("AutoModelForImageClassification")

config.json

@@ -2,6 +2,10 @@
   "architectures": [
     "ConvNextBaseModelForImageClassification"
   ],
+  "auto_map": {
+    "AutoConfig": "ConvNextBase.ConvNextBaseConfig",
+    "AutoModelForImageClassification": "ConvNextBase.ConvNextBaseModelForImageClassification"
+  },
   "model_type": "ConvNext",
   "torch_dtype": "float32",
   "transformers_version": "4.24.0"

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0df9805053983ed8a274bdbea1a3321d2f4a2c69deaf8fe7821badb70349281d
+oid sha256:4c59b8953c27c63ced868f8a94fac3d90b0261c8e5a002078ec4cb202d19c85f
 size 354474029