| 190 | |
| 191 | |
| 192 | class PitchPredictor(torch.nn.Module): |
| 193 | def __init__(self, idim, n_layers=5, n_chans=384, odim=2, kernel_size=5, |
| 194 | dropout_rate=0.1, padding='SAME'): |
| 195 | """Initilize pitch predictor module. |
| 196 | Args: |
| 197 | idim (int): Input dimension. |
| 198 | n_layers (int, optional): Number of convolutional layers. |
| 199 | n_chans (int, optional): Number of channels of convolutional layers. |
| 200 | kernel_size (int, optional): Kernel size of convolutional layers. |
| 201 | dropout_rate (float, optional): Dropout rate. |
| 202 | """ |
| 203 | super(PitchPredictor, self).__init__() |
| 204 | self.conv = torch.nn.ModuleList() |
| 205 | self.kernel_size = kernel_size |
| 206 | self.padding = padding |
| 207 | for idx in range(n_layers): |
| 208 | in_chans = idim if idx == 0 else n_chans |
| 209 | self.conv += [torch.nn.Sequential( |
| 210 | torch.nn.ConstantPad1d(((kernel_size - 1) // 2, (kernel_size - 1) // 2) |
| 211 | if padding == 'SAME' |
| 212 | else (kernel_size - 1, 0), 0), |
| 213 | torch.nn.Conv1d(in_chans, n_chans, kernel_size, stride=1, padding=0), |
| 214 | torch.nn.ReLU(), |
| 215 | LayerNorm(n_chans, dim=1), |
| 216 | torch.nn.Dropout(dropout_rate) |
| 217 | )] |
| 218 | self.linear = torch.nn.Linear(n_chans, odim) |
| 219 | self.embed_positions = SinusoidalPositionalEmbedding(idim, 0, init_size=4096) |
| 220 | self.pos_embed_alpha = nn.Parameter(torch.Tensor([1])) |
| 221 | |
| 222 | def forward(self, xs): |
| 223 | """ |
| 224 | |
| 225 | :param xs: [B, T, H] |
| 226 | :return: [B, T, H] |
| 227 | """ |
| 228 | positions = self.pos_embed_alpha * self.embed_positions(xs[..., 0]) |
| 229 | xs = xs + positions |
| 230 | xs = xs.transpose(1, -1) # (B, idim, Tmax) |
| 231 | for f in self.conv: |
| 232 | xs = f(xs) # (B, C, Tmax) |
| 233 | # NOTE: calculate in log domain |
| 234 | xs = self.linear(xs.transpose(1, -1)) # (B, Tmax, H) |
| 235 | return xs |
| 236 | |
| 237 | |
| 238 | class EnergyPredictor(PitchPredictor): |