Upload model

ViTMAETimm.py

@@ -0,0 +1,42 @@
+from transformers import PreTrainedModel
+from timm.models.resnet import BasicBlock, Bottleneck, ResNet
+
+from transformers import PretrainedConfig
+from typing import List
+import torch
+import timm
+
+
+class ViTMAEConfig(PretrainedConfig):
+    model_type = "vit_mae_custom"
+
+    def __init__( self,  model_name='timm/vit_base_patch16_224.mae', num_classes: int = 1000, **kwargs ):
+        self.model_name = model_name
+        self.num_classes = num_classes
+        super().__init__(**kwargs)
+        
+# 'timm/vit_huge_patch14_224.mae'
+# class ViTMAEModel(PreTrainedModel):
+#     config_class = ViTMAEConfig
+
+#     def __init__(self, config):
+#         super().__init__(config)
+#         self.model = timm.create_model(config.model_name, num_classes=config.num_classes, pretrained=True)
+
+#     def forward(self, tensor):
+#         return self.model.forward_features(tensor)
+
+class ViTMAEModelForImageClassification(PreTrainedModel):
+    config_class = ViTMAEConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+  
+        self.model = timm.create_model(config.model_name, num_classes=config.num_classes, pretrained=True)
+
+    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}

config.json

@@ -0,0 +1,14 @@
+{
+  "architectures": [
+    "ViTMAEModelForImageClassification"
+  ],
+  "auto_map": {
+    "AutoConfig": "ViTMAETimm.ViTMAEConfig",
+    "AutoModelForImageClassification": "ViTMAETimm.ViTMAEModelForImageClassification"
+  },
+  "model_name": "timm/vit_huge_patch14_224.mae",
+  "model_type": "vit_mae_custom",
+  "num_classes": 2,
+  "torch_dtype": "float32",
+  "transformers_version": "4.32.0.dev0"
+}

pytorch_model.bin

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