| 456 | |
| 457 | |
| 458 | class sparseKernelFT1d(nn.Module): |
| 459 | def __init__(self, |
| 460 | k, alpha, c=1, |
| 461 | nl=1, |
| 462 | initializer=None, |
| 463 | **kwargs): |
| 464 | super(sparseKernelFT1d, self).__init__() |
| 465 | |
| 466 | self.modes1 = alpha |
| 467 | self.scale = (1 / (c * k * c * k)) |
| 468 | self.weights1 = nn.Parameter(self.scale * torch.rand(c * k, c * k, self.modes1, dtype=torch.float)) |
| 469 | self.weights2 = nn.Parameter(self.scale * torch.rand(c * k, c * k, self.modes1, dtype=torch.float)) |
| 470 | self.weights1.requires_grad = True |
| 471 | self.weights2.requires_grad = True |
| 472 | self.k = k |
| 473 | |
| 474 | def compl_mul1d(self, order, x, weights): |
| 475 | x_flag = True |
| 476 | w_flag = True |
| 477 | if not torch.is_complex(x): |
| 478 | x_flag = False |
| 479 | x = torch.complex(x, torch.zeros_like(x).to(x.device)) |
| 480 | if not torch.is_complex(weights): |
| 481 | w_flag = False |
| 482 | weights = torch.complex(weights, torch.zeros_like(weights).to(weights.device)) |
| 483 | if x_flag or w_flag: |
| 484 | return torch.complex(torch.einsum(order, x.real, weights.real) - torch.einsum(order, x.imag, weights.imag), |
| 485 | torch.einsum(order, x.real, weights.imag) + torch.einsum(order, x.imag, weights.real)) |
| 486 | else: |
| 487 | return torch.einsum(order, x.real, weights.real) |
| 488 | |
| 489 | def forward(self, x): |
| 490 | B, N, c, k = x.shape # (B, N, c, k) |
| 491 | |
| 492 | x = x.view(B, N, -1) |
| 493 | x = x.permute(0, 2, 1) |
| 494 | x_fft = torch.fft.rfft(x) |
| 495 | # Multiply relevant Fourier modes |
| 496 | l = min(self.modes1, N // 2 + 1) |
| 497 | out_ft = torch.zeros(B, c * k, N // 2 + 1, device=x.device, dtype=torch.cfloat) |
| 498 | out_ft[:, :, :l] = self.compl_mul1d("bix,iox->box", x_fft[:, :, :l], |
| 499 | torch.complex(self.weights1, self.weights2)[:, :, :l]) |
| 500 | x = torch.fft.irfft(out_ft, n=N) |
| 501 | x = x.permute(0, 2, 1).view(B, N, c, k) |
| 502 | return x |
| 503 | |
| 504 | |
| 505 | # ## |