-
Notifications
You must be signed in to change notification settings - Fork 0
/
hgao.py
156 lines (148 loc) · 6.47 KB
/
hgao.py
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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
"""
Graph Representation Learning via Hard Attention Networks in DGL using Adam optimization.
References
----------
Paper: https://arxiv.org/abs/1907.04652
"""
import torch
import torch.nn as nn
import dgl.function as fn
from dgl.nn.pytorch import edge_softmax
from dgl.sampling import select_topk
from functools import partial
from dgl.nn.pytorch.utils import Identity
import torch.nn.functional as F
from dgl.base import DGLError
import dgl
class HardGAO(nn.Module):
def __init__(self,
in_feats,
out_feats,
num_heads=8,
feat_drop=0.,
attn_drop=0.,
negative_slope=0.2,
residual=True,
activation=F.elu,
k=8,):
super(HardGAO, self).__init__()
self.num_heads = num_heads
self.in_feats = in_feats
self.out_feats = out_feats
self.k = k
self.residual = residual
# Initialize Parameters for Additive Attention
self.fc = nn.Linear(
self.in_feats, self.out_feats * self.num_heads, bias=False)
self.attn_l = nn.Parameter(torch.FloatTensor(size=(1, self.num_heads, self.out_feats)))
self.attn_r = nn.Parameter(torch.FloatTensor(size=(1, self.num_heads, self.out_feats)))
# Initialize Parameters for Hard Projection
self.p = nn.Parameter(torch.FloatTensor(size=(1,in_feats)))
# Initialize Dropouts
self.feat_drop = nn.Dropout(feat_drop)
self.attn_drop = nn.Dropout(attn_drop)
self.leaky_relu = nn.LeakyReLU(negative_slope)
if self.residual:
if self.in_feats == self.out_feats:
self.residual_module = Identity()
else:
self.residual_module = nn.Linear(self.in_feats,self.out_feats*num_heads,bias=False)
self.reset_parameters()
self.activation = activation
def reset_parameters(self):
gain = nn.init.calculate_gain('relu')
nn.init.xavier_normal_(self.fc.weight, gain=gain)
nn.init.xavier_normal_(self.p,gain=gain)
nn.init.xavier_normal_(self.attn_l, gain=gain)
nn.init.xavier_normal_(self.attn_r, gain=gain)
if self.residual:
nn.init.xavier_normal_(self.residual_module.weight,gain=gain)
def forward(self, graph, feat, get_attention=False):
# Check in degree and generate error
if (graph.in_degrees()==0).any():
raise DGLError('There are 0-in-degree nodes in the graph, '
'output for those nodes will be invalid. '
'This is harmful for some applications, '
'causing silent performance regression. '
'Adding self-loop on the input graph by '
'calling `g = dgl.add_self_loop(g)` will resolve '
'the issue. Setting ``allow_zero_in_degree`` '
'to be `True` when constructing this module will '
'suppress the check and let the code run.')
# projection process to get importance vector y
graph.ndata['y'] = torch.abs(torch.matmul(self.p,feat.T).view(-1))/torch.norm(self.p,p=2)
# Use edge message passing function to get the weight from src node
graph.apply_edges(fn.copy_u('y','y'))
# Select Top k neighbors
subgraph = select_topk(graph.cpu(),self.k,'y').to(graph.device)
# Sigmoid as information threshold
subgraph.ndata['y'] = torch.sigmoid(subgraph.ndata['y'])
# Using vector matrix elementwise mul for acceleration
feat = subgraph.ndata['y'].view(-1,1)*feat
feat = self.feat_drop(feat)
h = self.fc(feat).view(-1, self.num_heads, self.out_feats)
el = (h * self.attn_l).sum(dim=-1).unsqueeze(-1)
er = (h * self.attn_r).sum(dim=-1).unsqueeze(-1)
# Assign the value on the subgraph
subgraph.srcdata.update({'ft': h, 'el': el})
subgraph.dstdata.update({'er': er})
# compute edge attention, el and er are a_l Wh_i and a_r Wh_j respectively.
subgraph.apply_edges(fn.u_add_v('el', 'er', 'e'))
e = self.leaky_relu(subgraph.edata.pop('e'))
# compute softmax
subgraph.edata['a'] = self.attn_drop(edge_softmax(subgraph, e))
# message passing
subgraph.update_all(fn.u_mul_e('ft', 'a', 'm'),
fn.sum('m', 'ft'))
rst = subgraph.dstdata['ft']
# activation
if self.activation:
rst = self.activation(rst)
# Residual
if self.residual:
rst = rst + self.residual_module(feat).view(feat.shape[0],-1,self.out_feats)
if get_attention:
return rst, subgraph.edata['a']
else:
return rst
class HardGAT(nn.Module):
def __init__(self,
g,
num_layers,
in_dim,
num_hidden,
num_classes,
heads,
activation,
feat_drop,
attn_drop,
negative_slope,
residual,
k):
super(HardGAT, self).__init__()
self.g = g
self.num_layers = num_layers
self.gat_layers = nn.ModuleList()
self.activation = activation
gat_layer = partial(HardGAO,k=k)
muls = heads
# input projection (no residual)
self.gat_layers.append(gat_layer(
in_dim, num_hidden, heads[0],
feat_drop, attn_drop, negative_slope, False, self.activation))
# hidden layers
for l in range(1, num_layers):
# due to multi-head, the in_dim = num_hidden * num_heads
self.gat_layers.append(gat_layer(
num_hidden*muls[l-1] , num_hidden, heads[l],
feat_drop, attn_drop, negative_slope, residual, self.activation))
# output projection
self.gat_layers.append(gat_layer(
num_hidden*muls[-2] , num_classes, heads[-1],
feat_drop, attn_drop, negative_slope, False, None))
def forward(self, inputs):
h = inputs
for l in range(self.num_layers):
h = self.gat_layers[l](self.g, h).flatten(1)
logits = self.gat_layers[-1](self.g, h).mean(1)
return logits