| 540 | |
| 541 | |
| 542 | class EncSALayer(nn.Module): |
| 543 | def __init__(self, c, num_heads, dropout, attention_dropout=0.1, |
| 544 | relu_dropout=0.1, kernel_size=9, padding='SAME', norm='ln', act='gelu'): |
| 545 | super().__init__() |
| 546 | self.c = c |
| 547 | self.dropout = dropout |
| 548 | self.num_heads = num_heads |
| 549 | if num_heads > 0: |
| 550 | if norm == 'ln': |
| 551 | self.layer_norm1 = LayerNorm(c) |
| 552 | elif norm == 'bn': |
| 553 | self.layer_norm1 = BatchNorm1dTBC(c) |
| 554 | self.self_attn = MultiheadAttention( |
| 555 | self.c, num_heads, self_attention=True, dropout=attention_dropout, bias=False, |
| 556 | ) |
| 557 | if norm == 'ln': |
| 558 | self.layer_norm2 = LayerNorm(c) |
| 559 | elif norm == 'bn': |
| 560 | self.layer_norm2 = BatchNorm1dTBC(c) |
| 561 | self.ffn = TransformerFFNLayer( |
| 562 | c, 4 * c, kernel_size=kernel_size, dropout=relu_dropout, padding=padding, act=act) |
| 563 | |
| 564 | def forward(self, x, encoder_padding_mask=None, **kwargs): |
| 565 | layer_norm_training = kwargs.get('layer_norm_training', None) |
| 566 | if layer_norm_training is not None: |
| 567 | self.layer_norm1.training = layer_norm_training |
| 568 | self.layer_norm2.training = layer_norm_training |
| 569 | if self.num_heads > 0: |
| 570 | residual = x |
| 571 | x = self.layer_norm1(x) |
| 572 | x, _, = self.self_attn( |
| 573 | query=x, |
| 574 | key=x, |
| 575 | value=x, |
| 576 | key_padding_mask=encoder_padding_mask |
| 577 | ) |
| 578 | x = F.dropout(x, self.dropout, training=self.training) |
| 579 | x = residual + x |
| 580 | x = x * (1 - encoder_padding_mask.float()).transpose(0, 1)[..., None] |
| 581 | |
| 582 | residual = x |
| 583 | x = self.layer_norm2(x) |
| 584 | x = self.ffn(x) |
| 585 | x = F.dropout(x, self.dropout, training=self.training) |
| 586 | x = residual + x |
| 587 | x = x * (1 - encoder_padding_mask.float()).transpose(0, 1)[..., None] |
| 588 | return x |
| 589 | |
| 590 | |
| 591 | class DecSALayer(nn.Module): |