torchaudio 2.0.2__cp310-cp310-manylinux2014_aarch64.whl → 2.1.1__cp310-cp310-manylinux2014_aarch64.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,
+    )

torchaudio/models/wav2vec2/components.py

@@ -208,18 +208,13 @@ class ConvolutionalPositionalEmbedding(Module):
             groups=groups,
         )
 
-        self.conv = nn.utils.weight_norm(self.conv, name="weight", dim=2)
+        self.conv = nn.utils.parametrizations.weight_norm(self.conv, name="weight", dim=2)
         self.num_remove: int = 1 if kernel_size % 2 == 0 else 0
 
     def __prepare_scriptable__(self):
-        for hook in self.conv._forward_pre_hooks.values():
-            # The hook we want to remove is an instance of WeightNorm class, so
-            # normally we would do `if isinstance(...)` but this class is not accessible
-            # because of shadowing, so we check the module name directly.
-            # https://github.com/pytorch/pytorch/blob/be0ca00c5ce260eb5bcec3237357f7a30cc08983/torch/nn/utils/__init__.py#L3
-            if hook.__module__ == "torch.nn.utils.weight_norm" and hook.__class__.__name__ == "WeightNorm":
-                _LG.warning("Removing weight_norm from %s", self.__class__.__name__)
-                torch.nn.utils.remove_weight_norm(self.conv)
+        if self.conv.__class__.__name__ == "ParametrizedConv1d":
+            _LG.warning("Removing weight_norm from %s", self.__class__.__name__)
+            torch.nn.utils.parametrize.remove_parametrizations(self.conv, "weight")
         return self
 
     def forward(self, x):
@@ -458,9 +453,10 @@ class Transformer(Module):
                 raise ValueError(f"`num_layers` must be between [1, {len(self.layers)}]")
 
         ret: List[Tensor] = []
+        position_bias = None
         x = self._preprocess(x)
         for layer in self.layers:
-            x, _ = layer(x, attention_mask)  # Ignore position_bias
+            x, position_bias = layer(x, attention_mask, position_bias=position_bias)
             ret.append(x)
             if num_layers is not None and len(ret) >= num_layers:
                 return ret

torchaudio/pipelines/__init__.py

@@ -4,6 +4,7 @@ from ._source_separation_pipeline import (
     HDEMUCS_HIGH_MUSDB_PLUS,
     SourceSeparationBundle,
 )
+from ._squim_pipeline import SQUIM_OBJECTIVE, SQUIM_SUBJECTIVE, SquimObjectiveBundle, SquimSubjectiveBundle
 from ._tts import (
     TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH,
     TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH,
@@ -17,6 +18,7 @@ from ._wav2vec2.impl import (
     HUBERT_BASE,
     HUBERT_LARGE,
     HUBERT_XLARGE,
+    MMS_FA,
     VOXPOPULI_ASR_BASE_10K_DE,
     VOXPOPULI_ASR_BASE_10K_EN,
     VOXPOPULI_ASR_BASE_10K_ES,
@@ -40,6 +42,7 @@ from ._wav2vec2.impl import (
     WAV2VEC2_XLSR_300M,
     Wav2Vec2ASRBundle,
     Wav2Vec2Bundle,
+    Wav2Vec2FABundle,
     WAVLM_BASE,
     WAVLM_BASE_PLUS,
     WAVLM_LARGE,
@@ -50,6 +53,7 @@ from .rnnt_pipeline import EMFORMER_RNNT_BASE_LIBRISPEECH, RNNTBundle
 __all__ = [
     "Wav2Vec2Bundle",
     "Wav2Vec2ASRBundle",
+    "Wav2Vec2FABundle",
     "WAV2VEC2_BASE",
     "WAV2VEC2_LARGE",
     "WAV2VEC2_LARGE_LV60K",
@@ -76,6 +80,7 @@ __all__ = [
     "HUBERT_XLARGE",
     "HUBERT_ASR_LARGE",
     "HUBERT_ASR_XLARGE",
+    "MMS_FA",
     "WAVLM_BASE",
     "WAVLM_BASE_PLUS",
     "WAVLM_LARGE",
@@ -90,4 +95,8 @@ __all__ = [
     "CONVTASNET_BASE_LIBRI2MIX",
     "HDEMUCS_HIGH_MUSDB_PLUS",
     "HDEMUCS_HIGH_MUSDB",
+    "SQUIM_OBJECTIVE",
+    "SQUIM_SUBJECTIVE",
+    "SquimObjectiveBundle",
+    "SquimSubjectiveBundle",
 ]

torchaudio/pipelines/_squim_pipeline.py

@@ -0,0 +1,176 @@
+from dataclasses import dataclass
+
+from torchaudio._internal import load_state_dict_from_url
+
+from torchaudio.models import squim_objective_base, squim_subjective_base, SquimObjective, SquimSubjective
+
+
+@dataclass
+class SquimObjectiveBundle:
+    """Data class that bundles associated information to use pretrained
+    :py:class:`~torchaudio.models.SquimObjective` model.
+
+    This class provides interfaces for instantiating the pretrained model along with
+    the information necessary to retrieve pretrained weights and additional data
+    to be used with the model.
+
+    Torchaudio library instantiates objects of this class, each of which represents
+    a different pretrained model. Client code should access pretrained models via these
+    instances.
+
+    This bundle can estimate objective metric scores for speech enhancement, such as STOI, PESQ, Si-SDR.
+    A typical use case would be a flow like `waveform -> list of scores`. Please see below for the code example.
+
+    Example: Estimate the objective metric scores for the input waveform.
+        >>> import torch
+        >>> import torchaudio
+        >>> from torchaudio.pipelines import SQUIM_OBJECTIVE as bundle
+        >>>
+        >>> # Load the SquimObjective bundle
+        >>> model = bundle.get_model()
+        Downloading: "https://download.pytorch.org/torchaudio/models/squim_objective_dns2020.pth"
+        100%|████████████| 28.2M/28.2M [00:03<00:00, 9.24MB/s]
+        >>>
+        >>> # Resample audio to the expected sampling rate
+        >>> waveform = torchaudio.functional.resample(waveform, sample_rate, bundle.sample_rate)
+        >>>
+        >>> # Estimate objective metric scores
+        >>> scores = model(waveform)
+        >>> print(f"STOI: {scores[0].item()},  PESQ: {scores[1].item()}, SI-SDR: {scores[2].item()}.")
+    """  # noqa: E501
+
+    _path: str
+    _sample_rate: float
+
+    def _get_state_dict(self, dl_kwargs):
+        url = f"https://download.pytorch.org/torchaudio/models/{self._path}"
+        dl_kwargs = {} if dl_kwargs is None else dl_kwargs
+        state_dict = load_state_dict_from_url(url, **dl_kwargs)
+        return state_dict
+
+    def get_model(self, *, dl_kwargs=None) -> SquimObjective:
+        """Construct the SquimObjective model, and load the pretrained weight.
+
+        The weight file is downloaded from the internet and cached with
+        :func:`torch.hub.load_state_dict_from_url`
+
+        Args:
+            dl_kwargs (dictionary of keyword arguments): Passed to :func:`torch.hub.load_state_dict_from_url`.
+
+        Returns:
+            Variation of :py:class:`~torchaudio.models.SquimObjective`.
+        """
+        model = squim_objective_base()
+        model.load_state_dict(self._get_state_dict(dl_kwargs))
+        model.eval()
+        return model
+
+    @property
+    def sample_rate(self):
+        """Sample rate of the audio that the model is trained on.
+
+        :type: float
+        """
+        return self._sample_rate
+
+
+SQUIM_OBJECTIVE = SquimObjectiveBundle(
+    "squim_objective_dns2020.pth",
+    _sample_rate=16000,
+)
+SQUIM_OBJECTIVE.__doc__ = """SquimObjective pipeline trained using approach described in
+    :cite:`kumar2023torchaudio` on the *DNS 2020 Dataset* :cite:`reddy2020interspeech`.
+
+    The underlying model is constructed by :py:func:`torchaudio.models.squim_objective_base`.
+    The weights are under `Creative Commons Attribution 4.0 International License
+    <https://github.com/microsoft/DNS-Challenge/blob/interspeech2020/master/LICENSE>`__.
+
+    Please refer to :py:class:`SquimObjectiveBundle` for usage instructions.
+    """
+
+
+@dataclass
+class SquimSubjectiveBundle:
+    """Data class that bundles associated information to use pretrained
+    :py:class:`~torchaudio.models.SquimSubjective` model.
+
+    This class provides interfaces for instantiating the pretrained model along with
+    the information necessary to retrieve pretrained weights and additional data
+    to be used with the model.
+
+    Torchaudio library instantiates objects of this class, each of which represents
+    a different pretrained model. Client code should access pretrained models via these
+    instances.
+
+    This bundle can estimate subjective metric scores for speech enhancement, such as MOS.
+    A typical use case would be a flow like `waveform -> score`. Please see below for the code example.
+
+    Example: Estimate the subjective metric scores for the input waveform.
+        >>> import torch
+        >>> import torchaudio
+        >>> from torchaudio.pipelines import SQUIM_SUBJECTIVE as bundle
+        >>>
+        >>> # Load the SquimSubjective bundle
+        >>> model = bundle.get_model()
+        Downloading: "https://download.pytorch.org/torchaudio/models/squim_subjective_bvcc_daps.pth"
+        100%|████████████| 360M/360M [00:09<00:00, 41.1MB/s]
+        >>>
+        >>> # Resample audio to the expected sampling rate
+        >>> waveform = torchaudio.functional.resample(waveform, sample_rate, bundle.sample_rate)
+        >>> # Use a clean reference (doesn't need to be the reference for the waveform) as the second input
+        >>> reference = torchaudio.functional.resample(reference, sample_rate, bundle.sample_rate)
+        >>>
+        >>> # Estimate subjective metric scores
+        >>> score = model(waveform, reference)
+        >>> print(f"MOS: {score}.")
+    """  # noqa: E501
+
+    _path: str
+    _sample_rate: float
+
+    def _get_state_dict(self, dl_kwargs):
+        url = f"https://download.pytorch.org/torchaudio/models/{self._path}"
+        dl_kwargs = {} if dl_kwargs is None else dl_kwargs
+        state_dict = load_state_dict_from_url(url, **dl_kwargs)
+        return state_dict
+
+    def get_model(self, *, dl_kwargs=None) -> SquimSubjective:
+        """Construct the SquimSubjective model, and load the pretrained weight.
+
+        The weight file is downloaded from the internet and cached with
+        :func:`torch.hub.load_state_dict_from_url`
+
+        Args:
+            dl_kwargs (dictionary of keyword arguments): Passed to :func:`torch.hub.load_state_dict_from_url`.
+
+        Returns:
+            Variation of :py:class:`~torchaudio.models.SquimObjective`.
+        """
+        model = squim_subjective_base()
+        model.load_state_dict(self._get_state_dict(dl_kwargs))
+        model.eval()
+        return model
+
+    @property
+    def sample_rate(self):
+        """Sample rate of the audio that the model is trained on.
+
+        :type: float
+        """
+        return self._sample_rate
+
+
+SQUIM_SUBJECTIVE = SquimSubjectiveBundle(
+    "squim_subjective_bvcc_daps.pth",
+    _sample_rate=16000,
+)
+SQUIM_SUBJECTIVE.__doc__ = """SquimSubjective pipeline trained
+    as described in :cite:`manocha2022speech` and :cite:`kumar2023torchaudio`
+    on the *BVCC* :cite:`cooper2021voices` and *DAPS* :cite:`mysore2014can` datasets.
+
+    The underlying model is constructed by :py:func:`torchaudio.models.squim_subjective_base`.
+    The weights are under `Creative Commons Attribution Non Commercial 4.0 International
+    <https://zenodo.org/record/4660670#.ZBtWPOxuerN>`__.
+
+    Please refer to :py:class:`SquimSubjectiveBundle` for usage instructions.
+    """

torchaudio/pipelines/_wav2vec2/aligner.py

@@ -0,0 +1,87 @@
+from abc import ABC, abstractmethod
+from typing import Dict, List
+
+import torch
+import torchaudio.functional as F
+from torch import Tensor
+from torchaudio.functional import TokenSpan
+
+
+class ITokenizer(ABC):
+    @abstractmethod
+    def __call__(self, transcript: List[str]) -> List[List[str]]:
+        """Tokenize the given transcript (list of word)
+
+        .. note::
+
+           The toranscript must be normalized.
+
+        Args:
+            transcript (list of str): Transcript (list of word).
+
+        Returns:
+            (list of int): List of token sequences
+        """
+
+
+class Tokenizer(ITokenizer):
+    def __init__(self, dictionary: Dict[str, int]):
+        self.dictionary = dictionary
+
+    def __call__(self, transcript: List[str]) -> List[List[int]]:
+        return [[self.dictionary[c] for c in word] for word in transcript]
+
+
+def _align_emission_and_tokens(emission: Tensor, tokens: List[int], blank: int = 0):
+    device = emission.device
+    emission = emission.unsqueeze(0)
+    targets = torch.tensor([tokens], dtype=torch.int32, device=device)
+
+    aligned_tokens, scores = F.forced_align(emission, targets, blank=blank)
+
+    scores = scores.exp()  # convert back to probability
+    aligned_tokens, scores = aligned_tokens[0], scores[0]  # remove batch dimension
+    return aligned_tokens, scores
+
+
+class IAligner(ABC):
+    @abstractmethod
+    def __call__(self, emission: Tensor, tokens: List[List[int]]) -> List[List[TokenSpan]]:
+        """Generate list of time-stamped token sequences
+
+        Args:
+            emission (Tensor): Sequence of token probability distributions in log-domain.
+                Shape: `(time, tokens)`.
+            tokens (list of integer sequence): Tokenized transcript.
+                Output from :py:class:`torchaudio.pipelines.Wav2Vec2FABundle.Tokenizer`.
+
+        Returns:
+            (list of TokenSpan sequence): Tokens with time stamps and scores.
+        """
+
+
+def _unflatten(list_, lengths):
+    assert len(list_) == sum(lengths)
+    i = 0
+    ret = []
+    for l in lengths:
+        ret.append(list_[i : i + l])
+        i += l
+    return ret
+
+
+def _flatten(nested_list):
+    return [item for list_ in nested_list for item in list_]
+
+
+class Aligner(IAligner):
+    def __init__(self, blank):
+        self.blank = blank
+
+    def __call__(self, emission: Tensor, tokens: List[List[int]]) -> List[List[TokenSpan]]:
+        if emission.ndim != 2:
+            raise ValueError(f"The input emission must be 2D. Found: {emission.shape}")
+
+        aligned_tokens, scores = _align_emission_and_tokens(emission, _flatten(tokens), self.blank)
+        spans = F.merge_tokens(aligned_tokens, scores)
+        return _unflatten(spans, [len(ts) for ts in tokens])