| 221 | |
| 222 | |
| 223 | class FourierCrossAttentionW(nn.Module): |
| 224 | def __init__(self, in_channels, out_channels, modes=16, activation='tanh'): |
| 225 | super(FourierCrossAttentionW, self).__init__() |
| 226 | print('corss fourier correlation used!') |
| 227 | self.in_channels = in_channels |
| 228 | self.out_channels = out_channels |
| 229 | self.modes1 = modes |
| 230 | self.activation = activation |
| 231 | |
| 232 | def forward(self, q, k, v, mask): |
| 233 | B, L, E, H = q.shape |
| 234 | |
| 235 | xq = q.permute(0, 3, 2, 1) # size = [B, H, E, L] torch.Size([3, 8, 64, 512]) |
| 236 | xk = k.permute(0, 3, 2, 1) |
| 237 | xv = v.permute(0, 3, 2, 1) |
| 238 | self.index_q = list(range(0, min(int(L // 2), self.modes1))) |
| 239 | self.index_k_v = list(range(0, min(int(xv.shape[3] // 2), self.modes1))) |
| 240 | |
| 241 | # Compute Fourier coefficients |
| 242 | xq_ft_ = torch.zeros(B, H, E, len(self.index_q), device=xq.device, dtype=torch.cfloat) |
| 243 | xq_ft = torch.fft.rfft(xq, dim=-1) |
| 244 | for i, j in enumerate(self.index_q): |
| 245 | xq_ft_[:, :, :, i] = xq_ft[:, :, :, j] |
| 246 | |
| 247 | xk_ft_ = torch.zeros(B, H, E, len(self.index_k_v), device=xq.device, dtype=torch.cfloat) |
| 248 | xk_ft = torch.fft.rfft(xk, dim=-1) |
| 249 | for i, j in enumerate(self.index_k_v): |
| 250 | xk_ft_[:, :, :, i] = xk_ft[:, :, :, j] |
| 251 | xqk_ft = (torch.einsum("bhex,bhey->bhxy", xq_ft_, xk_ft_)) |
| 252 | if self.activation == 'tanh': |
| 253 | xqk_ft = xqk_ft.tanh() |
| 254 | elif self.activation == 'softmax': |
| 255 | xqk_ft = torch.softmax(abs(xqk_ft), dim=-1) |
| 256 | xqk_ft = torch.complex(xqk_ft, torch.zeros_like(xqk_ft)) |
| 257 | else: |
| 258 | raise Exception('{} actiation function is not implemented'.format(self.activation)) |
| 259 | xqkv_ft = torch.einsum("bhxy,bhey->bhex", xqk_ft, xk_ft_) |
| 260 | |
| 261 | xqkvw = xqkv_ft |
| 262 | out_ft = torch.zeros(B, H, E, L // 2 + 1, device=xq.device, dtype=torch.cfloat) |
| 263 | for i, j in enumerate(self.index_q): |
| 264 | out_ft[:, :, :, j] = xqkvw[:, :, :, i] |
| 265 | |
| 266 | out = torch.fft.irfft(out_ft / self.in_channels / self.out_channels, n=xq.size(-1)).permute(0, 3, 2, 1) |
| 267 | # size = [B, L, H, E] |
| 268 | return (out, None) |
| 269 | |
| 270 | |
| 271 | class sparseKernelFT1d(nn.Module): |