pyg-nightly 2.7.0.dev20241009__py3-none-any.whl → 2.8.0.dev20251228__py3-none-any.whl

torch_geometric/explain/algorithm/pg_explainer.py

@@ -1,21 +1,26 @@
 import logging
-from typing import Optional, Union
+from typing import Dict, Optional, Tuple, Union, overload
 
 import torch
 from torch import Tensor
 from torch.nn import ReLU, Sequential
 
-from torch_geometric.explain import Explanation
+from torch_geometric.explain import Explanation, HeteroExplanation
 from torch_geometric.explain.algorithm import ExplainerAlgorithm
-from torch_geometric.explain.algorithm.utils import clear_masks, set_masks
+from torch_geometric.explain.algorithm.utils import (
+    clear_masks,
+    set_hetero_masks,
+    set_masks,
+)
 from torch_geometric.explain.config import (
     ExplanationType,
     ModelMode,
     ModelTaskLevel,
 )
-from torch_geometric.nn import Linear
+from torch_geometric.nn import HANConv, HeteroConv, HGTConv, Linear
 from torch_geometric.nn.inits import reset
-from torch_geometric.utils import get_embeddings
+from torch_geometric.typing import EdgeType, NodeType
+from torch_geometric.utils import get_embeddings, get_embeddings_hetero
 
 
 class PGExplainer(ExplainerAlgorithm):
@@ -62,6 +67,13 @@ class PGExplainer(ExplainerAlgorithm):
         'bias': 0.01,
     }
 
+    # NOTE: Add more in the future as needed.
+    SUPPORTED_HETERO_MODELS = [
+        HGTConv,
+        HANConv,
+        HeteroConv,
+    ]
+
     def __init__(self, epochs: int, lr: float = 0.003, **kwargs):
         super().__init__()
         self.epochs = epochs
@@ -75,11 +87,13 @@ class PGExplainer(ExplainerAlgorithm):
         )
         self.optimizer = torch.optim.Adam(self.mlp.parameters(), lr=lr)
         self._curr_epoch = -1
+        self.is_hetero = False
 
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
         reset(self.mlp)
 
+    @overload
     def train(
         self,
         epoch: int,
@@ -90,17 +104,44 @@ class PGExplainer(ExplainerAlgorithm):
         target: Tensor,
         index: Optional[Union[int, Tensor]] = None,
         **kwargs,
-    ):
+    ) -> float:
+        ...
+
+    @overload
+    def train(
+        self,
+        epoch: int,
+        model: torch.nn.Module,
+        x: Dict[NodeType, Tensor],
+        edge_index: Dict[EdgeType, Tensor],
+        *,
+        target: Tensor,
+        index: Optional[Union[int, Tensor]] = None,
+        **kwargs,
+    ) -> float:
+        ...
+
+    def train(
+        self,
+        epoch: int,
+        model: torch.nn.Module,
+        x: Union[Tensor, Dict[NodeType, Tensor]],
+        edge_index: Union[Tensor, Dict[EdgeType, Tensor]],
+        *,
+        target: Tensor,
+        index: Optional[Union[int, Tensor]] = None,
+        **kwargs,
+    ) -> float:
         r"""Trains the underlying explainer model.
         Needs to be called before being able to make predictions.
 
         Args:
             epoch (int): The current epoch of the training phase.
             model (torch.nn.Module): The model to explain.
-            x (torch.Tensor): The input node features of a
-                homogeneous graph.
-            edge_index (torch.Tensor): The input edge indices of a homogeneous
-                graph.
+            x (torch.Tensor or Dict[str, torch.Tensor]): The input node
+                features. Can be either homogeneous or heterogeneous.
+            edge_index (torch.Tensor or Dict[Tuple[str, str, str]): The input
+                edge indices. Can be either homogeneous or heterogeneous.
             target (torch.Tensor): The target of the model.
             index (int or torch.Tensor, optional): The index of the model
                 output to explain. Needs to be a single index.
@@ -108,9 +149,9 @@ class PGExplainer(ExplainerAlgorithm):
             **kwargs (optional): Additional keyword arguments passed to
                 :obj:`model`.
         """
-        if isinstance(x, dict) or isinstance(edge_index, dict):
-            raise ValueError(f"Heterogeneous graphs not yet supported in "
-                             f"'{self.__class__.__name__}'")
+        self.is_hetero = isinstance(x, dict)
+        if self.is_hetero:
+            assert isinstance(edge_index, dict)
 
         if self.model_config.task_level == ModelTaskLevel.node:
             if index is None:
@@ -121,35 +162,68 @@ class PGExplainer(ExplainerAlgorithm):
                 raise ValueError(f"Only scalars are supported for the 'index' "
                                  f"argument in '{self.__class__.__name__}'")
 
-        z = get_embeddings(model, x, edge_index, **kwargs)[-1]
+        # Get embeddings based on whether the graph is homogeneous or
+        # heterogeneous
+        node_embeddings = self._get_embeddings(model, x, edge_index, **kwargs)
 
+        # Train the model
         self.optimizer.zero_grad()
         temperature = self._get_temperature(epoch)
 
-        inputs = self._get_inputs(z, edge_index, index)
-        logits = self.mlp(inputs).view(-1)
-        edge_mask = self._concrete_sample(logits, temperature)
-        set_masks(model, edge_mask, edge_index, apply_sigmoid=True)
-
-        if self.model_config.task_level == ModelTaskLevel.node:
-            _, hard_edge_mask = self._get_hard_masks(model, index, edge_index,
-                                                     num_nodes=x.size(0))
-            edge_mask = edge_mask[hard_edge_mask]
-
+        # Process embeddings and generate edge masks
+        edge_mask = self._generate_edge_masks(node_embeddings, edge_index,
+                                              index, temperature)
+
+        # Apply masks to the model
+        if self.is_hetero:
+            set_hetero_masks(model, edge_mask, edge_index, apply_sigmoid=True)
+
+            # For node-level tasks, we can compute hard masks
+            if self.model_config.task_level == ModelTaskLevel.node:
+                # Process each edge type separately
+                for edge_type, mask in edge_mask.items():
+                    # Get the edge indices for this edge type
+                    edges = edge_index[edge_type]
+                    src_type, _, dst_type = edge_type
+
+                    # Get hard masks for this specific edge type
+                    _, hard_mask = self._get_hard_masks(
+                        model, index, edges,
+                        num_nodes=max(x[src_type].size(0),
+                                      x[dst_type].size(0)))
+
+                    edge_mask[edge_type] = mask[hard_mask]
+        else:
+            # Apply masks for homogeneous graphs
+            set_masks(model, edge_mask, edge_index, apply_sigmoid=True)
+
+            # For node-level tasks, we may need to apply hard masks
+            hard_edge_mask = None
+            if self.model_config.task_level == ModelTaskLevel.node:
+                _, hard_edge_mask = self._get_hard_masks(
+                    model, index, edge_index, num_nodes=x.size(0))
+                edge_mask = edge_mask[hard_edge_mask]
+
+        # Forward pass with masks applied
         y_hat, y = model(x, edge_index, **kwargs), target
 
         if index is not None:
             y_hat, y = y_hat[index], y[index]
 
+        # Calculate loss
         loss = self._loss(y_hat, y, edge_mask)
+
+        # Backward pass and optimization
         loss.backward()
         self.optimizer.step()
 
+        # Clean up
         clear_masks(model)
         self._curr_epoch = epoch
 
         return float(loss)
 
+    @overload
     def forward(
         self,
         model: torch.nn.Module,
@@ -160,9 +234,32 @@ class PGExplainer(ExplainerAlgorithm):
         index: Optional[Union[int, Tensor]] = None,
         **kwargs,
     ) -> Explanation:
-        if isinstance(x, dict) or isinstance(edge_index, dict):
-            raise ValueError(f"Heterogeneous graphs not yet supported in "
-                             f"'{self.__class__.__name__}'")
+        ...
+
+    @overload
+    def forward(
+        self,
+        model: torch.nn.Module,
+        x: Dict[NodeType, Tensor],
+        edge_index: Dict[EdgeType, Tensor],
+        *,
+        target: Tensor,
+        index: Optional[Union[int, Tensor]] = None,
+        **kwargs,
+    ) -> HeteroExplanation:
+        ...
+
+    def forward(
+        self,
+        model: torch.nn.Module,
+        x: Union[Tensor, Dict[NodeType, Tensor]],
+        edge_index: Union[Tensor, Dict[EdgeType, Tensor]],
+        *,
+        target: Tensor,
+        index: Optional[Union[int, Tensor]] = None,
+        **kwargs,
+    ) -> Union[Explanation, HeteroExplanation]:
+        self.is_hetero = isinstance(x, dict)
 
         if self._curr_epoch < self.epochs - 1:  # Safety check:
             raise ValueError(f"'{self.__class__.__name__}' is not yet fully "
@@ -171,7 +268,6 @@ class PGExplainer(ExplainerAlgorithm):
                              f"the underlying explainer model by running "
                              f"`explainer.algorithm.train(...)`.")
 
-        hard_edge_mask = None
         if self.model_config.task_level == ModelTaskLevel.node:
             if index is None:
                 raise ValueError(f"The 'index' argument needs to be provided "
@@ -181,20 +277,55 @@ class PGExplainer(ExplainerAlgorithm):
                 raise ValueError(f"Only scalars are supported for the 'index' "
                                  f"argument in '{self.__class__.__name__}'")
 
-            # We need to compute hard masks to properly clean up edges and
-            # nodes attributions not involved during message passing:
-            _, hard_edge_mask = self._get_hard_masks(model, index, edge_index,
-                                                     num_nodes=x.size(0))
+        # Get embeddings
+        node_embeddings = self._get_embeddings(model, x, edge_index, **kwargs)
+
+        # Generate explanations
+        if self.is_hetero:
+            # Generate edge masks for each edge type
+            edge_masks = {}
+
+            # Generate masks for each edge type
+            for edge_type, edge_idx in edge_index.items():
+                src_node_type, _, dst_node_type = edge_type
+
+                assert src_node_type in node_embeddings
+                assert dst_node_type in node_embeddings
+
+                inputs = self._get_inputs_hetero(node_embeddings, edge_type,
+                                                 edge_idx, index)
+                logits = self.mlp(inputs).view(-1)
+
+                # For node-level explanations, get hard masks for this
+                # specific edge type
+                hard_edge_mask = None
+                if self.model_config.task_level == ModelTaskLevel.node:
+                    _, hard_edge_mask = self._get_hard_masks(
+                        model, index, edge_idx,
+                        num_nodes=max(x[src_node_type].size(0),
+                                      x[dst_node_type].size(0)))
 
-        z = get_embeddings(model, x, edge_index, **kwargs)[-1]
+                # Apply hard mask if available and it has any True values
+                edge_masks[edge_type] = self._post_process_mask(
+                    logits, hard_edge_mask, apply_sigmoid=True)
 
-        inputs = self._get_inputs(z, edge_index, index)
-        logits = self.mlp(inputs).view(-1)
+            explanation = HeteroExplanation()
+            explanation.set_value_dict('edge_mask', edge_masks)
+            return explanation
+        else:
+            hard_edge_mask = None
+            if self.model_config.task_level == ModelTaskLevel.node:
+                # We need to compute hard masks to properly clean up edges
+                _, hard_edge_mask = self._get_hard_masks(
+                    model, index, edge_index, num_nodes=x.size(0))
 
-        edge_mask = self._post_process_mask(logits, hard_edge_mask,
-                                            apply_sigmoid=True)
+            inputs = self._get_inputs(node_embeddings, edge_index, index)
+            logits = self.mlp(inputs).view(-1)
 
-        return Explanation(edge_mask=edge_mask)
+            edge_mask = self._post_process_mask(logits, hard_edge_mask,
+                                                apply_sigmoid=True)
+
+            return Explanation(edge_mask=edge_mask)
 
     def supports(self) -> bool:
         explanation_type = self.explainer_config.explanation_type
@@ -222,6 +353,76 @@ class PGExplainer(ExplainerAlgorithm):
 
     ###########################################################################
 
+    def _get_embeddings(self, model: torch.nn.Module, x: Union[Tensor,
+                                                               Dict[NodeType,
+                                                                    Tensor]],
+                        edge_index: Union[Tensor, Dict[EdgeType, Tensor]],
+                        **kwargs) -> Union[Tensor, Dict[NodeType, Tensor]]:
+        """Get embeddings from the model based on input type."""
+        if self.is_hetero:
+            # For heterogeneous graphs, get embeddings for each node type
+            embeddings_dict = get_embeddings_hetero(
+                model,
+                self.SUPPORTED_HETERO_MODELS,
+                x,
+                edge_index,
+                **kwargs,
+            )
+
+            # Use the last layer's embeddings for each node type
+            last_embedding_dict = {
+                node_type: embs[-1] if embs and len(embs) > 0 else None
+                for node_type, embs in embeddings_dict.items()
+            }
+
+            # Skip if no embeddings were captured
+            if not any(emb is not None
+                       for emb in last_embedding_dict.values()):
+                raise ValueError(
+                    "No embeddings were captured from the model. "
+                    "Please check if the model architecture is supported.")
+
+            return last_embedding_dict
+        else:
+            # For homogeneous graphs, get embeddings directly
+            return get_embeddings(model, x, edge_index, **kwargs)[-1]
+
+    def _generate_edge_masks(
+            self, emb: Union[Tensor, Dict[NodeType, Tensor]],
+            edge_index: Union[Tensor,
+                              Dict[EdgeType,
+                                   Tensor]], index: Optional[Union[int,
+                                                                   Tensor]],
+            temperature: float) -> Union[Tensor, Dict[EdgeType, Tensor]]:
+        """Generate edge masks based on embeddings."""
+        if self.is_hetero:
+            # For heterogeneous graphs, generate masks for each edge type
+            edge_masks = {}
+
+            for edge_type, edge_idx in edge_index.items():
+                src, _, dst = edge_type
+
+                assert src in emb and dst in emb
+                # Generate inputs for this edge type
+                inputs = self._get_inputs_hetero(emb, edge_type, edge_idx,
+                                                 index)
+                logits = self.mlp(inputs).view(-1)
+                edge_masks[edge_type] = self._concrete_sample(
+                    logits, temperature)
+
+            # Ensure we have at least one valid edge mask
+            if not edge_masks:
+                raise ValueError(
+                    "Could not generate edge masks for any edge type. "
+                    "Please ensure the model architecture is supported.")
+
+            return edge_masks
+        else:
+            # For homogeneous graphs, generate a single mask
+            inputs = self._get_inputs(emb, edge_index, index)
+            logits = self.mlp(inputs).view(-1)
+            return self._concrete_sample(logits, temperature)
+
     def _get_inputs(self, embedding: Tensor, edge_index: Tensor,
                     index: Optional[int] = None) -> Tensor:
         zs = [embedding[edge_index[0]], embedding[edge_index[1]]]
@@ -230,6 +431,27 @@ class PGExplainer(ExplainerAlgorithm):
             zs.append(embedding[index].view(1, -1).repeat(zs[0].size(0), 1))
         return torch.cat(zs, dim=-1)
 
+    def _get_inputs_hetero(self, embedding_dict: Dict[NodeType, Tensor],
+                           edge_type: Tuple[str, str, str], edge_index: Tensor,
+                           index: Optional[int] = None) -> Tensor:
+        src, _, dst = edge_type
+
+        # Get embeddings for source and destination nodes
+        src_emb = embedding_dict[src]
+        dst_emb = embedding_dict[dst]
+
+        # Source and destination node embeddings
+        zs = [src_emb[edge_index[0]], dst_emb[edge_index[1]]]
+
+        # For node-level explanations, add the target node embedding
+        if self.model_config.task_level == ModelTaskLevel.node:
+            assert index is not None
+            # Assuming index refers to a node of type 'src'
+            target_emb = src_emb[index].view(1, -1).repeat(zs[0].size(0), 1)
+            zs.append(target_emb)
+
+        return torch.cat(zs, dim=-1)
+
     def _get_temperature(self, epoch: int) -> float:
         temp = self.coeffs['temp']
         return temp[0] * pow(temp[1] / temp[0], epoch / self.epochs)
@@ -240,19 +462,55 @@ class PGExplainer(ExplainerAlgorithm):
         eps = (1 - 2 * bias) * torch.rand_like(logits) + bias
         return (eps.log() - (1 - eps).log() + logits) / temperature
 
-    def _loss(self, y_hat: Tensor, y: Tensor, edge_mask: Tensor) -> Tensor:
+    def _loss(self, y_hat: Tensor, y: Tensor,
+              edge_mask: Union[Tensor, Dict[EdgeType, Tensor]]) -> Tensor:
+        # Calculate base loss based on model configuration
+        loss = self._calculate_base_loss(y_hat, y)
+
+        # Apply regularization based on graph type
+        if self.is_hetero:
+            loss = self._apply_hetero_regularization(loss, edge_mask)
+        else:
+            loss = self._apply_homo_regularization(loss, edge_mask)
+
+        return loss
+
+    def _calculate_base_loss(self, y_hat: Tensor, y: Tensor) -> Tensor:
+        """Calculate base loss based on model configuration."""
         if self.model_config.mode == ModelMode.binary_classification:
-            loss = self._loss_binary_classification(y_hat, y)
+            return self._loss_binary_classification(y_hat, y)
         elif self.model_config.mode == ModelMode.multiclass_classification:
-            loss = self._loss_multiclass_classification(y_hat, y)
+            return self._loss_multiclass_classification(y_hat, y)
         elif self.model_config.mode == ModelMode.regression:
-            loss = self._loss_regression(y_hat, y)
-
-        # Regularization loss:
-        mask = edge_mask.sigmoid()
+            return self._loss_regression(y_hat, y)
+        else:
+            raise ValueError(
+                f"Unsupported model mode: {self.model_config.mode}")
+
+    def _apply_hetero_regularization(
+            self, loss: Tensor, edge_mask: Dict[EdgeType, Tensor]) -> Tensor:
+        """Apply regularization for heterogeneous graph."""
+        for _, mask in edge_mask.items():
+            loss = self._add_mask_regularization(loss, mask)
+
+        return loss
+
+    def _apply_homo_regularization(self, loss: Tensor,
+                                   edge_mask: Tensor) -> Tensor:
+        """Apply regularization for homogeneous graph."""
+        return self._add_mask_regularization(loss, edge_mask)
+
+    def _add_mask_regularization(self, loss: Tensor, mask: Tensor) -> Tensor:
+        """Add size and entropy regularization for a mask."""
+        # Apply sigmoid for mask values
+        mask = mask.sigmoid()
+
+        # Size regularization
         size_loss = mask.sum() * self.coeffs['edge_size']
-        mask = 0.99 * mask + 0.005
-        mask_ent = -mask * mask.log() - (1 - mask) * (1 - mask).log()
+
+        # Entropy regularization
+        masked = 0.99 * mask + 0.005
+        mask_ent = -masked * masked.log() - (1 - masked) * (1 - masked).log()
         mask_ent_loss = mask_ent.mean() * self.coeffs['edge_ent']
 
         return loss + size_loss + mask_ent_loss

torch_geometric/explain/explainer.py

@@ -192,7 +192,7 @@ class Explainer:
             if target is not None:
                 warnings.warn(
                     f"The 'target' should not be provided for the explanation "
-                    f"type '{self.explanation_type.value}'")
+                    f"type '{self.explanation_type.value}'", stacklevel=2)
             prediction = self.get_prediction(x, edge_index, **kwargs)
             target = self.get_target(prediction)
 
@@ -265,7 +265,7 @@ class Explainer:
                 return (prediction > 0).long().view(-1)
             if self.model_config.return_type == ModelReturnType.probs:
                 return (prediction > 0.5).long().view(-1)
-            assert False
+            raise AssertionError()
 
         if self.model_config.mode == ModelMode.multiclass_classification:
             return prediction.argmax(dim=-1)

torch_geometric/explain/explanation.py

@@ -1,5 +1,5 @@
 import copy
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -8,7 +8,10 @@ from torch_geometric.data.data import Data, warn_or_raise
 from torch_geometric.data.hetero_data import HeteroData
 from torch_geometric.explain.config import ThresholdConfig, ThresholdType
 from torch_geometric.typing import EdgeType, NodeType
-from torch_geometric.visualization import visualize_graph
+from torch_geometric.visualization import (
+    visualize_graph,
+    visualize_hetero_graph,
+)
 
 
 class ExplanationMixin:
@@ -100,7 +103,7 @@ class ExplanationMixin:
                 out[index] = 1.0
             return out.view(mask.size())
 
-        assert False
+        raise AssertionError()
 
     def threshold(
         self,
@@ -362,6 +365,87 @@ class HeteroExplanation(HeteroData, ExplanationMixin):
 
         return _visualize_score(score, all_feat_labels, path, top_k)
 
+    def visualize_graph(
+            self,
+            path: Optional[str] = None,
+            node_labels: Optional[Dict[NodeType, List[str]]] = None,
+            node_size_range: Tuple[float, float] = (50, 500),
+            node_opacity_range: Tuple[float, float] = (0.2, 1.0),
+            edge_width_range: Tuple[float, float] = (0.1, 2.0),
+            edge_opacity_range: Tuple[float, float] = (0.2, 1.0),
+    ) -> None:
+        r"""Visualizes the explanation subgraph using networkx, with edge
+        opacity corresponding to edge importance and node colors
+        corresponding to node types.
+
+        Args:
+            path (str, optional): The path to where the plot is saved.
+                If set to :obj:`None`, will visualize the plot on-the-fly.
+                (default: :obj:`None`)
+            node_labels (Dict[NodeType, List[str]], optional): The display
+                names of nodes for each node type that will be shown in the
+                visualization. (default: :obj:`None`)
+            node_size_range (Tuple[float, float], optional): The minimum and
+                maximum node size in the visualization.
+                (default: :obj:`(50, 500)`)
+            node_opacity_range (Tuple[float, float], optional): The minimum and
+                maximum node opacity in the visualization.
+                (default: :obj:`(0.2, 1.0)`)
+            edge_width_range (Tuple[float, float], optional): The minimum and
+                maximum edge width in the visualization.
+                (default: :obj:`(0.1, 2.0)`)
+            edge_opacity_range (Tuple[float, float], optional): The minimum and
+                maximum edge opacity in the visualization.
+                (default: :obj:`(0.2, 1.0)`)
+        """
+        # Validate node labels if provided
+        if node_labels is not None:
+            for node_type, labels in node_labels.items():
+                if node_type not in self.node_types:
+                    raise ValueError(
+                        f"Node type '{node_type}' in node_labels "
+                        f"does not exist in the explanation graph")
+                if len(labels) != self[node_type].num_nodes:
+                    raise ValueError(f"Number of labels for node type "
+                                     f"'{node_type}' (got {len(labels)}) does "
+                                     f"not match the number of nodes "
+                                     f"(got {self[node_type].num_nodes})")
+        # Get the explanation subgraph
+        subgraph = self.get_explanation_subgraph()
+
+        # Prepare edge indices and weights for each edge type
+        edge_index_dict = {}
+        edge_weight_dict = {}
+        for edge_type in subgraph.edge_types:
+            if edge_type[0] == 'x' or edge_type[-1] == 'x':  # Skip edges
+                continue
+            edge_index_dict[edge_type] = subgraph[edge_type].edge_index
+            edge_weight_dict[edge_type] = subgraph[edge_type].get(
+                'edge_mask',
+                torch.ones(subgraph[edge_type].edge_index.size(1)))
+
+        # Prepare node weights for each node type
+        node_weight_dict = {}
+        for node_type in subgraph.node_types:
+            if node_type == 'x':  # Skip the global store
+                continue
+            node_weight_dict[node_type] = subgraph[node_type] \
+                .get('node_mask',
+                     torch.ones(subgraph[node_type].num_nodes)).squeeze(-1)
+
+        # Call the visualization function
+        visualize_hetero_graph(
+            edge_index_dict=edge_index_dict,
+            edge_weight_dict=edge_weight_dict,
+            path=path,
+            node_labels_dict=node_labels,
+            node_weight_dict=node_weight_dict,
+            node_size_range=node_size_range,
+            node_opacity_range=node_opacity_range,
+            edge_width_range=edge_width_range,
+            edge_opacity_range=edge_opacity_range,
+        )
+
 
 def _visualize_score(
     score: torch.Tensor,

torch_geometric/explain/metric/faithfulness.py

@@ -13,7 +13,7 @@ def unfaithfulness(
     top_k: Optional[int] = None,
 ) -> float:
     r"""Evaluates how faithful an :class:`~torch_geometric.explain.Explanation`
-    is to an underyling GNN predictor, as described in the
+    is to an underlying GNN predictor, as described in the
     `"Evaluating Explainability for Graph Neural Networks"
     <https://arxiv.org/abs/2208.09339>`_ paper.
 

torch_geometric/graphgym/config.py

@@ -16,8 +16,9 @@ try:  # Define global config object
     cfg = CN()
 except ImportError:
     cfg = None
-    warnings.warn("Could not define global config object. Please install "
-                  "'yacs' via 'pip install yacs' in order to use GraphGym")
+    warnings.warn(
+        "Could not define global config object. Please install "
+        "'yacs' via 'pip install yacs' in order to use GraphGym", stacklevel=2)
 
 
 def set_cfg(cfg):

torch_geometric/graphgym/imports.py

@@ -3,13 +3,24 @@ import warnings
 import torch
 
 try:
-    import pytorch_lightning as pl
+    import lightning.pytorch as pl
+    _pl_is_available = True
+except ImportError:
+    try:
+        import pytorch_lightning as pl
+        _pl_is_available = True
+    except ImportError:
+        _pl_is_available = False
+
+if _pl_is_available:
     LightningModule = pl.LightningModule
     Callback = pl.Callback
-except ImportError:
+else:
     pl = object
     LightningModule = torch.nn.Module
     Callback = object
 
-    warnings.warn("Please install 'pytorch_lightning' via  "
-                  "'pip install pytorch_lightning' in order to use GraphGym")
+    warnings.warn(
+        "To use GraphGym, install 'pytorch_lightning' or 'lightning' via "
+        "'pip install pytorch_lightning' or 'pip install lightning'",
+        stacklevel=2)

torch_geometric/graphgym/logger.py

@@ -239,7 +239,7 @@ def create_logger():
     r"""Create logger for the experiment."""
     loggers = []
     names = ['train', 'val', 'test']
-    for i, dataset in enumerate(range(cfg.share.num_splits)):
+    for i, _ in enumerate(range(cfg.share.num_splits)):
         loggers.append(Logger(name=names[i], task_type=infer_task()))
     return loggers