| 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] |