torchaudio 2.9.0__cp314-cp314-macosx_11_0_arm64.whl

torchaudio/models/squim/subjective.py

@@ -0,0 +1,150 @@
+from typing import Tuple
+
+import torch
+import torch.nn as nn
+import torchaudio
+
+
+class AttPool(nn.Module):
+    """Attention-Pooling module that estimates the attention score.
+
+    Args:
+        input_dim (int): Input feature dimension.
+        att_dim (int): Attention Tensor dimension.
+    """
+
+    def __init__(self, input_dim: int, att_dim: int):
+        super(AttPool, self).__init__()
+
+        self.linear1 = nn.Linear(input_dim, 1)
+        self.linear2 = nn.Linear(input_dim, att_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply attention and pooling.
+
+        Args:
+            x (torch.Tensor): Input Tensor with dimensions `(batch, time, feature_dim)`.
+
+        Returns:
+            (torch.Tensor): Attention score with dimensions `(batch, att_dim)`.
+        """
+
+        att = self.linear1(x)  # (batch, time, 1)
+        att = att.transpose(2, 1)  # (batch, 1, time)
+        att = nn.functional.softmax(att, dim=2)
+        x = torch.matmul(att, x).squeeze(1)  # (batch, input_dim)
+        x = self.linear2(x)  # (batch, att_dim)
+        return x
+
+
+class Predictor(nn.Module):
+    """Prediction module that apply pooling and attention, then predict subjective metric scores.
+
+    Args:
+        input_dim (int): Input feature dimension.
+        att_dim (int): Attention Tensor dimension.
+    """
+
+    def __init__(self, input_dim: int, att_dim: int):
+        super(Predictor, self).__init__()
+        self.att_pool_layer = AttPool(input_dim, att_dim)
+        self.att_dim = att_dim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Predict subjective evaluation metric score.
+
+        Args:
+            x (torch.Tensor): Input Tensor with dimensions `(batch, time, feature_dim)`.
+
+        Returns:
+            (torch.Tensor): Subjective metric score. Tensor with dimensions `(batch,)`.
+        """
+        x = self.att_pool_layer(x)
+        x = nn.functional.softmax(x, dim=1)
+        B = torch.linspace(0, 4, steps=self.att_dim, device=x.device)
+        x = (x * B).sum(dim=1)
+        return x
+
+
+class SquimSubjective(nn.Module):
+    """Speech Quality and Intelligibility Measures (SQUIM) model that predicts **subjective** metric scores
+    for speech enhancement (e.g., Mean Opinion Score (MOS)). The model is adopted from *NORESQA-MOS*
+    :cite:`manocha2022speech` which predicts MOS scores given the input speech and a non-matching reference.
+
+    Args:
+        ssl_model (torch.nn.Module): The self-supervised learning model for feature extraction.
+        projector (torch.nn.Module): Projection layer that projects SSL feature to a lower dimension.
+        predictor (torch.nn.Module): Predict the subjective scores.
+    """
+
+    def __init__(self, ssl_model: nn.Module, projector: nn.Module, predictor: nn.Module):
+        super(SquimSubjective, self).__init__()
+        self.ssl_model = ssl_model
+        self.projector = projector
+        self.predictor = predictor
+
+    def _align_shapes(self, waveform: torch.Tensor, reference: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Cut or pad the reference Tensor to make it aligned with waveform Tensor.
+
+        Args:
+            waveform (torch.Tensor): Input waveform for evaluation. Tensor with dimensions `(batch, time)`.
+            reference (torch.Tensor): Non-matching clean reference. Tensor with dimensions `(batch, time_ref)`.
+
+        Returns:
+            (torch.Tensor, torch.Tensor): The aligned waveform and reference Tensors
+                with same dimensions `(batch, time)`.
+        """
+        T_waveform = waveform.shape[-1]
+        T_reference = reference.shape[-1]
+        if T_reference < T_waveform:
+            num_padding = T_waveform // T_reference + 1
+            reference = torch.cat([reference for _ in range(num_padding)], dim=1)
+        return waveform, reference[:, :T_waveform]
+
+    def forward(self, waveform: torch.Tensor, reference: torch.Tensor):
+        """Predict subjective evaluation metric score.
+
+        Args:
+            waveform (torch.Tensor): Input waveform for evaluation. Tensor with dimensions `(batch, time)`.
+            reference (torch.Tensor): Non-matching clean reference. Tensor with dimensions `(batch, time_ref)`.
+
+        Returns:
+            (torch.Tensor): Subjective metric score. Tensor with dimensions `(batch,)`.
+        """
+        waveform, reference = self._align_shapes(waveform, reference)
+        waveform = self.projector(self.ssl_model.extract_features(waveform)[0][-1])
+        reference = self.projector(self.ssl_model.extract_features(reference)[0][-1])
+        concat = torch.cat((reference, waveform), dim=2)
+        score_diff = self.predictor(concat)  # Score difference compared to the reference
+        return 5 - score_diff
+
+
+def squim_subjective_model(
+    ssl_type: str,
+    feat_dim: int,
+    proj_dim: int,
+    att_dim: int,
+) -> SquimSubjective:
+    """Build a custome :class:`torchaudio.prototype.models.SquimSubjective` model.
+
+    Args:
+        ssl_type (str): Type of self-supervised learning (SSL) models.
+            Must be one of ["wav2vec2_base", "wav2vec2_large"].
+        feat_dim (int): Feature dimension of the SSL feature representation.
+        proj_dim (int): Output dimension of projection layer.
+        att_dim (int): Dimension of attention scores.
+    """
+    ssl_model = getattr(torchaudio.models, ssl_type)()
+    projector = nn.Linear(feat_dim, proj_dim)
+    predictor = Predictor(proj_dim * 2, att_dim)
+    return SquimSubjective(ssl_model, projector, predictor)
+
+
+def squim_subjective_base() -> SquimSubjective:
+    """Build :class:`torchaudio.prototype.models.SquimSubjective` model with default arguments."""
+    return squim_subjective_model(
+        ssl_type="wav2vec2_base",
+        feat_dim=768,
+        proj_dim=32,
+        att_dim=5,
+    )