| 116 | |
| 117 | class LstmGat(nn.Module): |
| 118 | def __init__(self, seq, n_fea, adj_dense, adj_sparse): |
| 119 | super(LstmGat, self).__init__() |
| 120 | self.nodes = adj_dense.shape[0] |
| 121 | self.gcn = nn.Linear(in_features=seq - n_fea + 1, out_features=seq - n_fea + 1, device=device) |
| 122 | self.encoder = nn.Conv2d(self.nodes, self.nodes, (n_fea, n_fea), device=device) |
| 123 | self.gat_l1 = models.MultiHeadsGATLayer(adj_sparse, seq - n_fea + 1, seq - n_fea + 1, 4, 0, 0.2) |
| 124 | self.gat_l2 = models.MultiHeadsGATLayer(adj_sparse, seq - n_fea + 1, seq - n_fea + 1, 4, 0, 0.2) |
| 125 | self.lstm = nn.LSTM(self.nodes, self.nodes, num_layers=2, batch_first=True) |
| 126 | self.decoder = nn.Linear(seq - n_fea + 1, 1, device=device) |
| 127 | |
| 128 | # Activation |
| 129 | self.dropout = nn.Dropout(p=0.5) |
| 130 | self.LeakyReLU = nn.LeakyReLU() |
| 131 | |
| 132 | def forward(self, occ, prc): # occ.shape = [batch, node, seq] |
| 133 | x = torch.stack([occ, prc], dim=3) |