1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
|
import pytest
import torch
from torch_geometric.explain import AttentionExplainer, Explainer
from torch_geometric.explain.config import ExplanationType, MaskType
from torch_geometric.nn import AttentiveFP, GATConv, GATv2Conv, TransformerConv
class AttentionGNN(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = GATConv(3, 16, heads=4)
self.conv2 = GATv2Conv(4 * 16, 16, heads=2)
self.conv3 = TransformerConv(2 * 16, 7, heads=1)
def forward(self, x, edge_index):
x = self.conv1(x, edge_index).relu()
x = self.conv2(x, edge_index)
x = self.conv3(x, edge_index)
return x
x = torch.randn(8, 3)
edge_index = torch.tensor([
[0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7],
[1, 0, 2, 1, 3, 2, 4, 3, 5, 4, 6, 5, 7, 6],
])
edge_attr = torch.randn(edge_index.size(1), 5)
batch = torch.tensor([0, 0, 0, 1, 1, 2, 2, 2])
@pytest.mark.parametrize('index', [None, 2, torch.arange(3)])
def test_attention_explainer(index, check_explanation):
explainer = Explainer(
model=AttentionGNN(),
algorithm=AttentionExplainer(),
explanation_type='model',
edge_mask_type='object',
model_config=dict(
mode='multiclass_classification',
task_level='node',
return_type='raw',
),
)
explanation = explainer(x, edge_index, index=index)
check_explanation(explanation, None, explainer.edge_mask_type)
@pytest.mark.parametrize('explanation_type', [e for e in ExplanationType])
@pytest.mark.parametrize('node_mask_type', [m for m in MaskType])
def test_attention_explainer_supports(explanation_type, node_mask_type):
with pytest.raises(ValueError, match="not support the given explanation"):
Explainer(
model=AttentionGNN(),
algorithm=AttentionExplainer(),
explanation_type=explanation_type,
node_mask_type=node_mask_type,
edge_mask_type='object',
model_config=dict(
mode='multiclass_classification',
task_level='node',
return_type='raw',
),
)
def test_attention_explainer_attentive_fp(check_explanation):
model = AttentiveFP(3, 16, 1, edge_dim=5, num_layers=2, num_timesteps=2)
explainer = Explainer(
model=model,
algorithm=AttentionExplainer(),
explanation_type='model',
edge_mask_type='object',
model_config=dict(
mode='binary_classification',
task_level='node',
return_type='raw',
),
)
explanation = explainer(x, edge_index, edge_attr=edge_attr, batch=batch)
check_explanation(explanation, None, explainer.edge_mask_type)
|
