GAT option for GCN

Author: edebrouwer

topognn/layers.py

@@ -2,7 +2,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch_geometric.nn import GCNConv, GINConv
+from torch_geometric.nn import GCNConv, GINConv, GATConv
 from torch_scatter import scatter
 from torch_persistent_homology.persistent_homology_cpu import compute_persistence_homology_batched_mt
 from topognn.data_utils import remove_duplicate_edges
@@ -67,6 +67,38 @@ def forward(self, x, edge_index, **kwargs):
         return self.dropout(h)
 
 
+class GATLayer(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        activation,
+        dropout,
+        batch_norm,
+        num_heads,
+        residual=True,
+        train_eps=False,
+        **kwargs
+    ):
+        super().__init__()
+
+
+        self.activation = activation
+        self.residual = residual
+        self.dropout = nn.Dropout(dropout)
+        self.batchnorm = nn.BatchNorm1d(
+            out_features * num_heads) if batch_norm else nn.Identity()
+        
+        self.conv = GATConv(in_features, out_features, heads = num_heads, dropout = dropout)
+
+    def forward(self, x, edge_index, **kwargs):
+        h = self.conv(x, edge_index)
+        h = self.batchnorm(h)
+        h = self.activation(h)
+        if self.residual:
+            h = h + x
+        return self.dropout(h)
+
 class DeepSetLayer(nn.Module):
     """Simple equivariant deep set layer."""
 

topognn/models.py

@@ -11,7 +11,7 @@
 
 from topognn import Tasks
 from topognn.cli_utils import str2bool, int_or_none
-from topognn.layers import GCNLayer, GINLayer, SimpleSetTopoLayer, fake_persistence_computation
+from topognn.layers import GCNLayer, GINLayer, GATLayer, SimpleSetTopoLayer, fake_persistence_computation
 from topognn.metrics import WeightedAccuracy
 from topognn.data_utils import remove_duplicate_edges
 from torch_persistent_homology.persistent_homology_cpu import compute_persistence_homology_batched_mt
@@ -529,24 +529,35 @@ def forward(self,x, **kwargs):
 
 class LargerGCNModel(pl.LightningModule):
     def __init__(self, hidden_dim, depth, num_node_features, num_classes, task,
-                 lr=0.001, dropout_p=0.2, GIN=False, batch_norm=False,
+                 lr=0.001, dropout_p=0.2, GIN=False, GAT = False, batch_norm=False,
                  residual=False, train_eps=True, save_filtration = False,
                  add_mlp=False, weight_decay = 0., **kwargs):
         super().__init__()
         self.save_hyperparameters()
-        self.embedding = torch.nn.Linear(num_node_features, hidden_dim)
         self.save_filtration = save_filtration
 
+        num_heads = 1
+
         if GIN:
             def build_gnn_layer():
                 return GINLayer( in_features = hidden_dim, out_features = hidden_dim, train_eps=train_eps, activation = F.relu, batch_norm = batch_norm, dropout = dropout_p, **kwargs)
             graph_pooling_operation = global_add_pool
+
+        elif GAT:
+            num_heads = 8
+            def build_gnn_layer():
+                return GATLayer( in_features = hidden_dim * num_heads, out_features = hidden_dim, train_eps=train_eps, activation = F.relu, batch_norm = batch_norm, dropout = dropout_p, num_heads = num_heads, **kwargs)
+            graph_pooling_operation = global_mean_pool
+            
         else:
             def build_gnn_layer():
                 return GCNLayer(
                     hidden_dim, hidden_dim, F.relu, dropout_p, batch_norm)
             graph_pooling_operation = global_mean_pool
 
+        
+        self.embedding = torch.nn.Linear(num_node_features, hidden_dim * num_heads)
+        
         layers = [build_gnn_layer() for _ in range(depth)]
 
         if add_mlp:
@@ -567,7 +578,7 @@ def fake_pool(x, batch):
         if (kwargs.get("dim1",False) and ("dim1_out_dim" in kwargs.keys()) and ( not kwargs.get("fake",False))):
             dim_before_class = hidden_dim + kwargs["dim1_out_dim"] #SimpleTopoGNN with dim1
         else:
-            dim_before_class = hidden_dim
+            dim_before_class = hidden_dim * num_heads
 
         self.classif = torch.nn.Sequential(
             nn.Linear(dim_before_class, hidden_dim // 2),
@@ -633,6 +644,7 @@ def configure_optimizers(self):
         return [optimizer], [scheduler]
 
     def forward(self, data):
+
         x, edge_index = data.x, data.edge_index
         x = self.embedding(x)
 
@@ -722,6 +734,7 @@ def add_model_specific_args(cls, parent):
         parser.add_argument("--min_lr", type=float, default=0.00001)
         parser.add_argument("--dropout_p", type=float, default=0.0)
         parser.add_argument('--GIN', type=str2bool, default=False)
+        parser.add_argument('--GAT', type=str2bool, default=False)
         parser.add_argument('--train_eps', type=str2bool, default=True)
         parser.add_argument('--batch_norm', type=str2bool, default=True)
         parser.add_argument('--residual', type=str2bool, default=True)