restructure resgated params - add pyg impl to compare

Author: aditya0by0

chebai_graph/models/base.py

@@ -52,16 +52,17 @@ def __init__(self, config: dict, **kwargs) -> None:
                 - 'hidden_length'
                 - 'dropout_rate'
                 - 'n_conv_layers'
-                - 'n_atom_properties'
+                - 'n_node_properties'
                 - 'n_bond_properties'
             **kwargs: Additional keyword arguments for torch.nn.Module.
         """
         super().__init__(**kwargs)
-        self.hidden_length = int(config["hidden_length"])
-        self.dropout_rate = float(config["dropout_rate"])
-        self.n_conv_layers = int(config["n_conv_layers"])
-        self.n_atom_properties = int(config["n_atom_properties"])
-        self.n_bond_properties = int(config["n_bond_properties"])
+        self.hidden_channels = int(config["hidden_channels"])
+        self.out_channels = int(config["out_channels"])
+        self.num_layers = int(config["num_layers"])
+        assert self.num_layers > 1, "Need atleast two convolution layers"
+        self.n_node_properties = int(config["n_node_properties"])  # in_channels
+        self.n_bond_properties = int(config["n_bond_properties"])  # edge_dim
 
 
 class GraphNetWrapper(GraphBaseNet, ABC):
@@ -83,9 +84,7 @@ def __init__(
         """
         super().__init__(**kwargs)
         self.gnn = self._get_gnn(config)
-        gnn_out_dim = (
-            config["out_dim"] if "out_dim" in config else config["hidden_length"]
-        )
+        gnn_out_dim = int(config["out_channels"])
         self.activation = torch.nn.ELU
         self.lin_input_dim = self._get_lin_seq_input_dim(
             gnn_out_dim=gnn_out_dim,

chebai_graph/models/resgated.py

@@ -1,8 +1,12 @@
+from typing import Final, Tuple, Union
+
 import torch
 import torch.nn.functional as F
-from torch import nn
+from torch.nn import ELU
 from torch_geometric import nn as tgnn
 from torch_geometric.data import Data as GraphData
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.nn.models.basic_gnn import BasicGNN
 
 from .base import GraphModelBase, GraphNetWrapper
 
@@ -24,38 +28,36 @@ def __init__(self, config: dict, **kwargs):
 
         Args:
             config (dict): Configuration dictionary with keys:
-                - 'in_length' (int): Intermediate feature length used in GNN layers.
+                - 'hidden_length' (int): Intermediate feature length used in GNN layers.
                 - Other parameters inherited from GraphModelBase.
             **kwargs: Additional keyword arguments passed to GraphModelBase.
         """
         super().__init__(config=config, **kwargs)
-        self.in_length = int(config["in_length"])
 
         self.activation = F.elu
-        self.dropout = nn.Dropout(self.dropout_rate)
-
         self.convs = torch.nn.ModuleList()
-        for i in range(self.n_conv_layers):
-            if i == 0:
-                # Initial layer uses atom features as input
-                self.convs.append(
-                    tgnn.ResGatedGraphConv(
-                        self.n_atom_properties,
-                        self.in_length,
-                        # dropout=self.dropout_rate,
-                        edge_dim=self.n_bond_properties,
-                    )
-                )
+        self.convs.append(
+            tgnn.ResGatedGraphConv(
+                self.n_node_properties,
+                self.hidden_channels,
+                # dropout=self.dropout,
+                edge_dim=self.n_bond_properties,
+            )
+        )
+
+        for _ in range(self.num_layers - 2):
             # Intermediate layers
             self.convs.append(
                 tgnn.ResGatedGraphConv(
-                    self.in_length, self.in_length, edge_dim=self.n_bond_properties
+                    self.hidden_channels,
+                    self.hidden_channels,
+                    edge_dim=self.n_bond_properties,
                 )
             )
 
         # Final projection layer to hidden dimension
         self.final_conv = tgnn.ResGatedGraphConv(
-            self.in_length, self.hidden_length, edge_dim=self.n_bond_properties
+            self.hidden_channels, self.out_channels, edge_dim=self.n_bond_properties
         )
 
     def forward(self, batch: dict) -> torch.Tensor:
@@ -109,3 +111,96 @@ def _get_gnn(self, config: dict) -> ResGatedGraphConvNetBase:
             ResGatedGraphConvNetBase: The GNN encoder.
         """
         return ResGatedGraphConvNetBase(config=config)
+
+
+class ResGatedModel(BasicGNN):
+    supports_edge_weight: Final[bool] = False
+    supports_edge_attr: Final[bool] = True
+    supports_norm_batch: Final[bool]
+
+    def init_conv(
+        self, in_channels: Union[int, Tuple[int, int]], out_channels: int, **kwargs
+    ) -> MessagePassing:
+        return tgnn.ResGatedGraphConv(
+            in_channels,
+            out_channels,
+            **kwargs,
+        )
+
+
+class ResGatedPyG(GraphModelBase):
+    """
+    Graph Attention Network (GAT) base module for graph convolution.
+
+    Uses PyTorch Geometric's `GAT` implementation to process atomic node features
+    and bond edge attributes through multiple attention heads and layers.
+    """
+
+    def __init__(self, config: dict, **kwargs):
+        """
+        Initialize the GATGraphConvNetBase.
+
+        Args:
+            config (dict): Model configuration containing:
+                - 'heads' (int): Number of attention heads.
+                - 'v2' (bool): Whether to use the GATv2 variant.
+                - Other required GraphModelBase parameters.
+            **kwargs: Additional arguments for the base class.
+        """
+        super().__init__(config=config, **kwargs)
+        self.activation = ELU()  # Instantiate ELU once for reuse.
+        self.gat = ResGatedModel(
+            in_channels=self.n_node_properties,
+            hidden_channels=self.hidden_channels,
+            out_channels=self.out_channels,
+            num_layers=self.num_layers,
+            edge_dim=self.n_bond_properties,
+            act=self.activation,
+        )
+
+    def forward(self, batch: dict) -> torch.Tensor:
+        """
+        Forward pass through the GAT network.
+
+        Processes atomic node features and edge attributes, and applies
+        an ELU activation to the output.
+
+        Args:
+            batch (dict): Input batch containing:
+                - 'features': A list with a `GraphData` object as its first element.
+
+        Returns:
+            torch.Tensor: Node embeddings after GAT and activation.
+        """
+        graph_data = batch["features"][0]
+        assert isinstance(graph_data, GraphData)
+
+        out = self.gat(
+            x=graph_data.x.float(),
+            edge_index=graph_data.edge_index.long(),
+            edge_attr=graph_data.edge_attr,
+        )
+
+        return self.activation(out)
+
+
+class ResGatedGraphPredPyG(GraphNetWrapper):
+    """
+    Residual Gated GNN for Graph Prediction.
+
+    Uses `ResGatedGraphConvNetBase` as the GNN encoder to compute node embeddings.
+    """
+
+    NAME = "ResGatedGraphPred"
+
+    def _get_gnn(self, config: dict) -> ResGatedPyG:
+        """
+        Instantiate the residual gated GNN backbone.
+
+        Args:
+            config (dict): Model configuration.
+
+        Returns:
+            ResGatedGraphConvNetBase: The GNN encoder.
+        """
+        return ResGatedPyG(config=config)

configs/model/resgated.yml

@@ -3,11 +3,10 @@ init_args:
   optimizer_kwargs:
     lr: 1e-3
   config:
-    in_length: 256
-    hidden_length: 512
-    dropout_rate: 0
-    n_conv_layers: 3
-    n_atom_properties: 158
-    n_bond_properties: 7
+    n_node_properties: 68 # in_channels
+    hidden_channels : 256
+    out_channels : 512
+    num_layers : 4
+    n_bond_properties: 4 # edge_dim
   n_molecule_properties: 0
   n_linear_layers: 2