| 5 | |
| 6 | |
| 7 | class Prenet(nn.Module): |
| 8 | def __init__(self, in_dim=80, out_dim=256, kernel=5, n_layers=3, strides=None): |
| 9 | super(Prenet, self).__init__() |
| 10 | padding = kernel // 2 |
| 11 | self.layers = [] |
| 12 | self.strides = strides if strides is not None else [1] * n_layers |
| 13 | for l in range(n_layers): |
| 14 | self.layers.append(nn.Sequential( |
| 15 | nn.Conv1d(in_dim, out_dim, kernel_size=kernel, padding=padding, stride=self.strides[l]), |
| 16 | nn.ReLU(), |
| 17 | nn.BatchNorm1d(out_dim) |
| 18 | )) |
| 19 | in_dim = out_dim |
| 20 | self.layers = nn.ModuleList(self.layers) |
| 21 | self.out_proj = nn.Linear(out_dim, out_dim) |
| 22 | |
| 23 | def forward(self, x): |
| 24 | """ |
| 25 | |
| 26 | :param x: [B, T, 80] |
| 27 | :return: [L, B, T, H], [B, T, H] |
| 28 | """ |
| 29 | padding_mask = x.abs().sum(-1).eq(0).data # [B, T] |
| 30 | nonpadding_mask_TB = 1 - padding_mask.float()[:, None, :] # [B, 1, T] |
| 31 | x = x.transpose(1, 2) |
| 32 | hiddens = [] |
| 33 | for i, l in enumerate(self.layers): |
| 34 | nonpadding_mask_TB = nonpadding_mask_TB[:, :, ::self.strides[i]] |
| 35 | x = l(x) * nonpadding_mask_TB |
| 36 | hiddens.append(x) |
| 37 | hiddens = torch.stack(hiddens, 0) # [L, B, H, T] |
| 38 | hiddens = hiddens.transpose(2, 3) # [L, B, T, H] |
| 39 | x = self.out_proj(x.transpose(1, 2)) # [B, T, H] |
| 40 | x = x * nonpadding_mask_TB.transpose(1, 2) |
| 41 | return hiddens, x |
| 42 | |
| 43 | |
| 44 | class ConvBlock(nn.Module): |