Add more gnn baselines

README.md

@@ -76,7 +76,7 @@ The list can be found in the `configs/data/chebi50_graph_properties.yml` file.
 python -m chebai fit --trainer=configs/training/default_trainer.yml --trainer.logger=configs/training/csv_logger.yml --model=../python-chebai-graph/configs/model/gnn_res_gated.yml --model.train_metrics=configs/metrics/micro-macro-f1.yml --model.test_metrics=configs/metrics/micro-macro-f1.yml --model.val_metrics=configs/metrics/micro-macro-f1.yml --data=../python-chebai-graph/configs/data/chebi50_graph_properties.yml --data.init_args.batch_size=128 --trainer.accumulate_grad_batches=4 --data.init_args.num_workers=10 --model.pass_loss_kwargs=false --data.init_args.chebi_version=241 --trainer.min_epochs=200 --trainer.max_epochs=200 --model.criterion=configs/loss/bce_weighted.yml
 ```
 
-## Augmented Graphs 
+## Augmented Graphs
 _See thesis related to this work [here](https://www.uni-osnabrueck.de/fileadmin/informatik/Arbeitsgruppen/Hybride_KI/mt_aditya_khedekar.pdf)_.
 
 Graph Neural Networks (GNNs) often fail to explicitly leverage the chemically meaningful substructures present within molecules (i.e. **functional groups (FGs)**).  To make this implicit information explicitly accessible to GNNs, we augment molecular graphs with **artificial nodes** that represent these substructures. The resulting graph are referred to as **augmented graphs**.

chebai_graph/models/__init__.py

@@ -1,19 +1,25 @@
+from .architectures.gat import GATGraphPred
+from .architectures.gine import GINEGraphPred
+from .architectures.resgated import ResGatedGraphPred
 from .augmented import (
-    GATAugNodePoolGraphPred,
-    GATGraphNodeFGNodePoolGraphPred,
-    ResGatedAugNodePoolGraphPred,
-    ResGatedGraphNodeFGNodePoolGraphPred,
+    GATAAPoolGraphPred,
+    GATAMGPoolGraphPred,
+    GINEAAPoolGraphPred,
+    GINEAMGPoolGraphPred,
+    ResGatedAAPoolGraphPred,
+    ResGatedAMGPoolGraphPred,
 )
 from .dynamic_gni import ResGatedDynamicGNIGraphPred
-from .gat import GATGraphPred
-from .resgated import ResGatedGraphPred
 
 __all__ = [
     "ResGatedGraphPred",
-    "ResGatedAugNodePoolGraphPred",
-    "ResGatedGraphNodeFGNodePoolGraphPred",
-    "GATGraphPred",
-    "GATAugNodePoolGraphPred",
-    "GATGraphNodeFGNodePoolGraphPred",
+    "ResGatedAAPoolGraphPred",
+    "ResGatedAMGPoolGraphPred",
     "ResGatedDynamicGNIGraphPred",
+    "GATGraphPred",
+    "GATAAPoolGraphPred",
+    "GATAMGPoolGraphPred",
+    "GINEGraphPred",
+    "GINEAAPoolGraphPred",
+    "GINEAMGPoolGraphPred",
 ]

chebai_graph/models/architectures/base.py

@@ -0,0 +1,185 @@
+from abc import ABC, abstractmethod
+
+import torch
+from chebai.models.base import ChebaiBaseNet
+from chebai.preprocessing.structures import XYData
+from torch_geometric.data import Data as GraphData
+from torch_scatter import scatter_add
+
+
+class GraphBaseNet(ChebaiBaseNet, ABC):
+    """
+    Base class for graph-based prediction networks.
+    """
+
+    def _get_prediction_and_labels(
+        self, data: XYData, labels: torch.Tensor, output: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Apply sigmoid activation to outputs and return processed labels.
+
+        Args:
+            data (XYData): Input batch data.
+            labels (torch.Tensor): Ground-truth labels.
+            output (torch.Tensor): Raw model output.
+
+        Returns:
+            tuple[torch.Tensor, torch.Tensor]: Tuple of (predictions, labels).
+        """
+        return torch.sigmoid(output), labels.int()
+
+    def _process_labels_in_batch(self, batch: XYData) -> torch.Tensor | None:
+        """
+        Process labels from XYData batch.
+
+        Returns:
+            torch.Tensor | None: Processed labels if present, else None.
+        """
+        return batch.y.float() if batch.y is not None else None
+
+
+class GraphModelBase(torch.nn.Module, ABC):
+    """
+    Abstract base class for graph models with configurable architecture.
+    """
+
+    def __init__(self, config: dict, **kwargs) -> None:
+        """
+        Initialize model hyperparameters from configuration.
+
+        Args:
+            config (dict): Configuration dictionary with keys:
+                - 'num_layers'
+                - 'in_channels'
+                - 'hidden_channels'
+                - 'out_channels'
+                - 'edge_dim'
+                - 'dropout'
+            **kwargs: Additional keyword arguments for torch.nn.Module.
+        """
+        super().__init__(**kwargs)
+        self.num_layers = int(config["num_layers"])
+        assert self.num_layers > 1, "Need atleast two convolution layers"
+        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):
+    """
+    Base wrapper class for GNNs with linear layers for graph classification
+    with standard pooling .
+    """
+
+    def __init__(
+        self,
+        config: dict,
+        n_linear_layers: int,
+        use_batch_norm: bool = False,
+        **kwargs,
+    ):
+        """
+        Args:
+            config (dict): Model configuration.
+            n_linear_layers (int): Number of linear layers.
+            **kwargs: Additional arguments.
+        """
+        super().__init__(**kwargs)
+        self.gnn = self._get_gnn(config)
+        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,
+        )
+        self.use_batch_norm = use_batch_norm
+        if self.use_batch_norm:
+            self.batch_norm = torch.nn.BatchNorm1d(self.lin_input_dim)
+
+        lin_hidden_dim = kwargs.get("lin_hidden_dim", gnn_out_dim)
+        self.lin_sequential: torch.nn.Sequential = self._get_linear_module_list(
+            n_linear_layers=n_linear_layers,
+            in_dim=self.lin_input_dim,
+            hidden_dim=lin_hidden_dim,
+            out_dim=self.out_dim,
+        )
+
+    @abstractmethod
+    def _get_gnn(self, config: dict) -> torch.nn.Module:
+        """
+        Create the graph neural network.
+
+        Args:
+            config (dict): Configuration dictionary.
+
+        Returns:
+            torch.nn.Module: Instantiated GNN module.
+        """
+        pass
+
+    def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int:
+        """
+        Compute input dimension for the linear layers.
+
+        Args:
+            gnn_out_dim (int): Output dimension of GNN.
+
+        Returns:
+            int: Total input dimension.
+        """
+        return gnn_out_dim
+
+    def _get_linear_module_list(
+        self,
+        n_linear_layers: int,
+        in_dim: int,
+        hidden_dim: int,
+        out_dim: int,
+    ) -> torch.nn.Sequential:
+        """
+        Construct a sequential module of linear layers.
+
+        Args:
+            n_linear_layers (int): Number of linear layers.
+            in_dim (int): Input dimension.
+            hidden_dim (int): Hidden dimension.
+            out_dim (int): Output dimension.
+
+        Returns:
+            torch.nn.Sequential: Linear layers with activations.
+        """
+        if n_linear_layers < 1:
+            raise ValueError("n_linear_layers must be at least 1")
+
+        layers = []
+        if n_linear_layers == 1:
+            layers.append(torch.nn.Linear(in_dim, out_dim))
+        else:
+            layers.append(torch.nn.Linear(in_dim, hidden_dim))
+            layers.append(self.activation())
+            for _ in range(n_linear_layers - 2):
+                layers.append(torch.nn.Linear(hidden_dim, hidden_dim))
+                layers.append(self.activation())
+            layers.append(torch.nn.Linear(hidden_dim, out_dim))
+
+        return torch.nn.Sequential(*layers)
+
+    def forward(self, batch: dict) -> torch.Tensor:
+        """
+        Forward pass through GNN, pooling and linear layers.
+
+        Args:
+            batch (dict): Input batch with graph features.
+
+        Returns:
+            torch.Tensor: Predicted output.
+        """
+        graph_data = batch["features"][0]
+        graph_data.to(self.device)
+        assert isinstance(graph_data, GraphData)
+        a = self.gnn(batch)
+        a = scatter_add(a, graph_data.batch, dim=0)
+        if self.use_batch_norm:
+            a = self.batch_norm(a)
+        return self.lin_sequential(a)

chebai_graph/models/architectures/gine.py

@@ -0,0 +1,146 @@
+from typing import Any, Final
+
+from torch import Tensor
+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 import MLP
+from torch_geometric.nn.models.basic_gnn import BasicGNN
+
+from .base import GraphModelBase, GraphNetWrapper
+
+
+class GINEModel(BasicGNN):
+    """
+    A GIN-based GNN model based on PyG's BasicGNN, using GINEConv layers so that
+    edge (bond) features are incorporated into the message-passing step.
+
+    See:
+        - https://pytorch-geometric.readthedocs.io/en/2.7.0/generated/torch_geometric.nn.conv.GINEConv.html
+        - https://arxiv.org/abs/1810.00826 (GIN)
+        - https://arxiv.org/abs/1905.12265 (GINE / edge-feature extension)
+        - https://github.com/pyg-team/pytorch_geometric/blob/master/examples/mutag_gin.py
+        - https://github.com/pyg-team/pytorch_geometric/issues/1311
+
+    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.
+    """
+
+    supports_edge_weight: Final[bool] = False
+    supports_edge_attr: Final[bool] = True
+    supports_norm_batch: Final[bool]
+
+    def init_conv(
+        self, in_channels: int | tuple[int, int], out_channels: int, **kwargs: Any
+    ) -> MessagePassing:
+        """
+        Initializes a GINEConv layer.
+
+        The inner network is a 2-layer MLP (Linear -> act -> Linear, no
+        activation on the last layer), matching both
+        `torch_geometric.nn.models.GIN.init_conv` and the GIN paper's
+        message-transform MLP. `edge_dim` (passed via **kwargs) lets
+        GINEConv linearly project bond features onto the node feature
+        space before adding them into the neighbor messages.
+
+        Args:
+            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
+                (e.g. `edge_dim`, `train_eps`).
+
+        Returns:
+            MessagePassing: A GINEConv layer instance.
+        """
+        mlp = MLP(
+            [in_channels, out_channels, out_channels],
+            act=self.act,
+            act_first=self.act_first,
+            norm=self.norm,
+            norm_kwargs=self.norm_kwargs,
+        )
+        return tgnn.GINEConv(mlp, **kwargs)
+
+
+class GINEConvNetBase(GraphModelBase):
+    """
+    Base model class for applying GINEConv layers to graph-structured data.
+
+    Based on:
+        - Xu et al., "How Powerful are Graph Neural Networks?"
+          (https://arxiv.org/abs/1810.00826)
+        - Hu et al., "Strategies for Pre-training Graph Neural Networks"
+          (https://arxiv.org/abs/1905.12265), reference implementation at
+          https://github.com/snap-stanford/pretrain-gnns/blob/master/chem/model.py
+
+    Args:
+        config (dict): Configuration dictionary containing model hyperparameters.
+            Also supports an optional `train_eps` (bool, default True) key,
+            which makes GINEConv's epsilon a learnable parameter, as
+            recommended in the original GIN paper.
+        **kwargs: Additional keyword arguments for parent class.
+    """
+
+    def __init__(self, config: dict[str, Any], **kwargs: Any):
+        super().__init__(config=config, **kwargs)
+        self.activation = ELU()  # Instantiate ELU once for reuse.
+        self.train_eps = bool(config.get("train_eps", True))
+
+        self.gine: BasicGNN = GINEModel(
+            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,
+            train_eps=self.train_eps,
+            act=self.activation,
+        )
+
+    def forward(self, batch: dict[str, Any]) -> Tensor:
+        """
+        Forward pass of the model.
+
+        Args:
+            batch (dict): A batch containing graph input features under the key "features".
+
+        Returns:
+            Tensor: The output node-level embeddings after the final activation.
+        """
+        graph_data = batch["features"][0]
+        assert isinstance(graph_data, GraphData), "Expected GraphData instance"
+
+        out = self.gine(
+            x=graph_data.x.float(),
+            edge_index=graph_data.edge_index.long(),
+            edge_attr=graph_data.edge_attr,
+        )
+
+        return self.activation(out)
+
+
+class GINEGraphPred(GraphNetWrapper):
+    """
+    Wrapper for graph-level prediction using GINEConvNetBase.
+
+    This class instantiates the core GNN model using the provided config.
+    Graph-level pooling (scatter-add over nodes) and the final linear
+    prediction head are handled by `GraphNetWrapper`, not here.
+    """
+
+    NAME = "GINEGraphPred"
+
+    def _get_gnn(self, config: dict[str, Any]) -> GINEConvNetBase:
+        """
+        Returns the core GINE GNN model.
+
+        Args:
+            config (dict): Configuration dictionary for the GNN model.
+
+        Returns:
+            GINEConvNetBase: The core graph convolutional network.
+        """
+        return GINEConvNetBase(config=config)