Initial commit: FunASR Speech Recognition Toolkit

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

funasr/models/specaug/mask_along_axis.py

@@ -0,0 +1,364 @@
+import math
+import torch
+from typing import Sequence
+from typing import Union
+
+
+def mask_along_axis(
+    spec: torch.Tensor,
+    spec_lengths: torch.Tensor,
+    mask_width_range: Sequence[int] = (0, 30),
+    dim: int = 1,
+    num_mask: int = 2,
+    replace_with_zero: bool = True,
+):
+    """Apply mask along the specified direction.
+
+    Args:
+        spec: (Batch, Length, Freq)
+        spec_lengths: (Length): Not using lengths in this implementation
+        mask_width_range: Select the width randomly between this range
+    """
+
+    org_size = spec.size()
+    if spec.dim() == 4:
+        # spec: (Batch, Channel, Length, Freq) -> (Batch * Channel, Length, Freq)
+        spec = spec.view(-1, spec.size(2), spec.size(3))
+
+    B = spec.shape[0]
+    # D = Length or Freq
+    D = spec.shape[dim]
+    # mask_length: (B, num_mask, 1)
+    mask_length = torch.randint(
+        mask_width_range[0],
+        mask_width_range[1],
+        (B, num_mask),
+        device=spec.device,
+    ).unsqueeze(2)
+
+    # mask_pos: (B, num_mask, 1)
+    mask_pos = torch.randint(
+        0, max(1, D - mask_length.max()), (B, num_mask), device=spec.device
+    ).unsqueeze(2)
+
+    # aran: (1, 1, D)
+    aran = torch.arange(D, device=spec.device)[None, None, :]
+    # mask: (Batch, num_mask, D)
+    mask = (mask_pos <= aran) * (aran < (mask_pos + mask_length))
+    # Multiply masks: (Batch, num_mask, D) -> (Batch, D)
+    mask = mask.any(dim=1)
+    if dim == 1:
+        # mask: (Batch, Length, 1)
+        mask = mask.unsqueeze(2)
+    elif dim == 2:
+        # mask: (Batch, 1, Freq)
+        mask = mask.unsqueeze(1)
+
+    if replace_with_zero:
+        value = 0.0
+    else:
+        value = spec.mean()
+
+    if spec.requires_grad:
+        spec = spec.masked_fill(mask, value)
+    else:
+        spec = spec.masked_fill_(mask, value)
+    spec = spec.view(*org_size)
+    return spec, spec_lengths
+
+
+def mask_along_axis_lfr(
+    spec: torch.Tensor,
+    spec_lengths: torch.Tensor,
+    mask_width_range: Sequence[int] = (0, 30),
+    dim: int = 1,
+    num_mask: int = 2,
+    replace_with_zero: bool = True,
+    lfr_rate: int = 1,
+):
+    """Apply mask along the specified direction.
+
+    Args:
+        spec: (Batch, Length, Freq)
+        spec_lengths: (Length): Not using lengths in this implementation
+        mask_width_range: Select the width randomly between this range
+        lfr_rate：low frame rate
+    """
+
+    org_size = spec.size()
+    if spec.dim() == 4:
+        # spec: (Batch, Channel, Length, Freq) -> (Batch * Channel, Length, Freq)
+        spec = spec.view(-1, spec.size(2), spec.size(3))
+
+    B = spec.shape[0]
+    # D = Length or Freq
+    D = spec.shape[dim] // lfr_rate
+    # mask_length: (B, num_mask, 1)
+    mask_length = torch.randint(
+        mask_width_range[0],
+        mask_width_range[1],
+        (B, num_mask),
+        device=spec.device,
+    ).unsqueeze(2)
+    if lfr_rate > 1:
+        mask_length = mask_length.repeat(1, lfr_rate, 1)
+    # mask_pos: (B, num_mask, 1)
+    mask_pos = torch.randint(
+        0, max(1, D - mask_length.max()), (B, num_mask), device=spec.device
+    ).unsqueeze(2)
+    if lfr_rate > 1:
+        mask_pos_raw = mask_pos.clone()
+        mask_pos = torch.zeros((B, 0, 1), device=spec.device, dtype=torch.int32)
+        for i in range(lfr_rate):
+            mask_pos_i = mask_pos_raw + D * i
+            mask_pos = torch.cat((mask_pos, mask_pos_i), dim=1)
+    # aran: (1, 1, D)
+    D = spec.shape[dim]
+    aran = torch.arange(D, device=spec.device)[None, None, :]
+    # mask: (Batch, num_mask, D)
+    mask = (mask_pos <= aran) * (aran < (mask_pos + mask_length))
+    # Multiply masks: (Batch, num_mask, D) -> (Batch, D)
+    mask = mask.any(dim=1)
+    if dim == 1:
+        # mask: (Batch, Length, 1)
+        mask = mask.unsqueeze(2)
+    elif dim == 2:
+        # mask: (Batch, 1, Freq)
+        mask = mask.unsqueeze(1)
+
+    if replace_with_zero:
+        value = 0.0
+    else:
+        value = spec.mean()
+
+    if spec.requires_grad:
+        spec = spec.masked_fill(mask, value)
+    else:
+        spec = spec.masked_fill_(mask, value)
+    spec = spec.view(*org_size)
+    return spec, spec_lengths
+
+
+class MaskAlongAxis(torch.nn.Module):
+    def __init__(
+        self,
+        mask_width_range: Union[int, Sequence[int]] = (0, 30),
+        num_mask: int = 2,
+        dim: Union[int, str] = "time",
+        replace_with_zero: bool = True,
+    ):
+        """Initialize MaskAlongAxis.
+        
+            Args:
+                mask_width_range: TODO.
+                num_mask: TODO.
+                dim: TODO.
+                replace_with_zero: TODO.
+            """
+        if isinstance(mask_width_range, int):
+            mask_width_range = (0, mask_width_range)
+        if len(mask_width_range) != 2:
+            raise TypeError(
+                f"mask_width_range must be a tuple of int and int values: " f"{mask_width_range}",
+            )
+
+        assert mask_width_range[1] > mask_width_range[0]
+        if isinstance(dim, str):
+            if dim == "time":
+                dim = 1
+            elif dim == "freq":
+                dim = 2
+            else:
+                raise ValueError("dim must be int, 'time' or 'freq'")
+        if dim == 1:
+            self.mask_axis = "time"
+        elif dim == 2:
+            self.mask_axis = "freq"
+        else:
+            self.mask_axis = "unknown"
+
+        super().__init__()
+        self.mask_width_range = mask_width_range
+        self.num_mask = num_mask
+        self.dim = dim
+        self.replace_with_zero = replace_with_zero
+
+    def extra_repr(self):
+        """Extra repr."""
+        return (
+            f"mask_width_range={self.mask_width_range}, "
+            f"num_mask={self.num_mask}, axis={self.mask_axis}"
+        )
+
+    def forward(self, spec: torch.Tensor, spec_lengths: torch.Tensor = None):
+        """Forward function.
+
+        Args:
+            spec: (Batch, Length, Freq)
+        """
+
+        return mask_along_axis(
+            spec,
+            spec_lengths,
+            mask_width_range=self.mask_width_range,
+            dim=self.dim,
+            num_mask=self.num_mask,
+            replace_with_zero=self.replace_with_zero,
+        )
+
+
+class MaskAlongAxisVariableMaxWidth(torch.nn.Module):
+    """Mask input spec along a specified axis with variable maximum width.
+
+    Formula:
+        max_width = max_width_ratio * seq_len
+    """
+
+    def __init__(
+        self,
+        mask_width_ratio_range: Union[float, Sequence[float]] = (0.0, 0.05),
+        num_mask: int = 2,
+        dim: Union[int, str] = "time",
+        replace_with_zero: bool = True,
+    ):
+        """Initialize MaskAlongAxisVariableMaxWidth.
+        
+            Args:
+                mask_width_ratio_range: TODO.
+                num_mask: TODO.
+                dim: TODO.
+                replace_with_zero: TODO.
+            """
+        if isinstance(mask_width_ratio_range, float):
+            mask_width_ratio_range = (0.0, mask_width_ratio_range)
+        if len(mask_width_ratio_range) != 2:
+            raise TypeError(
+                f"mask_width_ratio_range must be a tuple of float and float values: "
+                f"{mask_width_ratio_range}",
+            )
+
+        assert mask_width_ratio_range[1] > mask_width_ratio_range[0]
+        if isinstance(dim, str):
+            if dim == "time":
+                dim = 1
+            elif dim == "freq":
+                dim = 2
+            else:
+                raise ValueError("dim must be int, 'time' or 'freq'")
+        if dim == 1:
+            self.mask_axis = "time"
+        elif dim == 2:
+            self.mask_axis = "freq"
+        else:
+            self.mask_axis = "unknown"
+
+        super().__init__()
+        self.mask_width_ratio_range = mask_width_ratio_range
+        self.num_mask = num_mask
+        self.dim = dim
+        self.replace_with_zero = replace_with_zero
+
+    def extra_repr(self):
+        """Extra repr."""
+        return (
+            f"mask_width_ratio_range={self.mask_width_ratio_range}, "
+            f"num_mask={self.num_mask}, axis={self.mask_axis}"
+        )
+
+    def forward(self, spec: torch.Tensor, spec_lengths: torch.Tensor = None):
+        """Forward function.
+
+        Args:
+            spec: (Batch, Length, Freq)
+        """
+
+        max_seq_len = spec.shape[self.dim]
+        min_mask_width = math.floor(max_seq_len * self.mask_width_ratio_range[0])
+        min_mask_width = max([0, min_mask_width])
+        max_mask_width = math.floor(max_seq_len * self.mask_width_ratio_range[1])
+        max_mask_width = min([max_seq_len, max_mask_width])
+
+        if max_mask_width > min_mask_width:
+            return mask_along_axis(
+                spec,
+                spec_lengths,
+                mask_width_range=(min_mask_width, max_mask_width),
+                dim=self.dim,
+                num_mask=self.num_mask,
+                replace_with_zero=self.replace_with_zero,
+            )
+        return spec, spec_lengths
+
+
+class MaskAlongAxisLFR(torch.nn.Module):
+    def __init__(
+        self,
+        mask_width_range: Union[int, Sequence[int]] = (0, 30),
+        num_mask: int = 2,
+        dim: Union[int, str] = "time",
+        replace_with_zero: bool = True,
+        lfr_rate: int = 1,
+    ):
+        """Initialize MaskAlongAxisLFR.
+        
+            Args:
+                mask_width_range: TODO.
+                num_mask: TODO.
+                dim: TODO.
+                replace_with_zero: TODO.
+                lfr_rate: TODO.
+            """
+        if isinstance(mask_width_range, int):
+            mask_width_range = (0, mask_width_range)
+        if len(mask_width_range) != 2:
+            raise TypeError(
+                f"mask_width_range must be a tuple of int and int values: " f"{mask_width_range}",
+            )
+
+        assert mask_width_range[1] > mask_width_range[0]
+        if isinstance(dim, str):
+            if dim == "time":
+                dim = 1
+                lfr_rate = 1
+            elif dim == "freq":
+                dim = 2
+            else:
+                raise ValueError("dim must be int, 'time' or 'freq'")
+        if dim == 1:
+            self.mask_axis = "time"
+            lfr_rate = 1
+        elif dim == 2:
+            self.mask_axis = "freq"
+        else:
+            self.mask_axis = "unknown"
+
+        super().__init__()
+        self.mask_width_range = mask_width_range
+        self.num_mask = num_mask
+        self.dim = dim
+        self.replace_with_zero = replace_with_zero
+        self.lfr_rate = lfr_rate
+
+    def extra_repr(self):
+        """Extra repr."""
+        return (
+            f"mask_width_range={self.mask_width_range}, "
+            f"num_mask={self.num_mask}, axis={self.mask_axis}"
+        )
+
+    def forward(self, spec: torch.Tensor, spec_lengths: torch.Tensor = None):
+        """Forward function.
+
+        Args:
+            spec: (Batch, Length, Freq)
+        """
+
+        return mask_along_axis_lfr(
+            spec,
+            spec_lengths,
+            mask_width_range=self.mask_width_range,
+            dim=self.dim,
+            num_mask=self.num_mask,
+            replace_with_zero=self.replace_with_zero,
+            lfr_rate=self.lfr_rate,
+        )

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

funasr/models/specaug/specaug.py

@@ -0,0 +1,227 @@
+"""SpecAugment module."""
+
+from typing import Optional
+from typing import Sequence
+from typing import Union
+
+from funasr.models.specaug.mask_along_axis import MaskAlongAxis
+from funasr.models.specaug.mask_along_axis import MaskAlongAxisVariableMaxWidth
+from funasr.models.specaug.mask_along_axis import MaskAlongAxisLFR
+from funasr.models.specaug.time_warp import TimeWarp
+from funasr.register import tables
+
+import torch.nn as nn
+
+
+@tables.register("specaug_classes", "SpecAug")
+class SpecAug(nn.Module):
+    """Implementation of SpecAug.
+
+    Reference:
+        Daniel S. Park et al.
+        "SpecAugment: A Simple Data
+         Augmentation Method for Automatic Speech Recognition"
+
+    .. warning::
+        When using cuda mode, time_warp doesn't have reproducibility
+        due to `torch.nn.functional.interpolate`.
+
+    """
+
+    def __init__(
+        self,
+        apply_time_warp: bool = True,
+        time_warp_window: int = 5,
+        time_warp_mode: str = "bicubic",
+        apply_freq_mask: bool = True,
+        freq_mask_width_range: Union[int, Sequence[int]] = (0, 20),
+        num_freq_mask: int = 2,
+        apply_time_mask: bool = True,
+        time_mask_width_range: Optional[Union[int, Sequence[int]]] = None,
+        time_mask_width_ratio_range: Optional[Union[float, Sequence[float]]] = None,
+        num_time_mask: int = 2,
+    ):
+        """Initialize SpecAug.
+        
+            Args:
+                apply_time_warp: TODO.
+                time_warp_window: TODO.
+                time_warp_mode: TODO.
+                apply_freq_mask: TODO.
+                freq_mask_width_range: TODO.
+                num_freq_mask: TODO.
+                apply_time_mask: TODO.
+                time_mask_width_range: TODO.
+                time_mask_width_ratio_range: TODO.
+                num_time_mask: TODO.
+            """
+        if not apply_time_warp and not apply_time_mask and not apply_freq_mask:
+            raise ValueError("Either one of time_warp, time_mask, or freq_mask should be applied")
+        if (
+            apply_time_mask
+            and (time_mask_width_range is not None)
+            and (time_mask_width_ratio_range is not None)
+        ):
+            raise ValueError(
+                'Either one of "time_mask_width_range" or '
+                '"time_mask_width_ratio_range" can be used'
+            )
+        super().__init__()
+        self.apply_time_warp = apply_time_warp
+        self.apply_freq_mask = apply_freq_mask
+        self.apply_time_mask = apply_time_mask
+
+        if apply_time_warp:
+            self.time_warp = TimeWarp(window=time_warp_window, mode=time_warp_mode)
+        else:
+            self.time_warp = None
+
+        if apply_freq_mask:
+            self.freq_mask = MaskAlongAxis(
+                dim="freq",
+                mask_width_range=freq_mask_width_range,
+                num_mask=num_freq_mask,
+            )
+        else:
+            self.freq_mask = None
+
+        if apply_time_mask:
+            if time_mask_width_range is not None:
+                self.time_mask = MaskAlongAxis(
+                    dim="time",
+                    mask_width_range=time_mask_width_range,
+                    num_mask=num_time_mask,
+                )
+            elif time_mask_width_ratio_range is not None:
+                self.time_mask = MaskAlongAxisVariableMaxWidth(
+                    dim="time",
+                    mask_width_ratio_range=time_mask_width_ratio_range,
+                    num_mask=num_time_mask,
+                )
+            else:
+                raise ValueError(
+                    'Either one of "time_mask_width_range" or '
+                    '"time_mask_width_ratio_range" should be used.'
+                )
+        else:
+            self.time_mask = None
+
+    def forward(self, x, x_lengths=None):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+                x_lengths: Lengths of x.
+            """
+        if self.time_warp is not None:
+            x, x_lengths = self.time_warp(x, x_lengths)
+        if self.freq_mask is not None:
+            x, x_lengths = self.freq_mask(x, x_lengths)
+        if self.time_mask is not None:
+            x, x_lengths = self.time_mask(x, x_lengths)
+        return x, x_lengths
+
+
+@tables.register("specaug_classes", "SpecAugLFR")
+class SpecAugLFR(nn.Module):
+    """Implementation of SpecAug.
+    lfr_rate：low frame rate
+    """
+
+    def __init__(
+        self,
+        apply_time_warp: bool = True,
+        time_warp_window: int = 5,
+        time_warp_mode: str = "bicubic",
+        apply_freq_mask: bool = True,
+        freq_mask_width_range: Union[int, Sequence[int]] = (0, 20),
+        num_freq_mask: int = 2,
+        lfr_rate: int = 0,
+        apply_time_mask: bool = True,
+        time_mask_width_range: Optional[Union[int, Sequence[int]]] = None,
+        time_mask_width_ratio_range: Optional[Union[float, Sequence[float]]] = None,
+        num_time_mask: int = 2,
+    ):
+        """Initialize SpecAugLFR.
+        
+            Args:
+                apply_time_warp: TODO.
+                time_warp_window: TODO.
+                time_warp_mode: TODO.
+                apply_freq_mask: TODO.
+                freq_mask_width_range: TODO.
+                num_freq_mask: TODO.
+                lfr_rate: TODO.
+                apply_time_mask: TODO.
+                time_mask_width_range: TODO.
+                time_mask_width_ratio_range: TODO.
+                num_time_mask: TODO.
+            """
+        if not apply_time_warp and not apply_time_mask and not apply_freq_mask:
+            raise ValueError("Either one of time_warp, time_mask, or freq_mask should be applied")
+        if (
+            apply_time_mask
+            and (time_mask_width_range is not None)
+            and (time_mask_width_ratio_range is not None)
+        ):
+            raise ValueError(
+                'Either one of "time_mask_width_range" or '
+                '"time_mask_width_ratio_range" can be used'
+            )
+        super().__init__()
+        self.apply_time_warp = apply_time_warp
+        self.apply_freq_mask = apply_freq_mask
+        self.apply_time_mask = apply_time_mask
+
+        if apply_time_warp:
+            self.time_warp = TimeWarp(window=time_warp_window, mode=time_warp_mode)
+        else:
+            self.time_warp = None
+
+        if apply_freq_mask:
+            self.freq_mask = MaskAlongAxisLFR(
+                dim="freq",
+                mask_width_range=freq_mask_width_range,
+                num_mask=num_freq_mask,
+                lfr_rate=lfr_rate + 1,
+            )
+
+        else:
+            self.freq_mask = None
+
+        if apply_time_mask:
+            if time_mask_width_range is not None:
+                self.time_mask = MaskAlongAxisLFR(
+                    dim="time",
+                    mask_width_range=time_mask_width_range,
+                    num_mask=num_time_mask,
+                    lfr_rate=lfr_rate + 1,
+                )
+            elif time_mask_width_ratio_range is not None:
+                self.time_mask = MaskAlongAxisVariableMaxWidth(
+                    dim="time",
+                    mask_width_ratio_range=time_mask_width_ratio_range,
+                    num_mask=num_time_mask,
+                )
+            else:
+                raise ValueError(
+                    'Either one of "time_mask_width_range" or '
+                    '"time_mask_width_ratio_range" should be used.'
+                )
+        else:
+            self.time_mask = None
+
+    def forward(self, x, x_lengths=None):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+                x_lengths: Lengths of x.
+            """
+        if self.time_warp is not None:
+            x, x_lengths = self.time_warp(x, x_lengths)
+        if self.freq_mask is not None:
+            x, x_lengths = self.freq_mask(x, x_lengths)
+        if self.time_mask is not None:
+            x, x_lengths = self.time_mask(x, x_lengths)
+        return x, x_lengths

funasr/models/specaug/time_warp.py

@@ -0,0 +1,96 @@
+"""Time warp module."""
+
+import torch
+
+from funasr.models.transformer.utils.nets_utils import pad_list
+
+DEFAULT_TIME_WARP_MODE = "bicubic"
+
+
+def time_warp(x: torch.Tensor, window: int = 80, mode: str = DEFAULT_TIME_WARP_MODE):
+    """Time warping using torch.interpolate.
+
+    Args:
+        x: (Batch, Time, Freq)
+        window: time warp parameter
+        mode: Interpolate mode
+    """
+
+    # bicubic supports 4D or more dimension tensor
+    org_size = x.size()
+    if x.dim() == 3:
+        # x: (Batch, Time, Freq) -> (Batch, 1, Time, Freq)
+        x = x[:, None]
+
+    t = x.shape[2]
+    if t - window <= window:
+        return x.view(*org_size)
+
+    center = torch.randint(window, t - window, (1,))[0]
+    warped = torch.randint(center - window, center + window, (1,))[0] + 1
+
+    # left: (Batch, Channel, warped, Freq)
+    # right: (Batch, Channel, time - warped, Freq)
+    left = torch.nn.functional.interpolate(
+        x[:, :, :center], (warped, x.shape[3]), mode=mode, align_corners=False
+    )
+    right = torch.nn.functional.interpolate(
+        x[:, :, center:], (t - warped, x.shape[3]), mode=mode, align_corners=False
+    )
+
+    if x.requires_grad:
+        x = torch.cat([left, right], dim=-2)
+    else:
+        x[:, :, :warped] = left
+        x[:, :, warped:] = right
+
+    return x.view(*org_size)
+
+
+class TimeWarp(torch.nn.Module):
+    """Time warping using torch.interpolate.
+
+    Args:
+        window: time warp parameter
+        mode: Interpolate mode
+    """
+
+    def __init__(self, window: int = 80, mode: str = DEFAULT_TIME_WARP_MODE):
+        """Initialize TimeWarp.
+        
+            Args:
+                window: TODO.
+                mode: TODO.
+            """
+        super().__init__()
+        self.window = window
+        self.mode = mode
+
+    def extra_repr(self):
+        """Extra repr."""
+        return f"window={self.window}, mode={self.mode}"
+
+    def forward(self, x: torch.Tensor, x_lengths: torch.Tensor = None):
+        """Forward function.
+
+        Args:
+            x: (Batch, Time, Freq)
+            x_lengths: (Batch,)
+        """
+
+        if x_lengths is None or all(le == x_lengths[0] for le in x_lengths):
+            # Note that applying same warping for each sample
+            y = time_warp(x, window=self.window, mode=self.mode)
+        else:
+            # FIXME(kamo): I have no idea to batchify Timewarp
+            ys = []
+            for i in range(x.size(0)):
+                _y = time_warp(
+                    x[i][None, : x_lengths[i]],
+                    window=self.window,
+                    mode=self.mode,
+                )[0]
+                ys.append(_y)
+            y = pad_list(ys, 0.0)
+
+        return y, x_lengths