Updated stuff

sagvit.bin → SAG-ViT.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ff3f3f603954fd987783ae22320d9e941ddb1e53acb9ce5b341d8911bdf3d100
-size 27102094
+oid sha256:8bdfa65b805d74c284af254153960436014a6e26b740f25cf2c6c9289234d9d3
+size 27137010

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2b23a28871ba0814540b11999cb6dc70b9c9e4d0bf446d2e8770e8c26d244387
-size 32389992
+oid sha256:234c74e97091f54931e7050f775e56465b0abff0596906ad015c45df31e7b12a
+size 27009200

model_components.py

@@ -1,7 +1,7 @@
 import torch
 from torch import nn
 import torch.nn.functional as F
-from torch_geometric.nn import GATConv, global_mean_pool
+from torch_geometric.nn import GATConv, global_mean_pool, GCNConv
 
 from torchvision import models
 
@@ -52,15 +52,36 @@ class GATGNN(nn.Module):
     This module corresponds to the Graph Attention stage (Section 3.3),
     refining local relationships between patches in a learned manner.
     """
-    def __init__(self, in_channels, hidden_channels, out_channels, heads=8):
+    def __init__(self, in_channels, hidden_channels, out_channels, heads=4):
         super(GATGNN, self).__init__()
         # GAT layers: 
         # First layer maps raw patch embeddings to a higher-level representation.
         self.conv1 = GATConv(in_channels, hidden_channels, heads=heads)
-        # Second layer produces final node embeddings with a single head.
-        self.conv2 = GATConv(hidden_channels * heads, out_channels, heads=1)
+        # Final GCN layer for refined representation
+        self.conv2 = GCNConv(hidden_channels * heads, out_channels)
         self.pool = global_mean_pool
 
+    def forward(self, data):
+        """
+        Input:
+        - data (PyG Data): Contains x (node features), edge_index (graph edges), and batch indexing.
+
+        Output:
+        - x (Tensor): Aggregated graph-level embedding after mean pooling.
+        """
+        x, edge_index, batch = data.x, data.edge_index, data.batch
+        
+        # GAT layer with ReLU activation
+        x = F.relu(self.conv1(x, edge_index))
+        
+        # GCN layer for further aggregation
+        x = self.conv2(x, edge_index)
+        
+        # Global mean pooling to obtain graph-level representation
+        out = self.pool(x, batch)
+        
+        return out
+
     def forward(self, data):
         """
         Input:

push_model_to_hfhub.py

@@ -1,13 +1,25 @@
+import torch
 from transformers import AutoConfig, AutoModel
 from modeling_sagvit import SAGViTClassifier, SAGViTConfig
 
 
+print("Registering model...")
 AutoConfig.register("sagvit", SAGViTConfig)
 AutoModel.register(SAGViTConfig, SAGViTClassifier)
+print("Registered model...")
 
 # Load config and model
 config = SAGViTConfig()
 model = SAGViTClassifier(config)
 
+# Load the state dict into the model
+print("Loading model weights...")
+state_dict = torch.load('SAG-ViT.pth')
+model.load_state_dict(state_dict)
+print("Loaded model weights...")
+
 # Push model and code
+model.save_pretrained('.')
 model.push_to_hub("shravvvv/SAG-ViT")
+
+print("Pushed model to Hugging Face hub...")