use pyg model impl without changing the architecture

- #12

Author: aditya0by0

chebai_graph/models/base.py

@@ -49,20 +49,22 @@ def __init__(self, config: dict, **kwargs) -> None:
 
         Args:
             config (dict): Configuration dictionary with keys:
-                - 'hidden_length'
-                - 'dropout_rate'
-                - 'n_conv_layers'
-                - 'n_node_properties'
-                - 'n_bond_properties'
+                - 'num_layers'
+                - 'in_channels'
+                - 'hidden_channels'
+                - 'out_channels'
+                - 'edge_dim'
+                - 'dropout'
             **kwargs: Additional keyword arguments for torch.nn.Module.
         """
         super().__init__(**kwargs)
-        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
+        self.in_channels = int(config["in_channels"])  # number of node/atom properties
+        self.hidden_channels = int(config["hidden_channels"])
+        self.out_channels = int(config["out_channels"])
+        self.edge_dim = int(config["edge_dim"])  # number of bond properties
+        self.dropout = float(config["dropout"])
 
 
 class GraphNetWrapper(GraphBaseNet, ABC):

chebai_graph/models/gat.py

@@ -30,11 +30,12 @@ def __init__(self, config: dict, **kwargs):
         self.v2 = bool(config["v2"])
         self.activation = ELU()  # Instantiate ELU once for reuse.
         self.gat = GAT(
-            in_channels=self.n_atom_properties,
-            hidden_channels=self.hidden_length,
-            num_layers=self.n_conv_layers,
-            dropout=self.dropout_rate,
-            edge_dim=self.n_bond_properties,
+            in_channels=self.in_channels,
+            hidden_channels=self.hidden_channels,
+            out_channels=self.out_channels,
+            num_layers=self.num_layers,
+            dropout=self.dropout,
+            edge_dim=self.edge_dim,
             heads=self.heads,
             v2=self.v2,
             act=self.activation,

chebai_graph/models/resgated.py

@@ -1,7 +1,6 @@
-from typing import Final, Tuple, Union
+from typing import Any, Final
 
-import torch
-import torch.nn.functional as F
+from torch import Tensor
 from torch.nn import ELU
 from torch_geometric import nn as tgnn
 from torch_geometric.data import Data as GraphData
@@ -11,171 +10,77 @@
 from .base import GraphModelBase, GraphNetWrapper
 
 
-class ResGatedGraphConvNetBase(GraphModelBase):
-    """
-    Residual Gated Graph Convolutional Network with edge attributes support.
-
-    This model uses a stack of `ResGatedGraphConv` layers from PyTorch Geometric,
-    allowing edge attributes as part of message passing. A final projection layer maps
-    to the hidden length specified for downstream graph prediction tasks.
-    """
-
-    NAME = "ResGatedGraphConvNetBase"
-
-    def __init__(self, config: dict, **kwargs):
-        """
-        Initialize the ResGatedGraphConvNetBase.
-
-        Args:
-            config (dict): Configuration dictionary with keys:
-                - '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.activation = F.elu
-        self.convs = torch.nn.ModuleList()
-        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.hidden_channels,
-                    self.hidden_channels,
-                    edge_dim=self.n_bond_properties,
-                )
-            )
-
-        # Final projection layer to hidden dimension
-        self.final_conv = tgnn.ResGatedGraphConv(
-            self.hidden_channels, self.out_channels, edge_dim=self.n_bond_properties
-        )
-
-    def forward(self, batch: dict) -> torch.Tensor:
-        """
-        Forward pass through residual gated GNN layers.
-
-        Args:
-            batch (dict): A batch containing:
-                - 'features': A list with a `GraphData` instance as the first element.
-
-        Returns:
-            torch.Tensor: Node-level embeddings of shape [num_nodes, hidden_length].
-        """
-        graph_data = batch["features"][0]
-        assert isinstance(graph_data, GraphData)
-
-        x = graph_data.x.float()  # Atom features
-
-        for conv in self.convs:
-            assert isinstance(conv, tgnn.ResGatedGraphConv)
-            x = self.activation(
-                conv(x, graph_data.edge_index.long(), edge_attr=graph_data.edge_attr)
-            )
-
-        x = self.activation(
-            self.final_conv(
-                x, graph_data.edge_index.long(), edge_attr=graph_data.edge_attr
-            )
-        )
-
-        return x
-
-
-class ResGatedGraphPred(GraphNetWrapper):
+class ResGatedModel(BasicGNN):
     """
-    Residual Gated GNN for Graph Prediction.
+    A residual gated GNN model based on PyG's BasicGNN using ResGatedGraphConv layers.
 
-    Uses `ResGatedGraphConvNetBase` as the GNN encoder to compute node embeddings.
+    Attributes:
+        supports_edge_weight (bool): Indicates edge weights are not supported.
+        supports_edge_attr (bool): Indicates edge attributes are supported.
+        supports_norm_batch (bool): Indicates if batch normalization is supported.
     """
 
-    NAME = "ResGatedGraphPred"
+    supports_edge_weight: Final[bool] = False
+    supports_edge_attr: Final[bool] = True
+    supports_norm_batch: Final[bool]
 
-    def _get_gnn(self, config: dict) -> ResGatedGraphConvNetBase:
+    def init_conv(
+        self, in_channels: int | tuple[int, int], out_channels: int, **kwargs: Any
+    ) -> MessagePassing:
         """
-        Instantiate the residual gated GNN backbone.
+        Initializes a ResGatedGraphConv layer.
 
         Args:
-            config (dict): Model configuration.
+            in_channels (int or Tuple[int, int]): Number of input channels.
+            out_channels (int): Number of output channels.
+            **kwargs: Additional keyword arguments for the convolution layer.
 
         Returns:
-            ResGatedGraphConvNetBase: The GNN encoder.
+            MessagePassing: A ResGatedGraphConv layer instance.
         """
-        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):
+class ResGatedGraphConvNetBase(GraphModelBase):
     """
-    Graph Attention Network (GAT) base module for graph convolution.
+    Base model class for applying ResGatedGraphConv layers to graph-structured data.
 
-    Uses PyTorch Geometric's `GAT` implementation to process atomic node features
-    and bond edge attributes through multiple attention heads and layers.
+    Args:
+        config (dict): Configuration dictionary containing model hyperparameters.
+        **kwargs: Additional keyword arguments for parent class.
     """
 
-    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.
-        """
+    def __init__(self, config: dict[str, Any], **kwargs: Any):
         super().__init__(config=config, **kwargs)
         self.activation = ELU()  # Instantiate ELU once for reuse.
-        self.gat = ResGatedModel(
-            in_channels=self.n_node_properties,
+
+        self.resgated: BasicGNN = ResGatedModel(
+            in_channels=self.in_channels,
             hidden_channels=self.hidden_channels,
             out_channels=self.out_channels,
             num_layers=self.num_layers,
-            edge_dim=self.n_bond_properties,
+            edge_dim=self.edge_dim,
             act=self.activation,
         )
 
-    def forward(self, batch: dict) -> torch.Tensor:
+    def forward(self, batch: dict[str, Any]) -> Tensor:
         """
-        Forward pass through the GAT network.
-
-        Processes atomic node features and edge attributes, and applies
-        an ELU activation to the output.
+        Forward pass of the model.
 
         Args:
-            batch (dict): Input batch containing:
-                - 'features': A list with a `GraphData` object as its first element.
+            batch (dict): A batch containing graph input features under the key "features".
 
         Returns:
-            torch.Tensor: Node embeddings after GAT and activation.
+            Tensor: The output node-level embeddings after the final activation.
         """
         graph_data = batch["features"][0]
-        assert isinstance(graph_data, GraphData)
+        assert isinstance(graph_data, GraphData), "Expected GraphData instance"
 
-        out = self.gat(
+        out = self.resgated(
             x=graph_data.x.float(),
             edge_index=graph_data.edge_index.long(),
             edge_attr=graph_data.edge_attr,
@@ -184,23 +89,21 @@ def forward(self, batch: dict) -> torch.Tensor:
         return self.activation(out)
 
 
-class ResGatedGraphPredPyG(GraphNetWrapper):
+class ResGatedGraphPred(GraphNetWrapper):
     """
-    Residual Gated GNN for Graph Prediction.
+    Wrapper for graph-level prediction using ResGatedGraphConvNetBase.
 
-    Uses `ResGatedGraphConvNetBase` as the GNN encoder to compute node embeddings.
+    This class instantiates the core GNN model using the provided config.
     """
 
-    NAME = "ResGatedGraphPred"
-
-    def _get_gnn(self, config: dict) -> ResGatedPyG:
+    def _get_gnn(self, config: dict[str, Any]) -> ResGatedGraphConvNetBase:
         """
-        Instantiate the residual gated GNN backbone.
+        Returns the core ResGated GNN model.
 
         Args:
-            config (dict): Model configuration.
+            config (dict): Configuration dictionary for the GNN model.
 
         Returns:
-            ResGatedGraphConvNetBase: The GNN encoder.
+            ResGatedGraphConvNetBase: The core graph convolutional network.
         """
-        return ResGatedPyG(config=config)
+        return ResGatedGraphConvNetBase(config=config)

chebai_graph/preprocessing/datasets/chebi.py

@@ -267,8 +267,9 @@ def load_processed_data_from_file(self, filename: str) -> list[dict]:
         ]
         rank_zero_info(
             f"Finished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n"
-            f"Use n_atom_properties: {sum(p.encoder.get_encoding_length() for p in self.properties if isinstance(p, AtomProperty))}, "
-            f"n_bond_properties: {sum(p.encoder.get_encoding_length() for p in self.properties if isinstance(p, BondProperty))}, "
+            f"Use following values for given parameters for model configuration: \n\t"
+            f"in_channels: {sum(p.encoder.get_encoding_length() for p in self.properties if isinstance(p, AtomProperty))}, "
+            f"edge_dim: {sum(p.encoder.get_encoding_length() for p in self.properties if isinstance(p, BondProperty))}, "
             f"n_molecule_properties: {sum(p.encoder.get_encoding_length() for p in self.properties if isinstance(p, MoleculeProperty))}"
         )
 
@@ -337,7 +338,7 @@ def load_processed_data_from_file(self, filename: str) -> list[dict]:
             else:
                 raise TypeError(f"Unsupported property type: {type(prop).__name__}")
 
-        n_atom_properties = max(
+        n_node_properties = max(
             n_atom_node_properties, n_fg_node_properties, n_graph_node_properties
         )
         rank_zero_info(
@@ -347,7 +348,8 @@ def load_processed_data_from_file(self, filename: str) -> list[dict]:
             f"n_fg_node_properties: {n_fg_node_properties}, "
             f"n_bond_properties: {n_bond_properties}, "
             f"n_graph_node_properties: {n_graph_node_properties}\n"
-            f"Use n_atom_properties: {n_atom_properties}, n_bond_properties: {n_bond_properties}, n_molecule_properties: 0"
+            f"Use following values for given parameters for model configuration: \n\t"
+            f"in_channels: {n_node_properties}, edge_dim: {n_bond_properties}, n_molecule_properties: 0"
         )
 
         for property in self.properties:
@@ -368,7 +370,7 @@ def load_processed_data_from_file(self, filename: str) -> list[dict]:
         base_df["features"] = base_df.apply(
             lambda row: self._merge_props_into_base(
                 row,
-                max_len_node_properties=n_atom_properties,
+                max_len_node_properties=n_node_properties,
             ),
             axis=1,
         )

configs/model/gat.yml

@@ -3,13 +3,13 @@ init_args:
   optimizer_kwargs:
     lr: 1e-3
   config:
-    hidden_length: 512
-    dropout_rate: 0
-    n_conv_layers: 3
+    in_channels: 158 # number of node/atom properties
+    hidden_channels: 256
+    out_channels: 512
+    num_layers: 5
+    edge_dim: 7 # number of bond properties
     heads: 8  # the number of heads should be divisible by output channels (hidden channels if output channel not given)
     v2: False  # set True to use `torch_geometric.nn.conv.GATv2Conv` convolution layers, default is GATConv
-    n_atom_properties: 158
-    n_bond_properties: 7
-
-  n_molecule_properties: 200
-  n_linear_layers: 3
+    dropout: 0
+  n_molecule_properties: 0
+  n_linear_layers: 2

configs/model/resgated.yml

@@ -3,10 +3,11 @@ init_args:
   optimizer_kwargs:
     lr: 1e-3
   config:
-    n_node_properties: 68 # in_channels
-    hidden_channels : 256
-    out_channels : 512
-    num_layers : 4
-    n_bond_properties: 4 # edge_dim
+    in_channels: 158 # number of node/atom properties
+    hidden_channels: 256
+    out_channels: 512
+    num_layers: 5
+    edge_dim: 7 # number of bond properties
+    dropout: 0
   n_molecule_properties: 0
   n_linear_layers: 2