Added XGNN implementation to PyG

.gitignore

@@ -28,6 +28,5 @@ examples/**/*.png
 examples/**/*.pdf
 benchmark/results/
 .mypy_cache/
-
 !torch_geometric/data/
 !test/data/

CHANGELOG.md

@@ -38,6 +38,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added the `RCDD` dataset ([#8196](https://github.com/pyg-team/pytorch_geometric/pull/8196))
 - Added distributed `GAT + ogbn-products` example targeting XPU device ([#8032](https://github.com/pyg-team/pytorch_geometric/pull/8032))
 - Added the option to skip explanations of certain message passing layers via `conv.explain = False` ([#8216](https://github.com/pyg-team/pytorch_geometric/pull/8216))
+- Added XGNN implementation for graph explanation to `explain` module
 
 ### Changed
 

docs/source/modules/explain.rst

@@ -55,6 +55,10 @@ Explanations
    :show-inheritance:
    :members:
 
+.. autoclass:: torch_geometric.explain.GenerativeExplanation
+   :show-inheritance:
+   :members:
+
 Explainer Algorithms
 --------------------
 

examples/explain/xgnn/xgnn_model.py

@@ -0,0 +1,65 @@
+from torch_geometric.nn import GCNConv
+import torch
+import torch.nn.functional as F
+from torch_geometric.nn import global_mean_pool
+from torch.nn.parameter import Parameter
+import math 
+
+### GCN to predict graph property
+class GCN_Graph(torch.nn.Module):
+    def __init__(self, input_dim, output_dim, dropout, emb = False):
+        super(GCN_Graph, self).__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+
+        self.dropout = dropout
+        self.convs = torch.nn.ModuleList([GCNConv(in_channels = input_dim, out_channels = 32),
+                                          GCNConv(in_channels = 32,        out_channels = 48),
+                                          GCNConv(in_channels = 48,        out_channels = 64)])
+
+        self.pool = global_mean_pool # global averaging to obtain graph representation
+
+        self.fc1 = torch.nn.Linear(64, 32)
+        self.fc2 = torch.nn.Linear(32, output_dim)
+
+        self.loss = torch.nn.BCEWithLogitsLoss()
+        self.reset_parameters()
+
+    def reset_parameters(self):
+      for conv in self.convs:
+          conv.reset_parameters()
+          stdv = 1. / math.sqrt(conv.lin.weight.size(1))
+          torch.nn.init.uniform_(conv.lin.weight, -stdv, stdv)
+
+          conv.bias = Parameter(torch.FloatTensor(conv.out_channels))
+          conv.bias.data.uniform_(-stdv, stdv)
+
+      self.fc1.reset_parameters()
+      self.fc2.reset_parameters()
+
+    def forward(self, data):
+        # Extract important attributes of our mini-batch
+        x, edge_index = data.x, data.edge_index
+
+        for i in range(len(self.convs)):
+            x = F.relu(self.convs[i](x, edge_index))
+            if i < len(self.convs) - 1: # do not apply dropout on last layer
+                x = F.dropout(x, p=self.dropout, training=self.training)
+
+        # Check if 'batch' attribute is present
+        if hasattr(data, 'batch'):
+            batch = data.batch
+        else:
+            # For a single graph, use a zero tensor as the batch vector, 
+            # where its size equals the number of nodes.
+            batch = torch.zeros(data.num_nodes, dtype=torch.long, device=x.device)
+
+        x = self.pool(x, batch)
+
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, self.dropout, training=self.training)
+        x = self.fc2(x)
+        #x = F.sigmoid(x)
+        #x = F.softmax(x, dim=1)
+        return x

examples/explain/xgnn/xgnn_train.py

@@ -0,0 +1,125 @@
+from torch_geometric.data import Batch
+from torch_geometric.datasets import TUDataset
+import torch
+import torch.optim as optim
+import numpy as np
+from tqdm import trange
+import copy
+from tqdm.auto import trange
+import matplotlib.pyplot as plt
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from xgnn_model import GCN_Graph
+
+seed = 42
+np.random.seed(seed)
+torch.manual_seed(seed)
+
+def create_single_batch(dataset):
+    data_list = [data for data in dataset]
+    batched_data = Batch.from_data_list(data_list)
+    return batched_data
+
+def test(test_dataset, model):
+    model.eval()
+    with torch.no_grad():
+        logits = model(test_dataset).squeeze() # Logits for each graph
+        probabilities = torch.sigmoid(logits) # Convert logits to probabilities
+        predictions = probabilities > 0.5 # Convert probabilities to binary predictions
+        correct = (predictions == test_dataset.y).float() # Assumes labels are 0 or 1
+        accuracy = correct.mean()
+
+    return accuracy
+
+
+def train(dataset, args, train_indices, val_indices, test_indices):
+    # Split dataset into training and testing (validation is not used here)
+    train_dataset = create_single_batch([dataset[i] for i in train_indices]).to(device)
+    test_dataset = create_single_batch([dataset[i] for i in test_indices]).to(device)
+
+    # Model initialization
+    model = GCN_Graph(args.input_dim, output_dim=1, dropout=args.dropout).to(device)
+    opt = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # 
+
+    # Training loop
+    losses = []
+    test_accs = []
+    best_acc = 0
+    best_model = None
+    for epoch in trange(args.epochs, desc="Training", unit="Epoch"):
+        model.train()
+        opt.zero_grad()
+
+        pred = model(train_dataset)
+        label = train_dataset.y.float()
+        loss = model.loss(pred.squeeze(), label)
+        loss.backward()
+        opt.step()
+        total_loss = loss.item()
+        losses.append(total_loss)
+
+        # Test accuracy
+        if epoch % 10 == 0:
+            test_acc = test(test_dataset, model)
+
+            test_accs.append(test_acc)
+            if test_acc > best_acc:
+                best_acc = test_acc
+                best_model = copy.deepcopy(model)
+        else:
+            test_accs.append(test_accs[-1])
+
+
+
+    return test_accs, losses, best_model, best_acc
+
+class objectview(object):
+    def __init__(self, d):
+        self.__dict__ = d
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+args = {'device': device,
+        'dropout': 0.1,
+        'epochs': 5000,
+        'input_dim' : 7,
+        'opt': 'adam',
+        'opt_restart': 0,
+        'weight_decay': 1e-4,
+        'lr': 0.007}
+
+args = objectview(args)
+
+dataset = TUDataset(root='/tmp/MUTAG', name='MUTAG')
+num_graphs = len(dataset)
+
+# Define split percentages
+train_percentage = 0.7
+val_percentage = 0.0
+
+# Calculate split sizes
+train_size = int(num_graphs * train_percentage)
+val_size = int(num_graphs * val_percentage)
+test_size = num_graphs - train_size - val_size
+
+# Create shuffled indices
+indices = np.random.permutation(num_graphs)
+train_indices = indices[:train_size]
+val_indices = indices[train_size:train_size + val_size]
+test_indices = indices[train_size + val_size:]
+
+test_accs, losses, best_model, best_acc = train(dataset, args, train_indices, val_indices, test_indices)
+
+try:
+    torch.save(best_model.state_dict(), 'examples/explain/xgnn/mutag_model.pth')
+    print("Model saved successfully.")
+except Exception as e:
+    print("Error saving model:", e)
+
+print("Maximum test set accuracy: {0}".format(max(test_accs)))
+print("Minimum loss: {0}".format(min(losses)))
+
+plt.title(dataset.name)
+plt.plot(losses, label="training loss")
+plt.plot(test_accs, label="test accuracy")
+plt.legend()
+plt.show()

test/explain/algorithm/test_xgnn_explainer.py

@@ -0,0 +1,40 @@
+import pytest
+import torch
+from torch_geometric.explain import XGNNExplainer, GenerativeExplanation
+from abc import abstractmethod
+
+# Mock subclass of XGNNExplainer for testing
+class MockXGNNExplainer(XGNNExplainer):
+    def train_generative_model(self, model, for_class, **kwargs):
+        return None
+
+@pytest.fixture
+def model():
+    return torch.nn.Linear(3, 2)
+
+def test_xgnn_explainer_initialization():
+    explainer = MockXGNNExplainer(epochs=200, lr=0.005)
+    assert explainer.epochs == 200
+    assert explainer.lr == 0.005
+
+def test_xgnn_explainer_forward(model):
+    explainer = MockXGNNExplainer()
+    x = torch.rand(10, 3)
+    edge_index = torch.randint(0, 10, (2, 30))
+    target = torch.randint(0, 2, (10,))
+
+    explanation = explainer(model, x, edge_index, target=target, for_class=1)
+    assert isinstance(explanation, GenerativeExplanation)
+
+    # Test ValueError for missing 'for_class' argument
+    with pytest.raises(ValueError):
+        explainer(model, x, edge_index, target=target)
+
+def test_xgnn_explainer_abstract_method():
+    class IncompleteExplainer(XGNNExplainer):
+        pass
+    explainer = IncompleteExplainer()
+
+    # Ensure that instantiation fails due to the unimplemented abstract method
+    with pytest.raises(NotImplementedError):
+        explainer.train_generative_model(None, for_class=0)

test/explain/test_explainer.py

@@ -60,7 +60,8 @@ def test_forward(data, target, explanation_type):
         assert isinstance(explanation, Explanation)
         assert 'x' in explanation
         assert 'edge_index' in explanation
-        assert 'target' in explanation
+        if explanation_type != ExplanationType.generative:  # target is not used for generative explanation
+            assert 'target' in explanation
         assert 'node_mask' in explanation.available_explanations
         assert explanation.node_mask.size() == data.x.size()
 

test/explain/test_generative_explanation.py

@@ -0,0 +1,96 @@
+import pytest
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.data import Data
+from torch_geometric.explain import Explainer, XGNNExplainer
+
+
+# Mock model for testing
+class MLP_Graph(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super(MLP_Graph, self).__init__()
+        self.fc1 = nn.Linear(input_dim, 8)
+        self.fc2 = nn.Linear(8, output_dim)
+
+    def forward(self, x):
+        # Flatten the graph representation
+        x = x.view(x.size(0), -1)
+        x = F.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+
+# Mock explainer algorithm
+class ExampleExplainer(XGNNExplainer):
+    def __init__(self, epochs, lr, candidate_set, validity_args):
+        super(ExampleExplainer, self).__init__()
+        self.epochs = epochs
+        self.lr = lr
+        self.candidate_set = candidate_set
+        self.validity_args = validity_args
+
+    def train_generative_model(self, model_to_explain, for_class):
+        # For simplicity, this example does not include actual training logic
+
+        for epoch in range(self.epochs):
+            # Placeholder for training logic
+            pass
+
+        return Data()
+
+
+# Mock graph generator
+class ExampleGraphGenerator():
+    def __init__(self, graph):
+        self.graph = graph
+
+    def sample(self):
+        # has to return a list of Data objects
+        return [Data(), Data(), Data()]
+
+
+# Fixture for setting up XGNNExplainer
+@pytest.fixture
+def setup_xgnn_explainer():
+    mock_model = MLP_Graph(input_dim=7, output_dim=1)
+
+    explainer = Explainer(
+        model = mock_model,
+        algorithm = ExampleExplainer(epochs = 10, 
+                                lr = 0.01,
+                                candidate_set={'C': torch.tensor([1, 0, 0, 0, 0, 0, 0])},  # Simplified candidate set
+                                validity_args={'C': 4}),
+        explanation_type = 'generative',
+        model_config = dict(
+            mode = 'binary_classification',
+            task_level = 'graph',
+            return_type = 'probs',
+        )
+    )
+
+    class_index = 1
+    x = torch.tensor([])
+    edge_index = torch.tensor([[], []])
+
+    return explainer, x, edge_index, class_index
+
+
+# Test output of XGNNExplainer
+def test_explainer_output(setup_xgnn_explainer):
+    explainer, x, edge_index, class_index = setup_xgnn_explainer
+    explanation = explainer(x, edge_index, for_class=class_index)
+
+    # Check if explanation is of type Data
+    assert isinstance(explanation, Data), "Explanation is not of type Data"
+
+
+# Test output of ExampleExplainer
+def test_sampler_output():
+    sampled_graphs = ExampleGraphGenerator(Data()).sample()
+
+    # Check if sampled_graphs is a list of Data objects
+    assert isinstance(sampled_graphs, list), "Sampled graphs is not a list"
+    assert all(isinstance(graph, Data) for graph in sampled_graphs), "Sampled graphs is not a list of Data objects"
+
+

torch_geometric/explain/__init__.py

@@ -1,5 +1,5 @@
 from .config import ExplainerConfig, ModelConfig, ThresholdConfig
-from .explanation import Explanation, HeteroExplanation
+from .explanation import Explanation, HeteroExplanation, GenerativeExplanation, ExplanationSetSampler
 from .algorithm import *  # noqa
 from .explainer import Explainer
 from .metric import *  # noqa
@@ -11,4 +11,6 @@
     'Explanation',
     'HeteroExplanation',
     'Explainer',
+    'GenerativeExplanation',
+    'ExplanationSetSampler',
 ]

torch_geometric/explain/algorithm/__init__.py

@@ -5,6 +5,7 @@
 from .pg_explainer import PGExplainer
 from .attention_explainer import AttentionExplainer
 from .graphmask_explainer import GraphMaskExplainer
+from .xgnn_explainer import XGNNExplainer
 
 __all__ = classes = [
     'ExplainerAlgorithm',
@@ -14,4 +15,5 @@
     'PGExplainer',
     'AttentionExplainer',
     'GraphMaskExplainer',
+    'XGNNExplainer'
 ]