CTC module. Args: odim: dimension of outputs encoder_output_size: number of encoder projection units dropout_rate: dropout rate (0.0 ~ 1.0) reduce: reduce the CTC loss into a scalar
| 3 | |
| 4 | |
| 5 | class CTC(torch.nn.Module): |
| 6 | """CTC module. |
| 7 | |
| 8 | Args: |
| 9 | odim: dimension of outputs |
| 10 | encoder_output_size: number of encoder projection units |
| 11 | dropout_rate: dropout rate (0.0 ~ 1.0) |
| 12 | reduce: reduce the CTC loss into a scalar |
| 13 | """ |
| 14 | |
| 15 | def __init__( |
| 16 | self, |
| 17 | odim: int, |
| 18 | encoder_output_size: int, |
| 19 | dropout_rate: float = 0.0, |
| 20 | reduce: bool = True, |
| 21 | blank_id: int = 0, |
| 22 | **kwargs, |
| 23 | ): |
| 24 | super().__init__() |
| 25 | eprojs = encoder_output_size |
| 26 | self.dropout_rate = dropout_rate |
| 27 | self.ctc_lo = torch.nn.Linear(eprojs, odim) |
| 28 | self.blank_id = blank_id |
| 29 | self.ctc_loss = torch.nn.CTCLoss(reduction="none", blank=blank_id) |
| 30 | self.reduce = reduce |
| 31 | |
| 32 | def softmax(self, hs_pad): |
| 33 | """softmax of frame activations |
| 34 | |
| 35 | Args: |
| 36 | Tensor hs_pad: 3d tensor (B, Tmax, eprojs) |
| 37 | Returns: |
| 38 | torch.Tensor: softmax applied 3d tensor (B, Tmax, odim) |
| 39 | """ |
| 40 | return F.softmax(self.ctc_lo(hs_pad), dim=2) |
| 41 | |
| 42 | def log_softmax(self, hs_pad): |
| 43 | """log_softmax of frame activations |
| 44 | |
| 45 | Args: |
| 46 | Tensor hs_pad: 3d tensor (B, Tmax, eprojs) |
| 47 | Returns: |
| 48 | torch.Tensor: log softmax applied 3d tensor (B, Tmax, odim) |
| 49 | """ |
| 50 | return F.log_softmax(self.ctc_lo(hs_pad), dim=2) |
| 51 | |
| 52 | def argmax(self, hs_pad): |
| 53 | """argmax of frame activations |
| 54 | |
| 55 | Args: |
| 56 | torch.Tensor hs_pad: 3d tensor (B, Tmax, eprojs) |
| 57 | Returns: |
| 58 | torch.Tensor: argmax applied 2d tensor (B, Tmax) |
| 59 | """ |
| 60 | return torch.argmax(self.ctc_lo(hs_pad), dim=2) |