| 33 | |
| 34 | |
| 35 | class TextEncoder(nn.Module): |
| 36 | def __init__(self, channels, kernel_size, depth, n_symbols, actv=nn.LeakyReLU(0.2)): |
| 37 | super().__init__() |
| 38 | self.embedding = nn.Embedding(n_symbols, channels) |
| 39 | padding = (kernel_size - 1) // 2 |
| 40 | self.cnn = nn.ModuleList() |
| 41 | for _ in range(depth): |
| 42 | self.cnn.append(nn.Sequential( |
| 43 | weight_norm(nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding)), |
| 44 | LayerNorm(channels), |
| 45 | actv, |
| 46 | nn.Dropout(0.2), |
| 47 | )) |
| 48 | self.lstm = nn.LSTM(channels, channels//2, 1, batch_first=True, bidirectional=True) |
| 49 | |
| 50 | def forward(self, x, input_lengths, m): |
| 51 | x = self.embedding(x) # [B, T, emb] |
| 52 | x = x.transpose(1, 2) # [B, emb, T] |
| 53 | m = m.unsqueeze(1) |
| 54 | x.masked_fill_(m, 0.0) |
| 55 | for c in self.cnn: |
| 56 | x = c(x) |
| 57 | x.masked_fill_(m, 0.0) |
| 58 | x = x.transpose(1, 2) # [B, T, chn] |
| 59 | lengths = input_lengths if input_lengths.device == torch.device('cpu') else input_lengths.to('cpu') |
| 60 | x = nn.utils.rnn.pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False) |
| 61 | self.lstm.flatten_parameters() |
| 62 | x, _ = self.lstm(x) |
| 63 | x, _ = nn.utils.rnn.pad_packed_sequence(x, batch_first=True) |
| 64 | x = x.transpose(-1, -2) |
| 65 | x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]], device=x.device) |
| 66 | x_pad[:, :, :x.shape[-1]] = x |
| 67 | x = x_pad |
| 68 | x.masked_fill_(m, 0.0) |
| 69 | return x |
| 70 | |
| 71 | |
| 72 | class AdaLayerNorm(nn.Module): |