1D multiwavelet block.
| 19 | |
| 20 | |
| 21 | class MultiWaveletTransform(nn.Module): |
| 22 | """ |
| 23 | 1D multiwavelet block. |
| 24 | """ |
| 25 | def __init__(self, ich=1, k=8, alpha=16, c=128, |
| 26 | nCZ=1, L=0, base='legendre', attention_dropout=0.1): |
| 27 | super(MultiWaveletTransform, self).__init__() |
| 28 | print('base', base) |
| 29 | self.k = k |
| 30 | self.c = c |
| 31 | self.L = L |
| 32 | self.nCZ = nCZ |
| 33 | self.Lk0 = nn.Linear(ich, c * k) |
| 34 | self.Lk1 = nn.Linear(c * k, ich) |
| 35 | self.ich = ich |
| 36 | self.MWT_CZ = nn.ModuleList(MWT_CZ1d(k, alpha, L, c, base) for i in range(nCZ)) |
| 37 | |
| 38 | def forward(self, queries, keys, values, attn_mask): |
| 39 | B, L, H, E = queries.shape |
| 40 | _, S, _, D = values.shape |
| 41 | if L > S: |
| 42 | zeros = torch.zeros_like(queries[:, :(L - S), :]).float() |
| 43 | values = torch.cat([values, zeros], dim=1) |
| 44 | keys = torch.cat([keys, zeros], dim=1) |
| 45 | else: |
| 46 | values = values[:, :L, :, :] |
| 47 | keys = keys[:, :L, :, :] |
| 48 | values = values.view(B, L, -1) |
| 49 | |
| 50 | V = self.Lk0(values).view(B, L, self.c, -1) |
| 51 | for i in range(self.nCZ): |
| 52 | V = self.MWT_CZ[i](V) |
| 53 | if i < self.nCZ - 1: |
| 54 | V = F.relu(V) |
| 55 | |
| 56 | V = self.Lk1(V.view(B, L, -1)) |
| 57 | V = V.view(B, L, -1, D) |
| 58 | return (V.contiguous(), None) |
| 59 | |
| 60 | |
| 61 | class MultiWaveletCross(nn.Module): |