Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,540 Bytes
0102e16 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import torch
from typing import Tuple
from typing import Union
from funasr_detach.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):
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)
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 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), "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
|