CTC module. Args: odim: dimension of outputs encoder_output_size: number of encoder projection units dropout_rate: dropout rate (0.0 ~ 1.0) ctc_type: builtin or warpctc reduce: reduce the CTC loss into a scalar
| 5 | |
| 6 | |
| 7 | class CTC(torch.nn.Module): |
| 8 | """CTC module. |
| 9 | |
| 10 | Args: |
| 11 | odim: dimension of outputs |
| 12 | encoder_output_size: number of encoder projection units |
| 13 | dropout_rate: dropout rate (0.0 ~ 1.0) |
| 14 | ctc_type: builtin or warpctc |
| 15 | reduce: reduce the CTC loss into a scalar |
| 16 | """ |
| 17 | |
| 18 | def __init__( |
| 19 | self, |
| 20 | odim: int, |
| 21 | encoder_output_size: int, |
| 22 | dropout_rate: float = 0.0, |
| 23 | ctc_type: str = "builtin", |
| 24 | reduce: bool = True, |
| 25 | ignore_nan_grad: bool = True, |
| 26 | extra_linear: bool = True, |
| 27 | ): |
| 28 | """Initialize CTC. |
| 29 | |
| 30 | Args: |
| 31 | odim: TODO. |
| 32 | encoder_output_size: Size/dimension parameter. |
| 33 | dropout_rate: TODO. |
| 34 | ctc_type: TODO. |
| 35 | reduce: TODO. |
| 36 | ignore_nan_grad: TODO. |
| 37 | extra_linear: TODO. |
| 38 | """ |
| 39 | super().__init__() |
| 40 | eprojs = encoder_output_size |
| 41 | self.dropout_rate = dropout_rate |
| 42 | |
| 43 | if extra_linear: |
| 44 | self.ctc_lo = torch.nn.Linear(eprojs, odim) |
| 45 | else: |
| 46 | self.ctc_lo = None |
| 47 | |
| 48 | self.ctc_type = ctc_type |
| 49 | self.ignore_nan_grad = ignore_nan_grad |
| 50 | |
| 51 | if self.ctc_type == "builtin": |
| 52 | self.ctc_loss = torch.nn.CTCLoss(reduction="none") |
| 53 | elif self.ctc_type == "warpctc": |
| 54 | import warpctc_pytorch as warp_ctc |
| 55 | |
| 56 | if ignore_nan_grad: |
| 57 | logging.warning("ignore_nan_grad option is not supported for warp_ctc") |
| 58 | self.ctc_loss = warp_ctc.CTCLoss(size_average=True, reduce=reduce) |
| 59 | else: |
| 60 | raise ValueError(f'ctc_type must be "builtin" or "warpctc": {self.ctc_type}') |
| 61 | |
| 62 | self.reduce = reduce |
| 63 | |
| 64 | def loss_fn(self, th_pred, th_target, th_ilen, th_olen) -> torch.Tensor: |