| """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) |
|
|
