| 40 | |
| 41 | |
| 42 | class TextEncoder(nn.Module): |
| 43 | def __init__(self, channels, kernel_size, depth, n_symbols, actv=nn.LeakyReLU(0.2)): |
| 44 | super().__init__() |
| 45 | self.embedding = nn.Embedding(n_symbols, channels) |
| 46 | |
| 47 | padding = (kernel_size - 1) // 2 |
| 48 | self.cnn = nn.ModuleList() |
| 49 | for _ in range(depth): |
| 50 | self.cnn.append( |
| 51 | nn.Sequential( |
| 52 | weight_norm(nn.Conv1d(channels, channels, kernel_size=kernel_size, padding=padding)), |
| 53 | LayerNorm(channels), |
| 54 | actv, |
| 55 | nn.Dropout(0.2), |
| 56 | ) |
| 57 | ) |
| 58 | # self.cnn = nn.Sequential(*self.cnn) |
| 59 | |
| 60 | self.lstm = nn.LSTM(channels, channels // 2, 1, batch_first=True, bidirectional=True) |
| 61 | |
| 62 | def forward(self, x, input_lengths, m): |
| 63 | x = self.embedding(x) # [B, T, emb] |
| 64 | x = x.transpose(1, 2) # [B, emb, T] |
| 65 | m = m.to(input_lengths.device).unsqueeze(1) |
| 66 | x.masked_fill_(m, 0.0) |
| 67 | |
| 68 | for c in self.cnn: |
| 69 | x = c(x) |
| 70 | x.masked_fill_(m, 0.0) |
| 71 | |
| 72 | x = x.transpose(1, 2) # [B, T, chn] |
| 73 | |
| 74 | input_lengths = input_lengths.cpu().numpy() |
| 75 | x = nn.utils.rnn.pack_padded_sequence(x, input_lengths, batch_first=True, enforce_sorted=False) |
| 76 | |
| 77 | self.lstm.flatten_parameters() |
| 78 | x, _ = self.lstm(x) |
| 79 | x, _ = nn.utils.rnn.pad_packed_sequence(x, batch_first=True) |
| 80 | |
| 81 | x = x.transpose(-1, -2) |
| 82 | x_pad = torch.zeros([x.shape[0], x.shape[1], m.shape[-1]]) |
| 83 | |
| 84 | x_pad[:, :, : x.shape[-1]] = x |
| 85 | x = x_pad.to(x.device) |
| 86 | |
| 87 | x.masked_fill_(m, 0.0) |
| 88 | |
| 89 | return x |
| 90 | |
| 91 | def inference(self, x): |
| 92 | x = self.embedding(x) |
| 93 | x = x.transpose(1, 2) |
| 94 | x = self.cnn(x) |
| 95 | x = x.transpose(1, 2) |
| 96 | self.lstm.flatten_parameters() |
| 97 | x, _ = self.lstm(x) |
| 98 | return x |
| 99 | |