Registed model and added config

register_model.py

@@ -0,0 +1,12 @@
+from transformers import AutoConfig, AutoModel
+from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+from transformers.models.auto.modeling_auto import MODEL_MAPPING
+
+from sagvit_config import SAGViTConfig
+from sag_vit_model import SAGViTClassifier
+
+# Register the configuration
+CONFIG_MAPPING.register("sagvit", SAGViTConfig)
+
+# Register the model
+MODEL_MAPPING.register(SAGViTConfig, SAGViTClassifier)

sagvit_config.py

@@ -0,0 +1,28 @@
+from transformers import PretrainedConfig
+
+class SAGViTConfig(PretrainedConfig):
+    model_type = "sagvit"
+
+    def __init__(self,
+                 d_model=64,
+                 dim_feedforward=64,
+                 gcn_hidden=128,
+                 gcn_out=64,
+                 hidden_mlp_features=64,
+                 in_channels=2560,
+                 nhead=4,
+                 num_classes=10,
+                 num_layers=2,
+                 patch_size=(4, 4),
+                 **kwargs):
+        super().__init__(**kwargs)
+        self.d_model = d_model
+        self.dim_feedforward = dim_feedforward
+        self.gcn_hidden = gcn_hidden
+        self.gcn_out = gcn_out
+        self.hidden_mlp_features = hidden_mlp_features
+        self.in_channels = in_channels
+        self.nhead = nhead
+        self.num_classes = num_classes
+        self.num_layers = num_layers
+        self.patch_size = patch_size