| 259 | |
| 260 | |
| 261 | class sparseKernelFT1d(nn.Module): |
| 262 | def __init__(self, |
| 263 | k, alpha, c=1, |
| 264 | nl=1, |
| 265 | initializer=None, |
| 266 | **kwargs): |
| 267 | super(sparseKernelFT1d, self).__init__() |
| 268 | |
| 269 | self.modes1 = alpha |
| 270 | self.scale = (1 / (c * k * c * k)) |
| 271 | self.weights1 = nn.Parameter(self.scale * torch.rand(c * k, c * k, self.modes1, dtype=torch.cfloat)) |
| 272 | self.weights1.requires_grad = True |
| 273 | self.k = k |
| 274 | |
| 275 | def compl_mul1d(self, x, weights): |
| 276 | # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) |
| 277 | return torch.einsum("bix,iox->box", x, weights) |
| 278 | |
| 279 | def forward(self, x): |
| 280 | B, N, c, k = x.shape # (B, N, c, k) |
| 281 | |
| 282 | x = x.view(B, N, -1) |
| 283 | x = x.permute(0, 2, 1) |
| 284 | x_fft = torch.fft.rfft(x) |
| 285 | # Multiply relevant Fourier modes |
| 286 | l = min(self.modes1, N // 2 + 1) |
| 287 | # l = N//2+1 |
| 288 | out_ft = torch.zeros(B, c * k, N // 2 + 1, device=x.device, dtype=torch.cfloat) |
| 289 | out_ft[:, :, :l] = self.compl_mul1d(x_fft[:, :, :l], self.weights1[:, :, :l]) |
| 290 | x = torch.fft.irfft(out_ft, n=N) |
| 291 | x = x.permute(0, 2, 1).view(B, N, c, k) |
| 292 | return x |
| 293 | |
| 294 | |
| 295 | # ## |