zhifu gao
2024-04-24 861147c7308b91068ffa02724fdf74ee623a909e
funasr/models/mfcca/e2e_asr_mfcca.py
@@ -31,9 +31,12 @@
    @contextmanager
    def autocast(enabled=True):
        yield
import pdb
import random
import math
class MFCCA(FunASRModel):
    """
@@ -127,18 +130,15 @@
        assert text_lengths.dim() == 1, text_lengths.shape
        # Check that batch_size is unified
        assert (
                speech.shape[0]
                == speech_lengths.shape[0]
                == text.shape[0]
                == text_lengths.shape[0]
            speech.shape[0] == speech_lengths.shape[0] == text.shape[0] == text_lengths.shape[0]
        ), (speech.shape, speech_lengths.shape, text.shape, text_lengths.shape)
        # pdb.set_trace()
        if (speech.dim() == 3 and speech.size(2) == 8 and self.mask_ratio != 0):
        if speech.dim() == 3 and speech.size(2) == 8 and self.mask_ratio != 0:
            rate_num = random.random()
            # rate_num = 0.1
            if (rate_num <= self.mask_ratio):
            if rate_num <= self.mask_ratio:
                retain_channel = math.ceil(random.random() * 8)
                if (retain_channel > 1):
                if retain_channel > 1:
                    speech = speech[:, :, torch.randperm(8)[0:retain_channel].sort().values]
                else:
                    speech = speech[:, :, torch.randperm(8)[0]]
@@ -230,7 +230,7 @@
            encoder_out.size(),
            speech.size(0),
        )
        if (encoder_out.dim() == 4):
        if encoder_out.dim() == 4:
            assert encoder_out.size(2) <= encoder_out_lens.max(), (
                encoder_out.size(),
                encoder_out_lens.max(),
@@ -272,9 +272,7 @@
        ys_in_lens = ys_pad_lens + 1
        # 1. Forward decoder
        decoder_out, _ = self.decoder(
            encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens
        )
        decoder_out, _ = self.decoder(encoder_out, encoder_out_lens, ys_in_pad, ys_in_lens)
        # 2. Compute attention loss
        loss_att = self.criterion_att(decoder_out, ys_out_pad)
@@ -301,7 +299,7 @@
            ys_pad_lens: torch.Tensor,
    ):
        # Calc CTC loss
        if (encoder_out.dim() == 4):
        if encoder_out.dim() == 4:
            encoder_out = encoder_out.mean(1)
        loss_ctc = self.ctc(encoder_out, encoder_out_lens, ys_pad, ys_pad_lens)