Encoder + Decoder + Calc loss Args: speech: (Batch, Length, ...) speech_lengths: (Batch, ) text: (Batch, Length) text_lengths: (Batch,)
(
self,
speech: torch.Tensor,
speech_lengths: torch.Tensor,
text: torch.Tensor,
text_lengths: torch.Tensor,
**kwargs,
)
| 87 | self.length_normalized_loss = length_normalized_loss |
| 88 | |
| 89 | def forward( |
| 90 | self, |
| 91 | speech: torch.Tensor, |
| 92 | speech_lengths: torch.Tensor, |
| 93 | text: torch.Tensor, |
| 94 | text_lengths: torch.Tensor, |
| 95 | **kwargs, |
| 96 | ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: |
| 97 | """Encoder + Decoder + Calc loss |
| 98 | Args: |
| 99 | speech: (Batch, Length, ...) |
| 100 | speech_lengths: (Batch, ) |
| 101 | text: (Batch, Length) |
| 102 | text_lengths: (Batch,) |
| 103 | """ |
| 104 | # import pdb; |
| 105 | # pdb.set_trace() |
| 106 | if len(text_lengths.size()) > 1: |
| 107 | text_lengths = text_lengths[:, 0] |
| 108 | if len(speech_lengths.size()) > 1: |
| 109 | speech_lengths = speech_lengths[:, 0] |
| 110 | |
| 111 | batch_size = speech.shape[0] |
| 112 | |
| 113 | # 1. Encoder |
| 114 | encoder_out, encoder_out_lens = self.encode(speech, speech_lengths) |
| 115 | |
| 116 | loss_ctc, cer_ctc = None, None |
| 117 | stats = dict() |
| 118 | |
| 119 | loss_ctc, cer_ctc = self._calc_ctc_loss( |
| 120 | encoder_out, encoder_out_lens, text, text_lengths |
| 121 | ) |
| 122 | |
| 123 | loss = loss_ctc |
| 124 | |
| 125 | # Collect total loss stats |
| 126 | stats["loss"] = torch.clone(loss.detach()) |
| 127 | |
| 128 | # force_gatherable: to-device and to-tensor if scalar for DataParallel |
| 129 | if self.length_normalized_loss: |
| 130 | batch_size = int((text_lengths + 1).sum()) |
| 131 | loss, stats, weight = force_gatherable((loss, stats, batch_size), loss.device) |
| 132 | return loss, stats, weight |
| 133 | |
| 134 | def encode( |
| 135 | self, |
nothing calls this directly
no test coverage detected