Add AttentionExplainer example

CHANGELOG.md

@@ -34,6 +34,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `EdgeIndex.sparse_resize_` functionality ([#8983](https://github.com/pyg-team/pytorch_geometric/pull/8983))
 - Added approximate `faiss`-based KNN-search ([#8952](https://github.com/pyg-team/pytorch_geometric/pull/8952))
 - Added documentation on environment setup on XPU device ([#9407](https://github.com/pyg-team/pytorch_geometric/pull/9407))
+- Added a `attention explainer` example ([#9482](https://github.com/pyg-team/pytorch_geometric/pull/9482))
 
 ### Changed
 

examples/explain/attention_exaplainer.py

@@ -0,0 +1,66 @@
+import os.path as osp
+
+import torch
+import torch.nn.functional as F
+
+import torch_geometric
+from torch_geometric.datasets import Planetoid
+from torch_geometric.explain import AttentionExplainer, Explainer
+from torch_geometric.nn import GATConv
+
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+elif torch_geometric.is_xpu_available():
+    device = torch.device('xpu')
+else:
+    device = torch.device('cpu')
+
+path = osp.join(osp.dirname(osp.realpath(__file__)), 'data', 'Planetoid')
+dataset = Planetoid(path, name='Cora')
+data = dataset[0].to(device)
+
+# GAT Node Classification =====================================================
+
+
+class GAT(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv1 = GATConv(dataset.num_features, 8, heads=8, dropout=0.6)
+        self.conv2 = GATConv(64, dataset.num_classes, heads=1, concat=False,
+                             dropout=0.6)
+
+    def forward(self, x, edge_index):
+        x = F.dropout(x, p=0.6, training=self.training)
+        x = F.elu(self.conv1(x, edge_index))
+        x = F.dropout(x, p=0.6, training=self.training)
+        x = self.conv2(x, edge_index)
+        return x
+
+
+model = GAT().to(device)
+optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
+
+for epoch in range(1, 201):
+    model.train()
+    optimizer.zero_grad()
+    out = model(data.x, data.edge_index)
+    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
+    loss.backward()
+    optimizer.step()
+
+explainer = Explainer(
+    model=model,
+    algorithm=AttentionExplainer(),
+    explanation_type='model',
+    node_mask_type='attributes',
+    edge_mask_type='object',
+    model_config=dict(
+        mode='multiclass_classification',
+        task_level='node',
+        return_type='log_probs',
+    ),
+)
+
+node_index = torch.tensor([10, 20])
+explanation = explainer(data.x, data.edge_index, index=node_index)
+print(f'Generated explanations in {explanation.available_explanations}')

examples/explain/graphmask_explainer.py

@@ -3,11 +3,17 @@
 import torch
 import torch.nn.functional as F
 
+import torch_geometric
 from torch_geometric.datasets import Planetoid
 from torch_geometric.explain import Explainer, GraphMaskExplainer
-from torch_geometric.nn import GATConv, GCNConv
+from torch_geometric.nn import GCNConv
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+elif torch_geometric.is_xpu_available():
+    device = torch.device('xpu')
+else:
+    device = torch.device('cpu')
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), 'data', 'Planetoid')
 dataset = Planetoid(path, name='Cora')
@@ -56,47 +62,3 @@ def forward(self, x, edge_index):
 node_index = 10
 explanation = explainer(data.x, data.edge_index, index=node_index)
 print(f'Generated explanations in {explanation.available_explanations}')
-
-# GAT Node Classification =====================================================
-
-
-class GAT(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.conv1 = GATConv(dataset.num_features, 8, heads=8)
-        self.conv2 = GATConv(64, dataset.num_classes, heads=1, concat=False)
-
-    def forward(self, x, edge_index):
-        x = self.conv1(x, edge_index).relu()
-        x = F.dropout(x, training=self.training)
-        x = self.conv2(x, edge_index)
-        return F.log_softmax(x, dim=1)
-
-
-model = GAT().to(device)
-optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
-
-for epoch in range(1, 201):
-    model.train()
-    optimizer.zero_grad()
-    out = model(data.x, data.edge_index)
-    loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask])
-    loss.backward()
-    optimizer.step()
-
-explainer = Explainer(
-    model=model,
-    algorithm=GraphMaskExplainer(2, epochs=5),
-    explanation_type='model',
-    node_mask_type='attributes',
-    edge_mask_type='object',
-    model_config=dict(
-        mode='multiclass_classification',
-        task_level='node',
-        return_type='log_probs',
-    ),
-)
-
-node_index = torch.tensor([10, 20])
-explanation = explainer(data.x, data.edge_index, index=node_index)
-print(f'Generated explanations in {explanation.available_explanations}')