Initial commit: FunASR Speech Recognition Toolkit

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

funasr/models/specaug/profileaug.py

@@ -0,0 +1,184 @@
+from typing import Tuple, Optional
+import numpy as np
+import torch
+from torch.nn import functional as F
+import torch.nn as nn
+
+
+class ProfileAug(nn.Module):
+    """
+    Implement the augmentation for profiles including:
+    - Split aug: split one profile into two profiles, i.e., main and inaccurate, labels assigned to main
+    - Merge aug: merge two profiles into one, labels are also merged into one, the other set to zero
+    - Disturb aug: disturb some profile with others to simulate the inaccurate clustering centroids.
+    """
+
+    def __init__(
+        self,
+        apply_split_aug: bool = True,
+        split_aug_prob: float = 0.05,
+        apply_merge_aug: bool = True,
+        merge_aug_prob: float = 0.2,
+        apply_disturb_aug: bool = True,
+        disturb_aug_prob: float = 0.4,
+        disturb_alpha: float = 0.2,
+    ) -> None:
+        """Initialize ProfileAug.
+        
+            Args:
+                apply_split_aug: TODO.
+                split_aug_prob: TODO.
+                apply_merge_aug: TODO.
+                merge_aug_prob: TODO.
+                apply_disturb_aug: TODO.
+                disturb_aug_prob: TODO.
+                disturb_alpha: TODO.
+            """
+        super().__init__()
+        self.apply_split_aug = apply_split_aug
+        self.split_aug_prob = split_aug_prob
+        self.apply_merge_aug = apply_merge_aug
+        self.merge_aug_prob = merge_aug_prob
+        self.apply_disturb_aug = apply_disturb_aug
+        self.disturb_aug_prob = disturb_aug_prob
+        self.disturb_alpha = disturb_alpha
+
+    def split_aug(self, profile: torch.Tensor, binary_labels: torch.Tensor, mask: torch.Tensor):
+        # B, N
+        """Split aug.
+        
+            Args:
+                profile: TODO.
+                binary_labels: TODO.
+                mask: TODO.
+            """
+        bsz, dim = profile.shape[0], profile.shape[-1]
+        profile_norm = torch.linalg.norm(profile, dim=-1, keepdim=False)
+        spk_count = binary_labels.sum(dim=1)
+        prob = np.random.rand(bsz)
+        batch_indices = np.nonzero(prob < self.split_aug_prob)[0]
+        for idx in batch_indices:
+            valid_spk_idx = torch.nonzero(spk_count[idx] * mask[idx])
+            pad_spk_idx = torch.nonzero((spk_count[idx] == 0) * mask[idx])
+            if len(valid_spk_idx) == 0 or len(pad_spk_idx) == 0:
+                continue
+            split_spk_idx = valid_spk_idx[torch.randint(len(valid_spk_idx), ())]
+            to_cover_idx = pad_spk_idx[torch.randint(len(pad_spk_idx), ())]
+            disturb_vec = torch.randn((dim,)).to(profile)
+            disturb_vec = F.normalize(disturb_vec, dim=-1)
+            profile[idx, to_cover_idx] = F.normalize(
+                profile[idx, split_spk_idx] + self.disturb_alpha * disturb_vec
+            )
+            mask[idx, split_spk_idx] = 0
+            mask[idx, to_cover_idx] = 0
+        return profile, binary_labels, mask
+
+    def merge_aug(self, profile: torch.Tensor, binary_labels: torch.Tensor, mask: torch.Tensor):
+        """Merge aug.
+        
+            Args:
+                profile: TODO.
+                binary_labels: TODO.
+                mask: TODO.
+            """
+        bsz, dim = profile.shape[0], profile.shape[-1]
+        profile_norm = torch.linalg.norm(profile, dim=-1, keepdim=False)
+        spk_count = binary_labels.sum(dim=1)
+        prob = np.random.rand(bsz)
+        batch_indices = np.nonzero(prob < self.merge_aug_prob)[0]
+        for idx in batch_indices:
+            valid_spk_idx = torch.nonzero(profile_norm[idx] * mask[idx])
+            if len(valid_spk_idx) == 0:
+                continue
+            to_merge = torch.randint(len(valid_spk_idx), (2,))
+            spk_idx_1, spk_idx_2 = valid_spk_idx[to_merge[0]], valid_spk_idx[to_merge[1]]
+            # merge profile
+            profile[idx, spk_idx_1] = profile[idx, spk_idx_1] + profile[idx, spk_idx_2]
+            profile[idx, spk_idx_1] = F.normalize(profile[idx, spk_idx_1], dim=-1)
+            profile[idx, spk_idx_2] = 0
+            # merge binary labels
+            binary_labels[idx, :, spk_idx_1] = (
+                binary_labels[idx, :, spk_idx_1] + binary_labels[idx, :, spk_idx_2]
+            )
+            binary_labels[idx, :, spk_idx_1] = (binary_labels[idx, :, spk_idx_1] > 0).to(
+                binary_labels
+            )
+            binary_labels[idx, :, spk_idx_2] = 0
+
+            mask[idx, spk_idx_1] = 0
+            mask[idx, spk_idx_2] = 0
+
+        return profile, binary_labels, mask
+
+    def disturb_aug(self, profile: torch.Tensor, binary_labels: torch.Tensor, mask: torch.Tensor):
+        """Disturb aug.
+        
+            Args:
+                profile: TODO.
+                binary_labels: TODO.
+                mask: TODO.
+            """
+        bsz, dim = profile.shape[0], profile.shape[-1]
+        profile_norm = torch.linalg.norm(profile, dim=-1, keepdim=False)
+        spk_count = binary_labels.sum(dim=1)
+        prob = np.random.rand(bsz)
+        batch_indices = np.nonzero(prob < self.disturb_aug_prob)[0]
+        for idx in batch_indices:
+            pos_spk_idx = torch.nonzero(spk_count[idx] * mask[idx])
+            valid_spk_idx = torch.nonzero(profile_norm[idx] * mask[idx])
+            if len(pos_spk_idx) == 0 or len(valid_spk_idx) == 0:
+                continue
+            to_disturb_idx = pos_spk_idx[torch.randint(len(pos_spk_idx), ())]
+            disturb_idx = valid_spk_idx[torch.randint(len(valid_spk_idx), ())]
+            alpha = self.disturb_alpha * torch.rand(()).item()
+            profile[idx, to_disturb_idx] = (1 - alpha) * profile[
+                idx, to_disturb_idx
+            ] + alpha * profile[idx, disturb_idx]
+            profile[idx, to_disturb_idx] = F.normalize(profile[idx, to_disturb_idx], dim=-1)
+            mask[idx, to_disturb_idx] = 0
+
+        return profile, binary_labels, mask
+
+    def forward(
+        self,
+        speech: torch.Tensor,
+        speech_lengths: torch.Tensor = None,
+        profile: torch.Tensor = None,
+        profile_lengths: torch.Tensor = None,
+        binary_labels: torch.Tensor = None,
+        labels_length: torch.Tensor = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+
+        # copy inputs to avoid inplace-operation
+        """Forward pass for training.
+        
+            Args:
+                speech: Speech audio tensor, shape (batch, time).
+                speech_lengths: Length of each speech sample.
+                profile: TODO.
+                profile_lengths: Lengths of profile.
+                binary_labels: TODO.
+                labels_length: TODO.
+            """
+        speech, profile, binary_labels = (
+            torch.clone(speech),
+            torch.clone(profile),
+            torch.clone(binary_labels),
+        )
+        profile = F.normalize(profile, dim=-1)
+
+        profile_mask = torch.ones(profile.shape[:2]).to(profile)
+        if self.apply_disturb_aug:
+            profile, binary_labels, profile_mask = self.disturb_aug(
+                profile, binary_labels, profile_mask
+            )
+        if self.apply_split_aug:
+            profile, binary_labels, profile_mask = self.split_aug(
+                profile, binary_labels, profile_mask
+            )
+        if self.apply_merge_aug:
+            profile, binary_labels, profile_mask = self.merge_aug(
+                profile, binary_labels, profile_mask
+            )
+
+        return speech, profile, binary_labels