File size: 1,765 Bytes
35d483e | 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 | """Mask-aware temporal pooling shared by student and teacher models."""
from __future__ import annotations
import torch
from torch import Tensor, nn
class MaskedAttentiveStatisticsPooling(nn.Module):
"""Learned weighted mean and standard deviation over valid frames.
Input is ``[batch, channels, frames]`` and output is
``[batch, 2 * channels]``. Multiplicative masking keeps even an all-padded
item finite, which is useful when robustly serving malformed audio.
"""
def __init__(self, channels: int, attention_channels: int = 128) -> None:
super().__init__()
self.attention = nn.Sequential(
nn.Conv1d(channels, attention_channels, kernel_size=1),
nn.Tanh(),
nn.Conv1d(attention_channels, 1, kernel_size=1),
)
def forward(self, x: Tensor, mask: Tensor | None = None) -> Tensor:
if x.ndim != 3:
raise ValueError("pooling input must have shape [batch, channels, frames]")
if mask is None:
mask = torch.ones((x.shape[0], x.shape[-1]), dtype=torch.bool, device=x.device)
if mask.shape != (x.shape[0], x.shape[-1]):
raise ValueError("mask must have shape [batch, frames]")
valid = mask.unsqueeze(1).to(x.dtype)
scores = self.attention(x)
scores = scores.masked_fill(~mask.unsqueeze(1), -1e4)
scores = scores - scores.amax(dim=-1, keepdim=True)
weights = torch.exp(scores) * valid
weights = weights / weights.sum(dim=-1, keepdim=True).clamp_min(1e-8)
mean = (weights * x).sum(dim=-1)
second_moment = (weights * x.square()).sum(dim=-1)
std = (second_moment - mean.square()).clamp_min(1e-5).sqrt()
return torch.cat([mean, std], dim=1)
|