FUNASR训练
blame | 历史 | 补丁 | 提交 | 提交对比 | show whitespace
jmwang66
2023-05-16 6f7e27eb7c2d0a7649ec8f14d167c8da8e29f906 
 funasr/models/e2e_asr_paraformer.py


@@ -29,9 +29,8 @@


from funasr.modules.nets_utils import make_pad_mask, pad_list


from funasr.modules.nets_utils import th_accuracy


from funasr.torch_utils.device_funcs import force_gatherable


from funasr.train.abs_espnet_model import AbsESPnetModel


from funasr.models.base_model import FunASRModel


from funasr.models.predictor.cif import CifPredictorV3






if LooseVersion(torch.__version__) >= LooseVersion("1.6.0"):


    from torch.cuda.amp import autocast


@@ -42,7 +41,7 @@


        yield






class Paraformer(AbsESPnetModel):


class Paraformer(FunASRModel):


    """


    Author: Speech Lab of DAMO Academy, Alibaba Group


    Paraformer: Fast and Accurate Parallel Transformer for Non-autoregressive End-to-End Speech Recognition


@@ -56,9 +55,7 @@


            frontend: Optional[AbsFrontend],


            specaug: Optional[AbsSpecAug],


            normalize: Optional[AbsNormalize],


            preencoder: Optional[AbsPreEncoder],


            encoder: AbsEncoder,


            postencoder: Optional[AbsPostEncoder],


            decoder: AbsDecoder,


            ctc: CTC,


            ctc_weight: float = 0.5,


@@ -79,6 +76,8 @@


            predictor_bias: int = 0,


            sampling_ratio: float = 0.2,


            share_embedding: bool = False,


            preencoder: Optional[AbsPreEncoder] = None,


            postencoder: Optional[AbsPostEncoder] = None,


    ):


        assert check_argument_types()


        assert 0.0 <= ctc_weight <= 1.0, ctc_weight


@@ -153,7 +152,6 @@


            text_lengths: torch.Tensor,


    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:


        """Frontend + Encoder + Decoder + Calc loss




        Args:


                speech: (Batch, Length, ...)


                speech_lengths: (Batch, )


@@ -270,7 +268,6 @@


            self, speech: torch.Tensor, speech_lengths: torch.Tensor


    ) -> Tuple[torch.Tensor, torch.Tensor]:


        """Frontend + Encoder. Note that this method is used by asr_inference.py




        Args:


                speech: (Batch, Length, ...)


                speech_lengths: (Batch, )


@@ -368,9 +365,7 @@


            ys_pad_lens: torch.Tensor,


    ) -> torch.Tensor:


        """Compute negative log likelihood(nll) from transformer-decoder




        Normally, this function is called in batchify_nll.




        Args:


                encoder_out: (Batch, Length, Dim)


                encoder_out_lens: (Batch,)


@@ -407,7 +402,6 @@


            batch_size: int = 100,


    ):


        """Compute negative log likelihood(nll) from transformer-decoder




        To avoid OOM, this fuction seperate the input into batches.


        Then call nll for each batch and combine and return results.


        Args:


@@ -664,7 +658,10 @@


            self, speech: torch.Tensor, speech_lengths: torch.Tensor, cache: dict = None


    ) -> Tuple[torch.Tensor, torch.Tensor]:


        """Frontend + Encoder. Note that this method is used by asr_inference.py


<<<<<<< HEAD


=======




>>>>>>> 4cd79db451786548d8a100f25c3b03da0eb30f4b


        Args:


                speech: (Batch, Length, ...)


                speech_lengths: (Batch, )


@@ -738,9 +735,7 @@


            frontend: Optional[AbsFrontend],


            specaug: Optional[AbsSpecAug],


            normalize: Optional[AbsNormalize],


            preencoder: Optional[AbsPreEncoder],


            encoder: AbsEncoder,


            postencoder: Optional[AbsPostEncoder],


            decoder: AbsDecoder,


            ctc: CTC,


            ctc_weight: float = 0.5,


@@ -763,6 +758,8 @@


            embeds_id: int = 2,


            embeds_loss_weight: float = 0.0,


            embed_dims: int = 768,


            preencoder: Optional[AbsPreEncoder] = None,


            postencoder: Optional[AbsPostEncoder] = None,


    ):


        assert check_argument_types()


        assert 0.0 <= ctc_weight <= 1.0, ctc_weight


@@ -894,7 +891,6 @@


            embed_lengths: torch.Tensor = None,


    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:


        """Frontend + Encoder + Decoder + Calc loss




        Args:


                speech: (Batch, Length, ...)


                speech_lengths: (Batch, )


@@ -913,9 +909,9 @@


        self.step_cur += 1


        # for data-parallel


        text = text[:, : text_lengths.max()]


        speech = speech[:, :speech_lengths.max(), :]


        speech = speech[:, :speech_lengths.max()]


        if embed is not None:


            embed = embed[:, :embed_lengths.max(), :]


            embed = embed[:, :embed_lengths.max()]




        # 1. Encoder


        encoder_out, encoder_out_lens = self.encode(speech, speech_lengths)


@@ -1003,7 +999,6 @@






class BiCifParaformer(Paraformer):




    """


    Paraformer model with an extra cif predictor


    to conduct accurate timestamp prediction


@@ -1016,9 +1011,7 @@


        frontend: Optional[AbsFrontend],


        specaug: Optional[AbsSpecAug],


        normalize: Optional[AbsNormalize],


        preencoder: Optional[AbsPreEncoder],


        encoder: AbsEncoder,


        postencoder: Optional[AbsPostEncoder],


        decoder: AbsDecoder,


        ctc: CTC,


        ctc_weight: float = 0.5,


@@ -1038,6 +1031,8 @@


        predictor_weight: float = 0.0,


        predictor_bias: int = 0,


        sampling_ratio: float = 0.2,


            preencoder: Optional[AbsPreEncoder] = None,


            postencoder: Optional[AbsPostEncoder] = None,


    ):


        assert check_argument_types()


        assert 0.0 <= ctc_weight <= 1.0, ctc_weight


@@ -1150,7 +1145,9 @@




        encoder_out_mask = (~make_pad_mask(encoder_out_lens, maxlen=encoder_out.size(1))[:, None, :]).to(


            encoder_out.device)


        pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index, pre_token_length2 = self.predictor(encoder_out, None, encoder_out_mask,


        pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index, pre_token_length2 = self.predictor(encoder_out,


                                                                                                          None,


                                                                                                          encoder_out_mask,


                                                                                  ignore_id=self.ignore_id)


        return pre_acoustic_embeds, pre_token_length, alphas, pre_peak_index


    


@@ -1170,7 +1167,6 @@


            text_lengths: torch.Tensor,


    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:


        """Frontend + Encoder + Decoder + Calc loss




        Args:


                speech: (Batch, Length, ...)


                speech_lengths: (Batch, )


@@ -1253,7 +1249,8 @@


        elif self.ctc_weight == 1.0:


            loss = loss_ctc


        else:


            loss = self.ctc_weight * loss_ctc + (1 - self.ctc_weight) * loss_att + loss_pre * self.predictor_weight + loss_pre2 * self.predictor_weight * 0.5


            loss = self.ctc_weight * loss_ctc + (


                        1 - self.ctc_weight) * loss_att + loss_pre * self.predictor_weight + loss_pre2 * self.predictor_weight * 0.5




        # Collect Attn branch stats


        stats["loss_att"] = loss_att.detach() if loss_att is not None else None


@@ -1282,9 +1279,7 @@


            frontend: Optional[AbsFrontend],


            specaug: Optional[AbsSpecAug],


            normalize: Optional[AbsNormalize],


            preencoder: Optional[AbsPreEncoder],


            encoder: AbsEncoder,


            postencoder: Optional[AbsPostEncoder],


            decoder: AbsDecoder,


            ctc: CTC,


            ctc_weight: float = 0.5,


@@ -1314,6 +1309,8 @@


            bias_encoder_type: str = 'lstm',


            label_bracket: bool = False,


            use_decoder_embedding: bool = False,


            preencoder: Optional[AbsPreEncoder] = None,


            postencoder: Optional[AbsPostEncoder] = None,


    ):


        assert check_argument_types()


        assert 0.0 <= ctc_weight <= 1.0, ctc_weight


@@ -1377,7 +1374,6 @@


            text_lengths: torch.Tensor,


    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor], torch.Tensor]:


        """Frontend + Encoder + Decoder + Calc loss




        Args:


                speech: (Batch, Length, ...)


                speech_lengths: (Batch, )