| 10 | |
| 11 | |
| 12 | class MultiHeadsGATLayer(nn.Module): |
| 13 | def __init__(self, a_sparse, input_dim, out_dim, head_n, dropout, alpha): # input_dim = seq_length |
| 14 | super(MultiHeadsGATLayer, self).__init__() |
| 15 | |
| 16 | self.head_n = head_n |
| 17 | self.heads_dict = dict() |
| 18 | for n in range(head_n): |
| 19 | self.heads_dict[n, 0] = nn.Parameter(torch.zeros(size=(input_dim, out_dim), device=device)) |
| 20 | self.heads_dict[n, 1] = nn.Parameter(torch.zeros(size=(1, 2 * out_dim), device=device)) |
| 21 | nn.init.xavier_normal_(self.heads_dict[n, 0], gain=1.414) |
| 22 | nn.init.xavier_normal_(self.heads_dict[n, 1], gain=1.414) |
| 23 | self.linear = nn.Linear(head_n, 1, device=device) |
| 24 | |
| 25 | # regularization |
| 26 | self.leakyrelu = nn.LeakyReLU(alpha) |
| 27 | self.dropout = nn.Dropout(dropout) |
| 28 | self.softmax = nn.Softmax(dim=0) |
| 29 | |
| 30 | # sparse matrix |
| 31 | self.a_sparse = a_sparse |
| 32 | self.edges = a_sparse.indices() |
| 33 | self.values = a_sparse.values() |
| 34 | self.N = a_sparse.shape[0] |
| 35 | a_dense = a_sparse.to_dense() |
| 36 | a_dense[torch.where(a_dense == 0)] = -1000000000 |
| 37 | a_dense[torch.where(a_dense == 1)] = 0 |
| 38 | self.mask = a_dense |
| 39 | |
| 40 | def forward(self, x): |
| 41 | b, n, s = x.shape |
| 42 | x = x.reshape(b*n, s) |
| 43 | |
| 44 | atts_stack = [] |
| 45 | # multi-heads attention |
| 46 | for n in range(self.head_n): |
| 47 | h = torch.matmul(x, self.heads_dict[n, 0]) |
| 48 | edge_h = torch.cat((h[self.edges[0, :], :], h[self.edges[1, :], :]), dim=1).t() # [Ni, Nj] |
| 49 | atts = self.heads_dict[n, 1].mm(edge_h).squeeze() |
| 50 | atts = self.leakyrelu(atts) |
| 51 | atts_stack.append(atts) |
| 52 | |
| 53 | mt_atts = torch.stack(atts_stack, dim=1) |
| 54 | mt_atts = self.linear(mt_atts) |
| 55 | new_values = self.values * mt_atts.squeeze() |
| 56 | atts_mat = torch.sparse_coo_tensor(self.edges, new_values) |
| 57 | atts_mat = atts_mat.to_dense() + self.mask |
| 58 | atts_mat = self.softmax(atts_mat) |
| 59 | return atts_mat |
| 60 | |
| 61 | |
| 62 | class MLP(nn.Module): |