| 48 | |
| 49 | # ########## fourier layer ############# |
| 50 | class FourierBlock(nn.Module): |
| 51 | def __init__(self, in_channels, out_channels, seq_len, modes=0, mode_select_method='random'): |
| 52 | super(FourierBlock, self).__init__() |
| 53 | print('fourier enhanced block used!') |
| 54 | """ |
| 55 | 1D Fourier block. It performs representation learning on frequency domain, |
| 56 | it does FFT, linear transform, and Inverse FFT. |
| 57 | """ |
| 58 | # get modes on frequency domain |
| 59 | self.index = get_frequency_modes(seq_len, modes=modes, mode_select_method=mode_select_method) |
| 60 | print('modes={}, index={}'.format(modes, self.index)) |
| 61 | |
| 62 | self.scale = (1 / (in_channels * out_channels)) |
| 63 | self.weights1 = nn.Parameter( |
| 64 | self.scale * torch.rand(8, in_channels // 8, out_channels // 8, len(self.index), dtype=torch.cfloat)) |
| 65 | |
| 66 | # Complex multiplication |
| 67 | def compl_mul1d(self, input, weights): |
| 68 | # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x) |
| 69 | return torch.einsum("bhi,hio->bho", input, weights) |
| 70 | |
| 71 | def forward(self, q, k, v, mask): |
| 72 | # size = [B, L, H, E] |
| 73 | B, L, H, E = q.shape |
| 74 | x = q.permute(0, 2, 3, 1) |
| 75 | # Compute Fourier coefficients |
| 76 | x_ft = torch.fft.rfft(x, dim=-1) |
| 77 | # Perform Fourier neural operations |
| 78 | out_ft = torch.zeros(B, H, E, L // 2 + 1, device=x.device, dtype=torch.cfloat) |
| 79 | for wi, i in enumerate(self.index): |
| 80 | out_ft[:, :, :, i] = self.compl_mul1d(x_ft[:, :, :, i], self.weights1[:, :, :, wi]) |
| 81 | # Return to time domain |
| 82 | x = torch.fft.irfft(out_ft, n=x.size(-1)) |
| 83 | return (x, None) |
| 84 | |
| 85 | |
| 86 | # ########## Fourier Cross Former #################### |