| 157 | |
| 158 | class TPA(nn.Module): |
| 159 | def __init__(self, seq, n_fea): |
| 160 | super(TPA, self).__init__() |
| 161 | self.encoder = nn.Conv2d(self.nodes, self.nodes, (n_fea, n_fea), device=device) |
| 162 | # TPA |
| 163 | self.lstm = nn.LSTM(2, 2, num_layers=2, batch_first=True, device=device) |
| 164 | self.fc1 = nn.Linear(in_features=self.seq - 1, out_features=2, device=device) |
| 165 | self.fc2 = nn.Linear(in_features=2, out_features=2, device=device) |
| 166 | self.fc3 = nn.Linear(in_features=2 + 2, out_features=1, device=device) |
| 167 | self.decoder = nn.Linear(self.seq, 1, device=device) |
| 168 | |
| 169 | def forward(self, occ, prc): # occ.shape = [batch, node, seq] |
| 170 | x = torch.stack([occ, prc], dim=3) |