| 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 |