Method forward

funasr/models/whisper_lid/model.py:587–653

Forward pass for training. Args: speech: Speech audio tensor, shape (batch, time). speech_lengths: Length of each speech sample. lid: TODO. lid_lengths: Lengths of lid.

(
        self,
        speech: torch.Tensor,  # may be padding
        speech_lengths: torch.Tensor,  # actual length
        lid: torch.Tensor,  # lid label, (batch_size, 1)
        lid_lengths: torch.Tensor,
    )

            self.encoder.interctc_use_conditioning = False

    def forward(
        self,
        speech: torch.Tensor,  # may be padding
        speech_lengths: torch.Tensor,  # actual length
        lid: torch.Tensor,  # lid label, (batch_size, 1)
        lid_lengths: torch.Tensor,
    ):
        """Forward pass for training.
        
            Args:
                speech: Speech audio tensor, shape (batch, time).
                speech_lengths: Length of each speech sample.
                lid: TODO.
                lid_lengths: Lengths of lid.
            """
        assert lid.shape[1] == 1
        batch_size = speech.shape[0]
        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)

        # re-generate encoder_out
        if self.clip_frames is None:
            reduced_encoder_out = (
                torch.zeros(batch_size, encoder_out_lens.max(), encoder_out.shape[-1])
                .to(encoder_out.dtype)
                .to(encoder_out.device)
            )
            for i, enc_length in enumerate(encoder_out_lens):
                reduced_encoder_out[i, :enc_length] = encoder_out[i, :enc_length]
        else:
            reduced_encoder_out = (
                torch.zeros(batch_size, self.clip_frames, encoder_out.shape[-1])
                .to(encoder_out.dtype)
                .to(encoder_out.device)
            )
            if self.random_clip:
                for i, enc_length in enumerate(encoder_out_lens):
                    if enc_length <= self.clip_frames:
                        reduced_encoder_out[i, :enc_length] = encoder_out[i, :enc_length]
                        encoder_out_lens[i] = enc_length
                    else:
                        max_start_index = enc_length.item() - self.clip_frames
                        start_index = np.random.randint(0, max_start_index + 1)
                        reduced_encoder_out[i, : self.clip_frames] = encoder_out[
                            i, start_index : start_index + self.clip_frames
                        ]
                        encoder_out_lens[i] = self.clip_frames
            else:
                for i, enc_length in enumerate(encoder_out_lens):
                    enc_length = self.clip_frames if enc_length >= self.clip_frames else enc_length
                    reduced_encoder_out[i, :enc_length] = encoder_out[i, :enc_length]
                    encoder_out_lens[i] = enc_length
        if self.proj_layer is not None:
            reduced_encoder_out = self.proj_layer(reduced_encoder_out)
        lid_output = self.lid_predictor(reduced_encoder_out, encoder_out_lens)  # (B, D)
        lid_logits = self.output_layer(lid_output)  # (B, num_classes)
        loss = self.criterion_lid(lid_logits[:, None, :], lid)
        with torch.no_grad():
            _, predicted_lid = torch.max(lid_logits, 1)