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
|
import torch
from torch_geometric.datasets import KarateClub
from torch_geometric.nn import GCNConv
from torch_geometric.visualization import influence
class Net(torch.nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
self.conv1 = GCNConv(in_channels, out_channels)
self.conv2 = GCNConv(out_channels, out_channels)
def forward(self, x, edge_index):
x = torch.nn.functional.relu(self.conv1(x, edge_index))
x = self.conv2(x, edge_index)
return x
def test_influence():
data = KarateClub()[0]
x = torch.randn(data.num_nodes, 8)
out = influence(Net(x.size(1), 16), x, data.edge_index)
assert out.size() == (data.num_nodes, data.num_nodes)
assert torch.allclose(out.sum(dim=-1), torch.ones(data.num_nodes),
atol=1e-04)
|
