| 308 | |
| 309 | |
| 310 | class FastspeechEncoder(FFTBlocks): |
| 311 | def __init__(self, embed_tokens, hidden_size=None, num_layers=None, kernel_size=None, num_heads=2): |
| 312 | hidden_size = hparams['hidden_size'] if hidden_size is None else hidden_size |
| 313 | kernel_size = hparams['enc_ffn_kernel_size'] if kernel_size is None else kernel_size |
| 314 | num_layers = hparams['dec_layers'] if num_layers is None else num_layers |
| 315 | super().__init__(hidden_size, num_layers, kernel_size, num_heads=num_heads, |
| 316 | use_pos_embed=False) # use_pos_embed_alpha for compatibility |
| 317 | self.embed_tokens = embed_tokens |
| 318 | self.embed_scale = math.sqrt(hidden_size) |
| 319 | self.padding_idx = 0 |
| 320 | self.embed_positions = SinusoidalPositionalEmbedding( |
| 321 | hidden_size, self.padding_idx, init_size=DEFAULT_MAX_TARGET_POSITIONS, |
| 322 | ) |
| 323 | |
| 324 | def forward(self, txt_tokens): |
| 325 | """ |
| 326 | |
| 327 | :param txt_tokens: [B, T] |
| 328 | :return: { |
| 329 | 'encoder_out': [T x B x C] |
| 330 | } |
| 331 | """ |
| 332 | encoder_padding_mask = txt_tokens.eq(self.padding_idx).data |
| 333 | x = self.forward_embedding(txt_tokens) # [B, T, H] |
| 334 | x = super(FastspeechEncoder, self).forward(x, encoder_padding_mask) |
| 335 | return x |
| 336 | |
| 337 | def forward_embedding(self, txt_tokens): |
| 338 | # embed tokens and positions |
| 339 | x = self.embed_scale * self.embed_tokens(txt_tokens) |
| 340 | if hparams['use_pos_embed']: |
| 341 | positions = self.embed_positions(txt_tokens) |
| 342 | x = x + positions |
| 343 | x = F.dropout(x, p=self.dropout, training=self.training) |
| 344 | return x |
| 345 | |
| 346 | |
| 347 | class FastspeechDecoder(FFTBlocks): |