From 6427c834dfd97b1f05c6659cdc7ccf010bf82fe1 Mon Sep 17 00:00:00 2001
From: 嘉渊 <wangjiaming.wjm@alibaba-inc.com>
Date: 星期一, 24 四月 2023 19:50:07 +0800
Subject: [PATCH] update
---
funasr/models/pooling/statistic_pooling.py | 62 +++++++++++++++++++++++++++++++
1 files changed, 62 insertions(+), 0 deletions(-)
diff --git a/funasr/models/pooling/statistic_pooling.py b/funasr/models/pooling/statistic_pooling.py
index eeaed7d..8f85de9 100644
--- a/funasr/models/pooling/statistic_pooling.py
+++ b/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,62 @@
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
+ if len(xs_pad.shape) == 4:
+ features = F.pad(xs_pad, (0, 0, pad, pad), "reflect")
+ else:
+ features = F.pad(xs_pad, (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
--
Gitblit v1.9.1