Single decoder layer module. Args: size (int): Input dimension. self_attn (torch.nn.Module): Self-attention module instance. `MultiHeadedAttention` instance can be used as the argument. src_attn (torch.nn.Module): Inter-attention module instance.
| 23 | |
| 24 | |
| 25 | class DecoderLayer(nn.Module): |
| 26 | """Single decoder layer module. |
| 27 | |
| 28 | Args: |
| 29 | size (int): Input dimension. |
| 30 | self_attn (torch.nn.Module): Self-attention module instance. |
| 31 | `MultiHeadedAttention` instance can be used as the argument. |
| 32 | src_attn (torch.nn.Module): Inter-attention module instance. |
| 33 | `MultiHeadedAttention` instance can be used as the argument. |
| 34 | If `None` is passed, Inter-attention is not used, such as |
| 35 | CIF, GPT, and other decoder only model. |
| 36 | feed_forward (torch.nn.Module): Feed-forward module instance. |
| 37 | `PositionwiseFeedForward` instance can be used as the argument. |
| 38 | dropout_rate (float): Dropout rate. |
| 39 | normalize_before (bool): |
| 40 | True: use layer_norm before each sub-block. |
| 41 | False: to use layer_norm after each sub-block. |
| 42 | """ |
| 43 | |
| 44 | def __init__( |
| 45 | self, |
| 46 | size: int, |
| 47 | self_attn: nn.Module, |
| 48 | src_attn: Optional[nn.Module], |
| 49 | feed_forward: nn.Module, |
| 50 | dropout_rate: float, |
| 51 | normalize_before: bool = True, |
| 52 | layer_norm_type: str = 'layer_norm', |
| 53 | norm_eps: float = 1e-5, |
| 54 | ): |
| 55 | """Construct an DecoderLayer object.""" |
| 56 | super().__init__() |
| 57 | self.size = size |
| 58 | self.self_attn = self_attn |
| 59 | self.src_attn = src_attn |
| 60 | self.feed_forward = feed_forward |
| 61 | assert layer_norm_type in ['layer_norm', 'rms_norm'] |
| 62 | self.norm1 = WENET_NORM_CLASSES[layer_norm_type](size, eps=norm_eps) |
| 63 | self.norm2 = WENET_NORM_CLASSES[layer_norm_type](size, eps=norm_eps) |
| 64 | self.norm3 = WENET_NORM_CLASSES[layer_norm_type](size, eps=norm_eps) |
| 65 | self.dropout = nn.Dropout(dropout_rate) |
| 66 | self.normalize_before = normalize_before |
| 67 | |
| 68 | def forward( |
| 69 | self, |
| 70 | tgt: torch.Tensor, |
| 71 | tgt_mask: torch.Tensor, |
| 72 | memory: torch.Tensor, |
| 73 | memory_mask: torch.Tensor, |
| 74 | cache: Optional[Dict[str, Optional[T_CACHE]]] = None |
| 75 | ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 76 | """Compute decoded features. |
| 77 | |
| 78 | Args: |
| 79 | tgt (torch.Tensor): Input tensor (#batch, maxlen_out, size). |
| 80 | tgt_mask (torch.Tensor): Mask for input tensor |
| 81 | (#batch, maxlen_out). |
| 82 | memory (torch.Tensor): Encoded memory |