GraphML/gma.py at main · Anuoluwapo65/GraphML · GitHub

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
85
86
87
88
89
90
91
92
93
94
95
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision

class gma(nn.Module):
    def __init__(self, c_in, c_out, num_heads = 1, concat_head = True, alpha = 0.2):
        super().__init__()

        self.num_heads = num_heads
        self.concat_head = concat_head
        assert num_heads > 0, "num_heads must be greater than 0"
        if self.concat_head:
            assert c_out % num_heads == 0, f"the num_heads must be divisible by the c_out values"
            c_out = c_out // num_heads


        self.projection = nn.Linear(c_in, c_out * num_heads)
        self.a = nn.Parameter(torch.Tensor(num_heads, 2 * c_out))
        self.leaky_relu = nn.LeakyReLU(alpha)

        nn.init.xavier_uniform_(self.projection.weight, gain = 1.414)
        nn.init.xavier_uniform_(self.a, gain = 1.414)

    def forward(self, adjacent_matrix, node_feats, print_trob_feats = False):
        batch_size, num_nodes = node_feats.size(0), node_feats.size(1)
        node_feats = self.projection(node_feats)
        node_feats = node_feats.view(batch_size, num_nodes, self.num_heads, -1)
        edges = adjacent_matrix.nonzero(as_tuple = False)
        node_feats_view = node_feats.view(batch_size * num_nodes, self.num_heads, -1)

        edges_rows = edges[:, 0] * num_nodes + edges[:, 1]
        edges_cols = edges[:, 0] * num_nodes + edges[:, 2]

        a_input = torch.cat([torch.index_select(input = node_feats_view, index = edges_rows, dim = 0),
                             torch.index_select(input = node_feats_view, index = edges_cols, dim = 0)],
                             dim = -1)

        sums = torch.einsum("ehc,hc->eh", a_input, self.a)
        attn = self.leaky_relu(sums)

        attns = attn.new_zeros(adjacent_matrix.shape+(self.num_heads,)).fill_(-9e15)
        attns[adjacent_matrix[...,None].repeat(1,1,1,self.num_heads)==1] = attn.reshape(-1)


        attn = F.softmax(attns, dim = 2)
        if print_trob_feats:
            print("Attention prob\n", attn.permute(0,3,1,2))
        node_feats = torch.einsum("bijh,bjhc->bihc", attn, node_feats)


        if self.concat_head:
            node_feats = node_feats.reshape(batch_size, num_nodes, -1)
        else:
            node_feats = node_feats.mean(dim = 2)

        return node_feats


if __name__ == "__main__":
    layer = gma(2, 2, num_heads=2)
    layer.projection.weight.data = torch.Tensor([[1., 0.], [0., 1.]])
    layer.projection.bias.data = torch.Tensor([0., 0.])
    layer.a.data = torch.Tensor([[-0.2, 0.3], [0.1, -0.1]])

    node_feats = torch.Tensor([[
        [0., 1.],
        [2., 3.],
        [4., 5.],
        [6., 7.]
    ]])

    adjacent_matrix = torch.Tensor([[
        [1, 1, 0, 0],
        [1, 1, 1, 0],
        [0, 1, 1, 1],
        [0, 0, 1, 1]
    ]])

    with torch.no_grad():
        out_feats = layer(adjacent_matrix, node_feats, print_trob_feats=False)


    print("Adjacency matrix", adjacent_matrix)
    print("Input features", node_feats)
    print("Output features", out_feats)