Feat/standard 2 layer gcn

astroml/models/gcn.py

@@ -1,35 +1,22 @@
-import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch_geometric.nn import GCNConv
 
 
 class GCN(nn.Module):
-    """
-    Configurable Graph Convolutional Network for node classification.
+    """Standard 2-layer Graph Convolutional Network for node classification.
+
+    Architecture: GCNConv -> ReLU -> Dropout -> GCNConv -> log_softmax
     """
 
-    def __init__(self, input_dim, hidden_dims, output_dim, dropout=0.5):
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, dropout: float = 0.5):
         super().__init__()
-
-        self.convs = nn.ModuleList()
+        self.conv1 = GCNConv(input_dim, hidden_dim)
+        self.conv2 = GCNConv(hidden_dim, output_dim)
         self.dropout = dropout
 
-        # Input layer
-        self.convs.append(GCNConv(input_dim, hidden_dims[0]))
-
-        # Hidden layers
-        for i in range(len(hidden_dims) - 1):
-            self.convs.append(GCNConv(hidden_dims[i], hidden_dims[i + 1]))
-
-        # Output layer
-        self.convs.append(GCNConv(hidden_dims[-1], output_dim))
-
     def forward(self, x, edge_index):
-        for conv in self.convs[:-1]:
-            x = conv(x, edge_index)
-            x = F.relu(x)
-            x = F.dropout(x, p=self.dropout, training=self.training)
-
-        x = self.convs[-1](x, edge_index)
-        return F.log_softmax(x, dim=1)
+        x = F.relu(self.conv1(x, edge_index))
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = self.conv2(x, edge_index)
+        return F.log_softmax(x, dim=1)

astroml/training/train_gcn.py

@@ -2,6 +2,7 @@
 import torch.nn.functional as F
 from torch_geometric.datasets import Planetoid
 from torch_geometric.transforms import NormalizeFeatures
+
 from astroml.models.gcn import GCN
 
 
@@ -13,36 +14,34 @@ def train():
 
     model = GCN(
         input_dim=dataset.num_node_features,
-        hidden_dims=[64],
+        hidden_dim=16,
         output_dim=dataset.num_classes,
         dropout=0.5,
     ).to(device)
 
     optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
 
-    model.train()
-    for epoch in range(200):
+    for epoch in range(1, 201):
+        model.train()
         optimizer.zero_grad()
         out = model(data.x, data.edge_index)
         loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
         loss.backward()
         optimizer.step()
 
         if epoch % 20 == 0:
-            print(f"Epoch {epoch}, Loss: {loss.item():.4f}")
+            val_acc = _accuracy(model, data, data.val_mask)
+            print(f"Epoch {epoch:3d} | Loss: {loss.item():.4f} | Val Acc: {val_acc:.4f}")
 
-    test(model, data)
+    print(f"Test Accuracy: {_accuracy(model, data, data.test_mask):.4f}")
 
 
-def test(model, data):
+def _accuracy(model: GCN, data, mask) -> float:
     model.eval()
-    out = model(data.x, data.edge_index)
-    pred = out.argmax(dim=1)
-
-    correct = (pred[data.test_mask] == data.y[data.test_mask]).sum()
-    acc = int(correct) / int(data.test_mask.sum())
-    print(f"Test Accuracy: {acc:.4f}")
+    with torch.no_grad():
+        pred = model(data.x, data.edge_index).argmax(dim=1)
+    return float((pred[mask] == data.y[mask]).sum()) / float(mask.sum())
 
 
 if __name__ == "__main__":
-    train()
+    train()