(self, in_channels, out_channels, seq_len_q, seq_len_kv, modes, c=64,
k=8, ich=512,
L=0,
base='legendre',
mode_select_method='random',
initializer=None, activation='tanh',
**kwargs)
| 81 | """ |
| 82 | |
| 83 | def __init__(self, in_channels, out_channels, seq_len_q, seq_len_kv, modes, c=64, |
| 84 | k=8, ich=512, |
| 85 | L=0, |
| 86 | base='legendre', |
| 87 | mode_select_method='random', |
| 88 | initializer=None, activation='tanh', |
| 89 | **kwargs): |
| 90 | super(MultiWaveletCross, self).__init__() |
| 91 | print('base', base) |
| 92 | |
| 93 | self.c = c |
| 94 | self.k = k |
| 95 | self.L = L |
| 96 | H0, H1, G0, G1, PHI0, PHI1 = get_filter(base, k) |
| 97 | H0r = H0 @ PHI0 |
| 98 | G0r = G0 @ PHI0 |
| 99 | H1r = H1 @ PHI1 |
| 100 | G1r = G1 @ PHI1 |
| 101 | |
| 102 | H0r[np.abs(H0r) < 1e-8] = 0 |
| 103 | H1r[np.abs(H1r) < 1e-8] = 0 |
| 104 | G0r[np.abs(G0r) < 1e-8] = 0 |
| 105 | G1r[np.abs(G1r) < 1e-8] = 0 |
| 106 | self.max_item = 3 |
| 107 | |
| 108 | self.attn1 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, modes=modes, activation=activation) |
| 109 | self.attn2 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, modes=modes, activation=activation) |
| 110 | self.attn3 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, modes=modes, activation=activation) |
| 111 | self.attn4 = FourierCrossAttentionW(in_channels=in_channels, out_channels=out_channels, modes=modes, activation=activation) |
| 112 | |
| 113 | self.T0 = nn.Linear(k, k) |
| 114 | self.register_buffer('ec_s', torch.Tensor( |
| 115 | np.concatenate((H0.T, H1.T), axis=0))) |
| 116 | self.register_buffer('ec_d', torch.Tensor( |
| 117 | np.concatenate((G0.T, G1.T), axis=0))) |
| 118 | |
| 119 | self.register_buffer('rc_e', torch.Tensor( |
| 120 | np.concatenate((H0r, G0r), axis=0))) |
| 121 | self.register_buffer('rc_o', torch.Tensor( |
| 122 | np.concatenate((H1r, G1r), axis=0))) |
| 123 | |
| 124 | self.Lk = nn.Linear(ich, c * k) |
| 125 | self.Lq = nn.Linear(ich, c * k) |
| 126 | self.Lv = nn.Linear(ich, c * k) |
| 127 | self.out = nn.Linear(c * k, ich) |
| 128 | # self.modes1 = modes |
| 129 | |
| 130 | def forward(self, q, k, v, mask=None): |
| 131 | B, N, H, E = q.shape # (B, N, H, E) torch.Size([3, 768, 8, 2]) |
nothing calls this directly
no test coverage detected