Frontend + Encoder + Calc loss Args: speech: (Batch, Length, ...) speech_lengths: (Batch, )
(
self,
speech: torch.Tensor,
speech_lengths: torch.Tensor,
)
| 66 | self.num_updates = 0 |
| 67 | |
| 68 | def forward( |
| 69 | self, |
| 70 | speech: torch.Tensor, |
| 71 | speech_lengths: torch.Tensor, |
| 72 | ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]: |
| 73 | """Frontend + Encoder + Calc loss |
| 74 | Args: |
| 75 | speech: (Batch, Length, ...) |
| 76 | speech_lengths: (Batch, ) |
| 77 | """ |
| 78 | # Check that batch_size is unified |
| 79 | assert speech.shape[0] == speech_lengths.shape[0], (speech.shape, speech_lengths.shape) |
| 80 | |
| 81 | self.encoder.set_num_updates(self.num_updates) |
| 82 | |
| 83 | # 1. Encoder |
| 84 | encoder_out = self.encode(speech, speech_lengths) |
| 85 | |
| 86 | losses = encoder_out["losses"] |
| 87 | loss = sum(losses.values()) |
| 88 | sample_size = encoder_out["sample_size"] |
| 89 | loss = loss.sum() / sample_size |
| 90 | |
| 91 | target_var = float(encoder_out["target_var"]) |
| 92 | pred_var = float(encoder_out["pred_var"]) |
| 93 | ema_decay = float(encoder_out["ema_decay"]) |
| 94 | |
| 95 | stats = dict( |
| 96 | loss=torch.clone(loss.detach()), |
| 97 | target_var=target_var, |
| 98 | pred_var=pred_var, |
| 99 | ema_decay=ema_decay, |
| 100 | ) |
| 101 | |
| 102 | loss, stats, weight = force_gatherable((loss, stats, sample_size), loss.device) |
| 103 | return loss, stats, weight |
| 104 | |
| 105 | def collect_feats( |
| 106 | self, speech: torch.Tensor, speech_lengths: torch.Tensor |
nothing calls this directly
no test coverage detected