Initial commit: FunASR Speech Recognition Toolkit

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

funasr/models/sond/pooling/pooling_layers.py

@@ -0,0 +1,144 @@
+# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+
+""" This implementation is adapted from https://github.com/wenet-e2e/wespeaker."""
+
+import torch
+import torch.nn as nn
+
+
+class TAP(nn.Module):
+    """
+    Temporal average pooling, only first-order mean is considered
+    """
+
+    def __init__(self, **kwargs):
+        """Initialize TAP.
+        
+            Args:
+                **kwargs: Additional keyword arguments.
+            """
+        super(TAP, self).__init__()
+
+    def forward(self, x):
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        pooling_mean = x.mean(dim=-1)
+        # To be compatable with 2D input
+        pooling_mean = pooling_mean.flatten(start_dim=1)
+        return pooling_mean
+
+
+class TSDP(nn.Module):
+    """
+    Temporal standard deviation pooling, only second-order std is considered
+    """
+
+    def __init__(self, **kwargs):
+        """Initialize TSDP.
+        
+            Args:
+                **kwargs: Additional keyword arguments.
+            """
+        super(TSDP, self).__init__()
+
+    def forward(self, x):
+        # The last dimension is the temporal axis
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        pooling_std = torch.sqrt(torch.var(x, dim=-1) + 1e-8)
+        pooling_std = pooling_std.flatten(start_dim=1)
+        return pooling_std
+
+
+class TSTP(nn.Module):
+    """
+    Temporal statistics pooling, concatenate mean and std, which is used in
+    x-vector
+    Comment: simple concatenation can not make full use of both statistics
+    """
+
+    def __init__(self, **kwargs):
+        """Initialize TSTP.
+        
+            Args:
+                **kwargs: Additional keyword arguments.
+            """
+        super(TSTP, self).__init__()
+
+    def forward(self, x):
+        # The last dimension is the temporal axis
+        """Forward pass for training.
+        
+            Args:
+                x: TODO.
+            """
+        pooling_mean = x.mean(dim=-1)
+        pooling_std = torch.sqrt(torch.var(x, dim=-1) + 1e-8)
+        pooling_mean = pooling_mean.flatten(start_dim=1)
+        pooling_std = pooling_std.flatten(start_dim=1)
+
+        stats = torch.cat((pooling_mean, pooling_std), 1)
+        return stats
+
+
+class ASTP(nn.Module):
+    """Attentive statistics pooling: Channel- and context-dependent
+    statistics pooling, first used in ECAPA_TDNN.
+    """
+
+    def __init__(self, in_dim, bottleneck_dim=128, global_context_att=False):
+        """Initialize ASTP.
+        
+            Args:
+                in_dim: Size/dimension parameter.
+                bottleneck_dim: Size/dimension parameter.
+                global_context_att: TODO.
+            """
+        super(ASTP, self).__init__()
+        self.global_context_att = global_context_att
+
+        # Use Conv1d with stride == 1 rather than Linear, then we don't
+        # need to transpose inputs.
+        if global_context_att:
+            self.linear1 = nn.Conv1d(
+                in_dim * 3, bottleneck_dim, kernel_size=1
+            )  # equals W and b in the paper
+        else:
+            self.linear1 = nn.Conv1d(
+                in_dim, bottleneck_dim, kernel_size=1
+            )  # equals W and b in the paper
+        self.linear2 = nn.Conv1d(
+            bottleneck_dim, in_dim, kernel_size=1
+        )  # equals V and k in the paper
+
+    def forward(self, x):
+        """
+        x: a 3-dimensional tensor in tdnn-based architecture (B,F,T)
+            or a 4-dimensional tensor in resnet architecture (B,C,F,T)
+            0-dim: batch-dimension, last-dim: time-dimension (frame-dimension)
+        """
+        if len(x.shape) == 4:
+            x = x.reshape(x.shape[0], x.shape[1] * x.shape[2], x.shape[3])
+        assert len(x.shape) == 3
+
+        if self.global_context_att:
+            context_mean = torch.mean(x, dim=-1, keepdim=True).expand_as(x)
+            context_std = torch.sqrt(torch.var(x, dim=-1, keepdim=True) + 1e-10).expand_as(x)
+            x_in = torch.cat((x, context_mean, context_std), dim=1)
+        else:
+            x_in = x
+
+        # DON'T use ReLU here! ReLU may be hard to converge.
+        alpha = torch.tanh(self.linear1(x_in))  # alpha = F.relu(self.linear1(x_in))
+        alpha = torch.softmax(self.linear2(alpha), dim=2)
+        mean = torch.sum(alpha * x, dim=2)
+        var = torch.sum(alpha * (x**2), dim=2) - mean**2
+        std = torch.sqrt(var.clamp(min=1e-10))
+        return torch.cat([mean, std], dim=1)

funasr/models/sond/pooling/statistic_pooling.py

@@ -0,0 +1,126 @@
+import torch
+from typing import Tuple
+from typing import Union
+from funasr.models.transformer.utils.nets_utils import make_non_pad_mask
+from torch.nn import functional as F
+import math
+
+VAR2STD_EPSILON = 1e-12
+
+
+class StatisticPooling(torch.nn.Module):
+    def __init__(self, pooling_dim: Union[int, Tuple] = 2, eps=1e-12):
+        """Initialize StatisticPooling.
+        
+            Args:
+                pooling_dim: Size/dimension parameter.
+                eps: TODO.
+            """
+        super(StatisticPooling, self).__init__()
+        if isinstance(pooling_dim, int):
+            pooling_dim = (pooling_dim,)
+        self.pooling_dim = pooling_dim
+        self.eps = eps
+
+    def forward(self, xs_pad, ilens=None):
+        # xs_pad in (Batch, Channel, Time, Frequency)
+
+        """Forward pass for training.
+        
+            Args:
+                xs_pad: TODO.
+                ilens: TODO.
+            """
+        if ilens is None:
+            masks = torch.ones_like(xs_pad).to(xs_pad)
+        else:
+            masks = make_non_pad_mask(ilens, xs_pad, length_dim=2).to(xs_pad)
+        mean = torch.sum(xs_pad, dim=self.pooling_dim, keepdim=True) / torch.sum(
+            masks, dim=self.pooling_dim, keepdim=True
+        )
+        squared_difference = torch.pow(xs_pad - mean, 2.0)
+        variance = torch.sum(squared_difference, dim=self.pooling_dim, keepdim=True) / torch.sum(
+            masks, dim=self.pooling_dim, keepdim=True
+        )
+        for i in reversed(self.pooling_dim):
+            mean, variance = torch.squeeze(mean, dim=i), torch.squeeze(variance, dim=i)
+
+        mask = torch.less_equal(variance, self.eps).float()
+        variance = (1.0 - mask) * variance + mask * self.eps
+        stddev = torch.sqrt(variance)
+
+        stat_pooling = torch.cat([mean, stddev], dim=1)
+
+        return stat_pooling
+
+
+def statistic_pooling(
+    xs_pad: torch.Tensor, ilens: torch.Tensor = None, pooling_dim: Tuple = (2, 3)
+) -> torch.Tensor:
+    # xs_pad in (Batch, Channel, Time, Frequency)
+
+    """Statistic pooling.
+    
+        Args:
+            xs_pad: TODO.
+            ilens: TODO.
+            pooling_dim: Size/dimension parameter.
+        """
+    if ilens is None:
+        seq_mask = torch.ones_like(xs_pad).to(xs_pad)
+    else:
+        seq_mask = make_non_pad_mask(ilens, xs_pad, length_dim=2).to(xs_pad)
+    mean = torch.sum(xs_pad, dim=pooling_dim, keepdim=True) / torch.sum(
+        seq_mask, dim=pooling_dim, keepdim=True
+    )
+    squared_difference = torch.pow(xs_pad - mean, 2.0)
+    variance = torch.sum(squared_difference, dim=pooling_dim, keepdim=True) / torch.sum(
+        seq_mask, dim=pooling_dim, keepdim=True
+    )
+    for i in reversed(pooling_dim):
+        mean, variance = torch.squeeze(mean, dim=i), torch.squeeze(variance, dim=i)
+
+    value_mask = torch.less_equal(variance, VAR2STD_EPSILON).float()
+    variance = (1.0 - value_mask) * variance + value_mask * VAR2STD_EPSILON
+    stddev = torch.sqrt(variance)
+
+    stat_pooling = torch.cat([mean, stddev], dim=1)
+
+    return stat_pooling
+
+
+def windowed_statistic_pooling(
+    xs_pad: torch.Tensor,
+    ilens: torch.Tensor = None,
+    pooling_dim: Tuple = (2, 3),
+    pooling_size: int = 20,
+    pooling_stride: int = 1,
+) -> Tuple[torch.Tensor, int]:
+    # xs_pad in (Batch, Channel, Time, Frequency)
+
+    """Windowed statistic pooling.
+    
+        Args:
+            xs_pad: TODO.
+            ilens: TODO.
+            pooling_dim: Size/dimension parameter.
+            pooling_size: Size/dimension parameter.
+            pooling_stride: TODO.
+        """
+    tt = xs_pad.shape[2]
+    num_chunk = int(math.ceil(tt / pooling_stride))
+    pad = pooling_size // 2
+    if len(xs_pad.shape) == 4:
+        features = F.pad(xs_pad, (0, 0, pad, pad), "replicate")
+    else:
+        features = F.pad(xs_pad, (pad, pad), "replicate")
+    stat_list = []
+
+    for i in range(num_chunk):
+        # B x C
+        st, ed = i * pooling_stride, i * pooling_stride + pooling_size
+        stat = statistic_pooling(features[:, :, st:ed], pooling_dim=pooling_dim)
+        stat_list.append(stat.unsqueeze(2))
+
+    # B x C x T
+    return torch.cat(stat_list, dim=2), ilens / pooling_stride