| 4 | |
| 5 | |
| 6 | class ConformerEncoder(nn.Module): |
| 7 | def __init__(self, idim, n_layers, n_head, d_model, |
| 8 | residual_dropout=0.1, dropout_rate=0.1, kernel_size=33, |
| 9 | pe_maxlen=5000): |
| 10 | super().__init__() |
| 11 | self.odim = d_model |
| 12 | |
| 13 | self.input_preprocessor = Conv2dSubsampling(idim, d_model) |
| 14 | self.positional_encoding = RelPositionalEncoding(d_model) |
| 15 | self.dropout = nn.Dropout(residual_dropout) |
| 16 | |
| 17 | self.layer_stack = nn.ModuleList() |
| 18 | for l in range(n_layers): |
| 19 | block = RelPosEmbConformerBlock(d_model, n_head, |
| 20 | residual_dropout, |
| 21 | dropout_rate, kernel_size) |
| 22 | self.layer_stack.append(block) |
| 23 | |
| 24 | def forward(self, padded_input, input_lengths, pad=True): |
| 25 | if pad: |
| 26 | padded_input = F.pad(padded_input, |
| 27 | (0, 0, 0, self.input_preprocessor.context - 1), 'constant', 0.0) |
| 28 | src_mask = self.padding_position_is_0(padded_input, input_lengths) |
| 29 | |
| 30 | embed_output, input_lengths, src_mask = self.input_preprocessor( |
| 31 | padded_input, src_mask, input_lengths |
| 32 | ) |
| 33 | enc_output = self.dropout(embed_output) |
| 34 | |
| 35 | pos_emb = self.dropout(self.positional_encoding(embed_output)) |
| 36 | |
| 37 | enc_outputs = [] |
| 38 | for enc_layer in self.layer_stack: |
| 39 | enc_output = enc_layer(enc_output, pos_emb, slf_attn_mask=src_mask, |
| 40 | pad_mask=src_mask) |
| 41 | enc_outputs.append(enc_output) |
| 42 | |
| 43 | return enc_output, input_lengths, src_mask |
| 44 | |
| 45 | def padding_position_is_0(self, padded_input, input_lengths): |
| 46 | N, T = padded_input.size()[:2] |
| 47 | mask = torch.ones((N, T)).to(padded_input.device) |
| 48 | for i in range(N): |
| 49 | mask[i, input_lengths[i]:] = 0 |
| 50 | mask = mask.unsqueeze(dim=1) |
| 51 | return mask.to(torch.bool) |
| 52 | |
| 53 | |
| 54 | class RelPosEmbConformerBlock(nn.Module): |