Initial commit: FunASR Speech Recognition Toolkit

Add complete FunASR codebase including models, runtime, and documentation.

funasr/losses/cr_ctc.py

@@ -0,0 +1,55 @@
+"""Consistency-Regularized CTC (CR-CTC) loss.
+
+Based on: "Improving CTC-based Speech Recognition via Consistency Regularization"
+Key idea: Run encoder twice (with/without SpecAug), compute KL divergence between
+the two CTC outputs as a consistency regularization term.
+
+Usage in training:
+    cr_loss = cr_ctc_loss(ctc_logprobs_aug, ctc_logprobs_clean, input_lengths)
+    total_loss = ctc_loss + cr_loss_scale * cr_loss
+"""
+
+import torch
+import torch.nn.functional as F
+
+
+def cr_ctc_loss(
+    log_probs_aug: torch.Tensor,
+    log_probs_clean: torch.Tensor,
+    input_lengths: torch.Tensor,
+) -> torch.Tensor:
+    """Compute CR-CTC consistency regularization loss.
+
+    Computes symmetric KL divergence between augmented and clean encoder outputs.
+
+    Args:
+        log_probs_aug: CTC log probabilities from augmented input (B, T, V)
+        log_probs_clean: CTC log probabilities from clean input (B, T, V)
+        input_lengths: Valid lengths for each sample (B,)
+
+    Returns:
+        Scalar loss value (mean over batch and time).
+    """
+    batch_size, max_len, _ = log_probs_aug.shape
+
+    # Create mask for valid positions
+    mask = torch.arange(max_len, device=input_lengths.device)[None, :] < input_lengths[:, None]
+    mask = mask.unsqueeze(-1)  # (B, T, 1)
+
+    # Convert log probs to probs for KL computation
+    probs_aug = log_probs_aug.exp()
+    probs_clean = log_probs_clean.exp()
+
+    # Symmetric KL divergence: 0.5 * (KL(p||q) + KL(q||p))
+    # KL(p||q) = sum(p * (log_p - log_q))
+    kl_aug_to_clean = (probs_aug * (log_probs_aug - log_probs_clean)) * mask
+    kl_clean_to_aug = (probs_clean * (log_probs_clean - log_probs_aug)) * mask
+
+    # Mean over valid positions
+    num_valid = mask.sum()
+    if num_valid > 0:
+        loss = 0.5 * (kl_aug_to_clean.sum() + kl_clean_to_aug.sum()) / num_valid
+    else:
+        loss = torch.tensor(0.0, device=log_probs_aug.device)
+
+    return loss

funasr/losses/label_smoothing_loss.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+# Copyright 2019 Shigeki Karita
+#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)
+
+"""Label smoothing module."""
+
+import torch
+from torch import nn
+from funasr.models.transformer.utils.nets_utils import make_pad_mask
+
+
+class LabelSmoothingLoss(nn.Module):
+    """Label-smoothing loss.
+
+    :param int size: the number of class
+    :param int padding_idx: ignored class id
+    :param float smoothing: smoothing rate (0.0 means the conventional CE)
+    :param bool normalize_length: normalize loss by sequence length if True
+    :param torch.nn.Module criterion: loss function to be smoothed
+    """
+
+    def __init__(
+        self,
+        size,
+        padding_idx,
+        smoothing,
+        normalize_length=False,
+        criterion=nn.KLDivLoss(reduction="none"),
+    ):
+        """Construct an LabelSmoothingLoss object."""
+        super(LabelSmoothingLoss, self).__init__()
+        self.criterion = criterion
+        self.padding_idx = padding_idx
+        self.confidence = 1.0 - smoothing
+        self.smoothing = smoothing
+        self.size = size
+        self.true_dist = None
+        self.normalize_length = normalize_length
+
+    def forward(self, x, target):
+        """Compute loss between x and target.
+
+        :param torch.Tensor x: prediction (batch, seqlen, class)
+        :param torch.Tensor target:
+            target signal masked with self.padding_id (batch, seqlen)
+        :return: scalar float value
+        :rtype torch.Tensor
+        """
+        assert x.size(2) == self.size
+        batch_size = x.size(0)
+        x = x.contiguous().view(-1, self.size)
+        target = target.contiguous().view(-1)
+        with torch.no_grad():
+            true_dist = x.clone()
+            true_dist.fill_(self.smoothing / (self.size - 1))
+            ignore = target == self.padding_idx  # (B,)
+            total = len(target) - ignore.sum().item()
+            target = target.masked_fill(ignore, 0)  # avoid -1 index
+            true_dist.scatter_(1, target.unsqueeze(1), self.confidence)
+        kl = self.criterion(torch.log_softmax(x, dim=1), true_dist)
+        denom = total if self.normalize_length else batch_size
+        return kl.masked_fill(ignore.unsqueeze(1), 0).sum() / denom
+
+
+class SequenceBinaryCrossEntropy(nn.Module):
+    def __init__(self, normalize_length=False, criterion=nn.BCEWithLogitsLoss(reduction="none")):
+        """Initialize SequenceBinaryCrossEntropy.
+        
+            Args:
+                normalize_length: TODO.
+                criterion: TODO.
+            """
+        super().__init__()
+        self.normalize_length = normalize_length
+        self.criterion = criterion
+
+    def forward(self, pred, label, lengths):
+        """Forward pass for training.
+        
+            Args:
+                pred: TODO.
+                label: TODO.
+                lengths: TODO.
+            """
+        pad_mask = make_pad_mask(lengths, maxlen=pred.shape[1]).to(pred.device)
+        loss = self.criterion(pred, label)
+        denom = (~pad_mask).sum() if self.normalize_length else pred.shape[0]
+        return loss.masked_fill(pad_mask.unsqueeze(-1), 0).sum() / denom
+
+
+class NllLoss(nn.Module):
+    """Nll loss.
+
+    :param int size: the number of class
+    :param int padding_idx: ignored class id
+    :param bool normalize_length: normalize loss by sequence length if True
+    :param torch.nn.Module criterion: loss function
+    """
+
+    def __init__(
+        self,
+        size,
+        padding_idx,
+        normalize_length=False,
+        criterion=nn.NLLLoss(reduction="none"),
+    ):
+        """Construct an NllLoss object."""
+        super(NllLoss, self).__init__()
+        self.criterion = criterion
+        self.padding_idx = padding_idx
+        self.size = size
+        self.true_dist = None
+        self.normalize_length = normalize_length
+
+    def forward(self, x, target):
+        """Compute loss between x and target.
+
+        :param torch.Tensor x: prediction (batch, seqlen, class)
+        :param torch.Tensor target:
+            target signal masked with self.padding_id (batch, seqlen)
+        :return: scalar float value
+        :rtype torch.Tensor
+        """
+        assert x.size(2) == self.size
+        batch_size = x.size(0)
+        x = x.view(-1, self.size)
+        target = target.view(-1)
+        with torch.no_grad():
+            ignore = target == self.padding_idx  # (B,)
+            total = len(target) - ignore.sum().item()
+            target = target.masked_fill(ignore, 0)  # avoid -1 index
+        kl = self.criterion(x, target)
+        denom = total if self.normalize_length else batch_size
+        return kl.masked_fill(ignore, 0).sum() / denom