| 380 | return z |
| 381 | |
| 382 | def forward(self, occ, prc): |
| 383 | x = torch.stack([occ, prc], dim=3) |
| 384 | x = self.encoder(x) |
| 385 | x = torch.squeeze(x) |
| 386 | |
| 387 | B, N, L = x.shape |
| 388 | # B*N*L ==> B*NL |
| 389 | x = x.reshape(B, -1) |
| 390 | # embedding B*NL ==> B*NL*D |
| 391 | x = self.tokenEmb(x) |
| 392 | |
| 393 | # FFT B*NL*D ==> B*NT/2*D |
| 394 | x = torch.fft.rfft(x, dim=1, norm='ortho') |
| 395 | |
| 396 | x = x.reshape(B, (N*L)//2+1, self.frequency_size) |
| 397 | |
| 398 | bias = x |
| 399 | |
| 400 | # FourierGNN |
| 401 | x = self.fourierGC(x, B, N, L) |
| 402 | |
| 403 | x = x + bias |
| 404 | |
| 405 | x = x.reshape(B, (N*L)//2+1, self.embed_size) |
| 406 | |
| 407 | # ifft |
| 408 | x = torch.fft.irfft(x, n=N*L, dim=1, norm="ortho") |
| 409 | |
| 410 | x = x.reshape(B, N, L, self.embed_size) |
| 411 | x = x.permute(0, 1, 3, 2) # B, N, D, L |
| 412 | |
| 413 | # projection |
| 414 | x = torch.matmul(x, self.embeddings_10) |
| 415 | x = x.reshape(B, N, -1) |
| 416 | x = self.fc(x) |
| 417 | x = torch.squeeze(x) |
| 418 | return x |
| 419 | |
| 420 | # Other baselines refer to its own original code. |