| 23 | |
| 24 | |
| 25 | class LSTM(nn.Module): |
| 26 | def __init__(self, seq, n_fea, node=247): |
| 27 | super(LSTM, self).__init__() |
| 28 | self.nodes = node |
| 29 | self.encoder = nn.Conv2d(self.nodes, self.nodes, (n_fea, n_fea)) # input.shape: [batch, channel, width, height] |
| 30 | self.lstm = nn.LSTM(self.nodes, self.nodes, num_layers=2, batch_first=True) |
| 31 | self.decoder = nn.Linear(seq-n_fea+1, 1) |
| 32 | |
| 33 | def forward(self, occ, prc): # occ.shape = [batch, node, seq] |
| 34 | x = torch.stack([occ, prc], dim=3) |
| 35 | x = self.encoder(x) |
| 36 | x = torch.transpose(x.squeeze(), 1, 2) # shape [batch, seq-n_fea+1, node] |
| 37 | x, _ = self.lstm(x) |
| 38 | x = torch.transpose(x, 1, 2) # shape [batch, node, seq-n_fea+1] |
| 39 | x = self.decoder(x) |
| 40 | x = torch.squeeze(x) |
| 41 | return x |
| 42 | |
| 43 | |
| 44 | class GCN(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected