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