add sond model

志浩

2023-02-10 f6a1cdaf3488c9ec572e1f753b50cb58a0f8fd79

 funasr/models/pooling/statistic_pooling.py

@@ -2,7 +2,10 @@
from typing import Tuple
from typing import Union
from funasr.modules.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):
@@ -34,3 +37,59 @@
        stat_pooling = torch.cat([mean, stddev], dim=1)

        return stat_pooling

    def convert_tf2torch(self, var_dict_tf, var_dict_torch):
        return {}


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)

    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)

    tt = xs_pad.shape[2]
    num_chunk = int(math.ceil(tt / pooling_stride))
    pad = pooling_size // 2
    features = F.pad(xs_pad, (0, 0, pad, pad), "reflect")
    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