.
+ kernel_size (int): The convolution kernel size of the middle layer, .
+ padding (int): Padding value of the convolution in the middle layer.
+ dilation (int, optional): Dilation value of the convolution in the middle layer.
+ no_redisual (bool, optional): Disable residual block/output.
+
+ Note:
+ This implementation corresponds to the "non-causal" setting in the paper.
+ """
+
+ def __init__(
+ self,
+ io_channels: int,
+ hidden_channels: int,
+ kernel_size: int,
+ padding: int,
+ dilation: int = 1,
+ no_residual: bool = False,
+ ):
+ super().__init__()
+
+ self.conv_layers = torch.nn.Sequential(
+ torch.nn.Conv1d(in_channels=io_channels, out_channels=hidden_channels, kernel_size=1),
+ torch.nn.PReLU(),
+ torch.nn.GroupNorm(num_groups=1, num_channels=hidden_channels, eps=1e-08),
+ torch.nn.Conv1d(
+ in_channels=hidden_channels,
+ out_channels=hidden_channels,
+ kernel_size=kernel_size,
+ padding=padding,
+ dilation=dilation,
+ groups=hidden_channels,
+ ),
+ torch.nn.PReLU(),
+ torch.nn.GroupNorm(num_groups=1, num_channels=hidden_channels, eps=1e-08),
+ )
+
+ self.res_out = (
+ None
+ if no_residual
+ else torch.nn.Conv1d(in_channels=hidden_channels, out_channels=io_channels, kernel_size=1)
+ )
+ self.skip_out = torch.nn.Conv1d(in_channels=hidden_channels, out_channels=io_channels, kernel_size=1)
+
+ def forward(self, input: torch.Tensor) -> Tuple[Optional[torch.Tensor], torch.Tensor]:
+ feature = self.conv_layers(input)
+ if self.res_out is None:
+ residual = None
+ else:
+ residual = self.res_out(feature)
+ skip_out = self.skip_out(feature)
+ return residual, skip_out
+
+
+class MaskGenerator(torch.nn.Module):
+ """TCN (Temporal Convolution Network) Separation Module
+
+ Generates masks for separation.
+
+ Args:
+ input_dim (int): Input feature dimension, .
+ num_sources (int): The number of sources to separate.
+ kernel_size (int): The convolution kernel size of conv blocks, .
+ num_featrs (int): Input/output feature dimenstion of conv blocks, .
+ num_hidden (int): Intermediate feature dimention of conv blocks,
+ num_layers (int): The number of conv blocks in one stack, .
+ num_stacks (int): The number of conv block stacks, .
+ msk_activate (str): The activation function of the mask output.
+
+ Note:
+ This implementation corresponds to the "non-causal" setting in the paper.
+ """
+
+ def __init__(
+ self,
+ input_dim: int,
+ num_sources: int,
+ kernel_size: int,
+ num_feats: int,
+ num_hidden: int,
+ num_layers: int,
+ num_stacks: int,
+ msk_activate: str,
+ ):
+ super().__init__()
+
+ self.input_dim = input_dim
+ self.num_sources = num_sources
+
+ self.input_norm = torch.nn.GroupNorm(num_groups=1, num_channels=input_dim, eps=1e-8)
+ self.input_conv = torch.nn.Conv1d(in_channels=input_dim, out_channels=num_feats, kernel_size=1)
+
+ self.receptive_field = 0
+ self.conv_layers = torch.nn.ModuleList([])
+ for s in range(num_stacks):
+ for l in range(num_layers):
+ multi = 2**l
+ self.conv_layers.append(
+ ConvBlock(
+ io_channels=num_feats,
+ hidden_channels=num_hidden,
+ kernel_size=kernel_size,
+ dilation=multi,
+ padding=multi,
+ # The last ConvBlock does not need residual
+ no_residual=(l == (num_layers - 1) and s == (num_stacks - 1)),
+ )
+ )
+ self.receptive_field += kernel_size if s == 0 and l == 0 else (kernel_size - 1) * multi
+ self.output_prelu = torch.nn.PReLU()
+ self.output_conv = torch.nn.Conv1d(
+ in_channels=num_feats,
+ out_channels=input_dim * num_sources,
+ kernel_size=1,
+ )
+ if msk_activate == "sigmoid":
+ self.mask_activate = torch.nn.Sigmoid()
+ elif msk_activate == "relu":
+ self.mask_activate = torch.nn.ReLU()
+ else:
+ raise ValueError(f"Unsupported activation {msk_activate}")
+
+ def forward(self, input: torch.Tensor) -> torch.Tensor:
+ """Generate separation mask.
+
+ Args:
+ input (torch.Tensor): 3D Tensor with shape [batch, features, frames]
+
+ Returns:
+ Tensor: shape [batch, num_sources, features, frames]
+ """
+ batch_size = input.shape[0]
+ feats = self.input_norm(input)
+ feats = self.input_conv(feats)
+ output = 0.0
+ for layer in self.conv_layers:
+ residual, skip = layer(feats)
+ if residual is not None: # the last conv layer does not produce residual
+ feats = feats + residual
+ output = output + skip
+ output = self.output_prelu(output)
+ output = self.output_conv(output)
+ output = self.mask_activate(output)
+ return output.view(batch_size, self.num_sources, self.input_dim, -1)
+
+
+class ConvTasNet(torch.nn.Module):
+ """Conv-TasNet architecture introduced in
+ *Conv-TasNet: Surpassing Ideal Time–Frequency Magnitude Masking for Speech Separation*
+ :cite:`Luo_2019`.
+
+ Note:
+ This implementation corresponds to the "non-causal" setting in the paper.
+
+ See Also:
+ * :class:`torchaudio.pipelines.SourceSeparationBundle`: Source separation pipeline with pre-trained models.
+
+ Args:
+ num_sources (int, optional): The number of sources to split.
+ enc_kernel_size (int, optional): The convolution kernel size of the encoder/decoder, .
+ enc_num_feats (int, optional): The feature dimensions passed to mask generator, .
+ msk_kernel_size (int, optional): The convolution kernel size of the mask generator, .
+ msk_num_feats (int, optional): The input/output feature dimension of conv block in the mask generator, .
+ msk_num_hidden_feats (int, optional): The internal feature dimension of conv block of the mask generator, .
+ msk_num_layers (int, optional): The number of layers in one conv block of the mask generator, .
+ msk_num_stacks (int, optional): The numbr of conv blocks of the mask generator, .
+ msk_activate (str, optional): The activation function of the mask output (Default: ``sigmoid``).
+ """
+
+ def __init__(
+ self,
+ num_sources: int = 2,
+ # encoder/decoder parameters
+ enc_kernel_size: int = 16,
+ enc_num_feats: int = 512,
+ # mask generator parameters
+ msk_kernel_size: int = 3,
+ msk_num_feats: int = 128,
+ msk_num_hidden_feats: int = 512,
+ msk_num_layers: int = 8,
+ msk_num_stacks: int = 3,
+ msk_activate: str = "sigmoid",
+ ):
+ super().__init__()
+
+ self.num_sources = num_sources
+ self.enc_num_feats = enc_num_feats
+ self.enc_kernel_size = enc_kernel_size
+ self.enc_stride = enc_kernel_size // 2
+
+ self.encoder = torch.nn.Conv1d(
+ in_channels=1,
+ out_channels=enc_num_feats,
+ kernel_size=enc_kernel_size,
+ stride=self.enc_stride,
+ padding=self.enc_stride,
+ bias=False,
+ )
+ self.mask_generator = MaskGenerator(
+ input_dim=enc_num_feats,
+ num_sources=num_sources,
+ kernel_size=msk_kernel_size,
+ num_feats=msk_num_feats,
+ num_hidden=msk_num_hidden_feats,
+ num_layers=msk_num_layers,
+ num_stacks=msk_num_stacks,
+ msk_activate=msk_activate,
+ )
+ self.decoder = torch.nn.ConvTranspose1d(
+ in_channels=enc_num_feats,
+ out_channels=1,
+ kernel_size=enc_kernel_size,
+ stride=self.enc_stride,
+ padding=self.enc_stride,
+ bias=False,
+ )
+
+ def _align_num_frames_with_strides(self, input: torch.Tensor) -> Tuple[torch.Tensor, int]:
+ """Pad input Tensor so that the end of the input tensor corresponds with
+
+ 1. (if kernel size is odd) the center of the last convolution kernel
+ or 2. (if kernel size is even) the end of the first half of the last convolution kernel
+
+ Assumption:
+ The resulting Tensor will be padded with the size of stride (== kernel_width // 2)
+ on the both ends in Conv1D
+
+ |<--- k_1 --->|
+ | | |<-- k_n-1 -->|
+ | | | |<--- k_n --->|
+ | | | | |
+ | | | | |
+ | v v v |
+ |<---->|<--- input signal --->|<--->|<---->|
+ stride PAD stride
+
+ Args:
+ input (torch.Tensor): 3D Tensor with shape (batch_size, channels==1, frames)
+
+ Returns:
+ Tensor: Padded Tensor
+ int: Number of paddings performed
+ """
+ batch_size, num_channels, num_frames = input.shape
+ is_odd = self.enc_kernel_size % 2
+ num_strides = (num_frames - is_odd) // self.enc_stride
+ num_remainings = num_frames - (is_odd + num_strides * self.enc_stride)
+ if num_remainings == 0:
+ return input, 0
+
+ num_paddings = self.enc_stride - num_remainings
+ pad = torch.zeros(
+ batch_size,
+ num_channels,
+ num_paddings,
+ dtype=input.dtype,
+ device=input.device,
+ )
+ return torch.cat([input, pad], 2), num_paddings
+
+ def forward(self, input: torch.Tensor) -> torch.Tensor:
+ """Perform source separation. Generate audio source waveforms.
+
+ Args:
+ input (torch.Tensor): 3D Tensor with shape [batch, channel==1, frames]
+
+ Returns:
+ Tensor: 3D Tensor with shape [batch, channel==num_sources, frames]
+ """
+ if input.ndim != 3 or input.shape[1] != 1:
+ raise ValueError(f"Expected 3D tensor (batch, channel==1, frames). Found: {input.shape}")
+
+ # B: batch size
+ # L: input frame length
+ # L': padded input frame length
+ # F: feature dimension
+ # M: feature frame length
+ # S: number of sources
+
+ padded, num_pads = self._align_num_frames_with_strides(input) # B, 1, L'
+ batch_size, num_padded_frames = padded.shape[0], padded.shape[2]
+ feats = self.encoder(padded) # B, F, M
+ masked = self.mask_generator(feats) * feats.unsqueeze(1) # B, S, F, M
+ masked = masked.view(batch_size * self.num_sources, self.enc_num_feats, -1) # B*S, F, M
+ decoded = self.decoder(masked) # B*S, 1, L'
+ output = decoded.view(batch_size, self.num_sources, num_padded_frames) # B, S, L'
+ if num_pads > 0:
+ output = output[..., :-num_pads] # B, S, L
+ return output
+
+
+def conv_tasnet_base(num_sources: int = 2) -> ConvTasNet:
+ r"""Builds non-causal version of :class:`~torchaudio.models.ConvTasNet`.
+
+ The parameter settings follow the ones with the highest Si-SNR metirc score in the paper,
+ except the mask activation function is changed from "sigmoid" to "relu" for performance improvement.
+
+ Args:
+ num_sources (int, optional): Number of sources in the output.
+ (Default: 2)
+ Returns:
+ ConvTasNet:
+ ConvTasNet model.
+ """
+ return ConvTasNet(
+ num_sources=num_sources,
+ enc_kernel_size=16,
+ enc_num_feats=512,
+ msk_kernel_size=3,
+ msk_num_feats=128,
+ msk_num_hidden_feats=512,
+ msk_num_layers=8,
+ msk_num_stacks=3,
+ msk_activate="relu",
+ )
diff --git a/lib/python3.12/site-packages/torchaudio/models/decoder/__init__.py b/lib/python3.12/site-packages/torchaudio/models/decoder/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d2e9b06d52ef7af302a000bb0f572b4c563e12bd
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/decoder/__init__.py
@@ -0,0 +1,46 @@
+_CTC_DECODERS = [
+ "CTCHypothesis",
+ "CTCDecoder",
+ "CTCDecoderLM",
+ "CTCDecoderLMState",
+ "ctc_decoder",
+ "download_pretrained_files",
+]
+_CUDA_CTC_DECODERS = [
+ "CUCTCDecoder",
+ "CUCTCHypothesis",
+ "cuda_ctc_decoder",
+]
+
+
+def __getattr__(name: str):
+ if name in _CTC_DECODERS:
+ try:
+ from . import _ctc_decoder
+ except Exception as err:
+ raise RuntimeError(
+ "CTC Decoder suit requires flashlight-text package and optionally KenLM. Please install them."
+ ) from err
+
+ item = getattr(_ctc_decoder, name)
+ globals()[name] = item
+ return item
+ elif name in _CUDA_CTC_DECODERS:
+ try:
+ from . import _cuda_ctc_decoder
+ except AttributeError as err:
+ raise RuntimeError(
+ "To use CUCTC decoder, please set BUILD_CUDA_CTC_DECODER=1 when building from source."
+ ) from err
+
+ item = getattr(_cuda_ctc_decoder, name)
+ globals()[name] = item
+ return item
+ raise AttributeError(f"module {__name__} has no attribute {name}")
+
+
+def __dir__():
+ return sorted(__all__)
+
+
+__all__ = _CTC_DECODERS + _CUDA_CTC_DECODERS
diff --git a/lib/python3.12/site-packages/torchaudio/models/decoder/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/decoder/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..855f38d812a19cda6a8b0fecdf11b4bc233ab532
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/decoder/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/decoder/__pycache__/_ctc_decoder.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/decoder/__pycache__/_ctc_decoder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..17c93fd642f6413940fadcc48da3857601d82998
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/decoder/__pycache__/_ctc_decoder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/decoder/__pycache__/_cuda_ctc_decoder.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/decoder/__pycache__/_cuda_ctc_decoder.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4d09ee37c90fa9933af0239c8084246e037a8a17
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/decoder/__pycache__/_cuda_ctc_decoder.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/decoder/_ctc_decoder.py b/lib/python3.12/site-packages/torchaudio/models/decoder/_ctc_decoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..4d45f12f523cc7748e1552ad410557fe9a1f6664
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/decoder/_ctc_decoder.py
@@ -0,0 +1,568 @@
+from __future__ import annotations
+
+import itertools as it
+
+from abc import abstractmethod
+from collections import namedtuple
+from typing import Dict, List, NamedTuple, Optional, Tuple, Union
+
+import torch
+
+from flashlight.lib.text.decoder import (
+ CriterionType as _CriterionType,
+ LexiconDecoder as _LexiconDecoder,
+ LexiconDecoderOptions as _LexiconDecoderOptions,
+ LexiconFreeDecoder as _LexiconFreeDecoder,
+ LexiconFreeDecoderOptions as _LexiconFreeDecoderOptions,
+ LM as _LM,
+ LMState as _LMState,
+ SmearingMode as _SmearingMode,
+ Trie as _Trie,
+ ZeroLM as _ZeroLM,
+)
+from flashlight.lib.text.dictionary import (
+ create_word_dict as _create_word_dict,
+ Dictionary as _Dictionary,
+ load_words as _load_words,
+)
+from torchaudio.utils import download_asset
+
+try:
+ from flashlight.lib.text.decoder.kenlm import KenLM as _KenLM
+except Exception:
+ try:
+ from flashlight.lib.text.decoder import KenLM as _KenLM
+ except Exception:
+ _KenLM = None
+
+__all__ = [
+ "CTCHypothesis",
+ "CTCDecoder",
+ "CTCDecoderLM",
+ "CTCDecoderLMState",
+ "ctc_decoder",
+ "download_pretrained_files",
+]
+
+_PretrainedFiles = namedtuple("PretrainedFiles", ["lexicon", "tokens", "lm"])
+
+
+def _construct_trie(tokens_dict, word_dict, lexicon, lm, silence):
+ vocab_size = tokens_dict.index_size()
+ trie = _Trie(vocab_size, silence)
+ start_state = lm.start(False)
+
+ for word, spellings in lexicon.items():
+ word_idx = word_dict.get_index(word)
+ _, score = lm.score(start_state, word_idx)
+ for spelling in spellings:
+ spelling_idx = [tokens_dict.get_index(token) for token in spelling]
+ trie.insert(spelling_idx, word_idx, score)
+ trie.smear(_SmearingMode.MAX)
+ return trie
+
+
+def _get_word_dict(lexicon, lm, lm_dict, tokens_dict, unk_word):
+ word_dict = None
+ if lm_dict is not None:
+ word_dict = _Dictionary(lm_dict)
+
+ if lexicon and word_dict is None:
+ word_dict = _create_word_dict(lexicon)
+ elif not lexicon and word_dict is None and type(lm) == str:
+ d = {tokens_dict.get_entry(i): [[tokens_dict.get_entry(i)]] for i in range(tokens_dict.index_size())}
+ d[unk_word] = [[unk_word]]
+ word_dict = _create_word_dict(d)
+
+ return word_dict
+
+
+class CTCHypothesis(NamedTuple):
+ r"""Represents hypothesis generated by CTC beam search decoder :class:`CTCDecoder`."""
+ tokens: torch.LongTensor
+ """Predicted sequence of token IDs. Shape `(L, )`, where `L` is the length of the output sequence"""
+
+ words: List[str]
+ """List of predicted words.
+
+ Note:
+ This attribute is only applicable if a lexicon is provided to the decoder. If
+ decoding without a lexicon, it will be blank. Please refer to :attr:`tokens` and
+ :func:`~torchaudio.models.decoder.CTCDecoder.idxs_to_tokens` instead.
+ """
+
+ score: float
+ """Score corresponding to hypothesis"""
+
+ timesteps: torch.IntTensor
+ """Timesteps corresponding to the tokens. Shape `(L, )`, where `L` is the length of the output sequence"""
+
+
+class CTCDecoderLMState(_LMState):
+ """Language model state."""
+
+ @property
+ def children(self) -> Dict[int, CTCDecoderLMState]:
+ """Map of indices to LM states"""
+ return super().children
+
+ def child(self, usr_index: int) -> CTCDecoderLMState:
+ """Returns child corresponding to usr_index, or creates and returns a new state if input index
+ is not found.
+
+ Args:
+ usr_index (int): index corresponding to child state
+
+ Returns:
+ CTCDecoderLMState: child state corresponding to usr_index
+ """
+ return super().child(usr_index)
+
+ def compare(self, state: CTCDecoderLMState) -> CTCDecoderLMState:
+ """Compare two language model states.
+
+ Args:
+ state (CTCDecoderLMState): LM state to compare against
+
+ Returns:
+ int: 0 if the states are the same, -1 if self is less, +1 if self is greater.
+ """
+ pass
+
+
+class CTCDecoderLM(_LM):
+ """Language model base class for creating custom language models to use with the decoder."""
+
+ @abstractmethod
+ def start(self, start_with_nothing: bool) -> CTCDecoderLMState:
+ """Initialize or reset the language model.
+
+ Args:
+ start_with_nothing (bool): whether or not to start sentence with sil token.
+
+ Returns:
+ CTCDecoderLMState: starting state
+ """
+ raise NotImplementedError
+
+ @abstractmethod
+ def score(self, state: CTCDecoderLMState, usr_token_idx: int) -> Tuple[CTCDecoderLMState, float]:
+ """Evaluate the language model based on the current LM state and new word.
+
+ Args:
+ state (CTCDecoderLMState): current LM state
+ usr_token_idx (int): index of the word
+
+ Returns:
+ (CTCDecoderLMState, float)
+ CTCDecoderLMState:
+ new LM state
+ float:
+ score
+ """
+ raise NotImplementedError
+
+ @abstractmethod
+ def finish(self, state: CTCDecoderLMState) -> Tuple[CTCDecoderLMState, float]:
+ """Evaluate end for language model based on current LM state.
+
+ Args:
+ state (CTCDecoderLMState): current LM state
+
+ Returns:
+ (CTCDecoderLMState, float)
+ CTCDecoderLMState:
+ new LM state
+ float:
+ score
+ """
+ raise NotImplementedError
+
+
+class CTCDecoder:
+ """CTC beam search decoder from *Flashlight* :cite:`kahn2022flashlight`.
+
+ .. devices:: CPU
+
+ Note:
+ To build the decoder, please use the factory function :func:`ctc_decoder`.
+ """
+
+ def __init__(
+ self,
+ nbest: int,
+ lexicon: Optional[Dict],
+ word_dict: _Dictionary,
+ tokens_dict: _Dictionary,
+ lm: CTCDecoderLM,
+ decoder_options: Union[_LexiconDecoderOptions, _LexiconFreeDecoderOptions],
+ blank_token: str,
+ sil_token: str,
+ unk_word: str,
+ ) -> None:
+ """
+ Args:
+ nbest (int): number of best decodings to return
+ lexicon (Dict or None): lexicon mapping of words to spellings, or None for lexicon-free decoder
+ word_dict (_Dictionary): dictionary of words
+ tokens_dict (_Dictionary): dictionary of tokens
+ lm (CTCDecoderLM): language model. If using a lexicon, only word level LMs are currently supported
+ decoder_options (_LexiconDecoderOptions or _LexiconFreeDecoderOptions):
+ parameters used for beam search decoding
+ blank_token (str): token corresopnding to blank
+ sil_token (str): token corresponding to silence
+ unk_word (str): word corresponding to unknown
+ """
+
+ self.nbest = nbest
+ self.word_dict = word_dict
+ self.tokens_dict = tokens_dict
+ self.blank = self.tokens_dict.get_index(blank_token)
+ silence = self.tokens_dict.get_index(sil_token)
+ transitions = []
+
+ if lexicon:
+ trie = _construct_trie(tokens_dict, word_dict, lexicon, lm, silence)
+ unk_word = word_dict.get_index(unk_word)
+ token_lm = False # use word level LM
+
+ self.decoder = _LexiconDecoder(
+ decoder_options,
+ trie,
+ lm,
+ silence,
+ self.blank,
+ unk_word,
+ transitions,
+ token_lm,
+ )
+ else:
+ self.decoder = _LexiconFreeDecoder(decoder_options, lm, silence, self.blank, transitions)
+ # https://github.com/pytorch/audio/issues/3218
+ # If lm is passed like rvalue reference, the lm object gets garbage collected,
+ # and later call to the lm fails.
+ # This ensures that lm object is not deleted as long as the decoder is alive.
+ # https://github.com/pybind/pybind11/discussions/4013
+ self.lm = lm
+
+ def _get_tokens(self, idxs: torch.IntTensor) -> torch.LongTensor:
+ idxs = (g[0] for g in it.groupby(idxs))
+ idxs = filter(lambda x: x != self.blank, idxs)
+ return torch.LongTensor(list(idxs))
+
+ def _get_timesteps(self, idxs: torch.IntTensor) -> torch.IntTensor:
+ """Returns frame numbers corresponding to non-blank tokens."""
+
+ timesteps = []
+ for i, idx in enumerate(idxs):
+ if idx == self.blank:
+ continue
+ if i == 0 or idx != idxs[i - 1]:
+ timesteps.append(i)
+ return torch.IntTensor(timesteps)
+
+ def decode_begin(self):
+ """Initialize the internal state of the decoder.
+
+ See :py:meth:`decode_step` for the usage.
+
+ .. note::
+
+ This method is required only when performing online decoding.
+ It is not necessary when performing batch decoding with :py:meth:`__call__`.
+ """
+ self.decoder.decode_begin()
+
+ def decode_end(self):
+ """Finalize the internal state of the decoder.
+
+ See :py:meth:`decode_step` for the usage.
+
+ .. note::
+
+ This method is required only when performing online decoding.
+ It is not necessary when performing batch decoding with :py:meth:`__call__`.
+ """
+ self.decoder.decode_end()
+
+ def decode_step(self, emissions: torch.FloatTensor):
+ """Perform incremental decoding on top of the curent internal state.
+
+ .. note::
+
+ This method is required only when performing online decoding.
+ It is not necessary when performing batch decoding with :py:meth:`__call__`.
+
+ Args:
+ emissions (torch.FloatTensor): CPU tensor of shape `(frame, num_tokens)` storing sequences of
+ probability distribution over labels; output of acoustic model.
+
+ Example:
+ >>> decoder = torchaudio.models.decoder.ctc_decoder(...)
+ >>> decoder.decode_begin()
+ >>> decoder.decode_step(emission1)
+ >>> decoder.decode_step(emission2)
+ >>> decoder.decode_end()
+ >>> result = decoder.get_final_hypothesis()
+ """
+ if emissions.dtype != torch.float32:
+ raise ValueError("emissions must be float32.")
+
+ if not emissions.is_cpu:
+ raise RuntimeError("emissions must be a CPU tensor.")
+
+ if not emissions.is_contiguous():
+ raise RuntimeError("emissions must be contiguous.")
+
+ if emissions.ndim != 2:
+ raise RuntimeError(f"emissions must be 2D. Found {emissions.shape}")
+
+ T, N = emissions.size()
+ self.decoder.decode_step(emissions.data_ptr(), T, N)
+
+ def _to_hypo(self, results) -> List[CTCHypothesis]:
+ return [
+ CTCHypothesis(
+ tokens=self._get_tokens(result.tokens),
+ words=[self.word_dict.get_entry(x) for x in result.words if x >= 0],
+ score=result.score,
+ timesteps=self._get_timesteps(result.tokens),
+ )
+ for result in results
+ ]
+
+ def get_final_hypothesis(self) -> List[CTCHypothesis]:
+ """Get the final hypothesis
+
+ Returns:
+ List[CTCHypothesis]:
+ List of sorted best hypotheses.
+
+ .. note::
+
+ This method is required only when performing online decoding.
+ It is not necessary when performing batch decoding with :py:meth:`__call__`.
+ """
+ results = self.decoder.get_all_final_hypothesis()
+ return self._to_hypo(results[: self.nbest])
+
+ def __call__(
+ self, emissions: torch.FloatTensor, lengths: Optional[torch.Tensor] = None
+ ) -> List[List[CTCHypothesis]]:
+ """
+ Performs batched offline decoding.
+
+ .. note::
+
+ This method performs offline decoding in one go. To perform incremental decoding,
+ please refer to :py:meth:`decode_step`.
+
+ Args:
+ emissions (torch.FloatTensor): CPU tensor of shape `(batch, frame, num_tokens)` storing sequences of
+ probability distribution over labels; output of acoustic model.
+ lengths (Tensor or None, optional): CPU tensor of shape `(batch, )` storing the valid length of
+ in time axis of the output Tensor in each batch.
+
+ Returns:
+ List[List[CTCHypothesis]]:
+ List of sorted best hypotheses for each audio sequence in the batch.
+ """
+
+ if emissions.dtype != torch.float32:
+ raise ValueError("emissions must be float32.")
+
+ if not emissions.is_cpu:
+ raise RuntimeError("emissions must be a CPU tensor.")
+
+ if not emissions.is_contiguous():
+ raise RuntimeError("emissions must be contiguous.")
+
+ if emissions.ndim != 3:
+ raise RuntimeError(f"emissions must be 3D. Found {emissions.shape}")
+
+ if lengths is not None and not lengths.is_cpu:
+ raise RuntimeError("lengths must be a CPU tensor.")
+
+ B, T, N = emissions.size()
+ if lengths is None:
+ lengths = torch.full((B,), T)
+
+ float_bytes = 4
+ hypos = []
+
+ for b in range(B):
+ emissions_ptr = emissions.data_ptr() + float_bytes * b * emissions.stride(0)
+ results = self.decoder.decode(emissions_ptr, lengths[b], N)
+ hypos.append(self._to_hypo(results[: self.nbest]))
+ return hypos
+
+ def idxs_to_tokens(self, idxs: torch.LongTensor) -> List:
+ """
+ Map raw token IDs into corresponding tokens
+
+ Args:
+ idxs (LongTensor): raw token IDs generated from decoder
+
+ Returns:
+ List: tokens corresponding to the input IDs
+ """
+ return [self.tokens_dict.get_entry(idx.item()) for idx in idxs]
+
+
+def ctc_decoder(
+ lexicon: Optional[str],
+ tokens: Union[str, List[str]],
+ lm: Union[str, CTCDecoderLM] = None,
+ lm_dict: Optional[str] = None,
+ nbest: int = 1,
+ beam_size: int = 50,
+ beam_size_token: Optional[int] = None,
+ beam_threshold: float = 50,
+ lm_weight: float = 2,
+ word_score: float = 0,
+ unk_score: float = float("-inf"),
+ sil_score: float = 0,
+ log_add: bool = False,
+ blank_token: str = "-",
+ sil_token: str = "|",
+ unk_word: str = "",
+) -> CTCDecoder:
+ """Builds an instance of :class:`CTCDecoder`.
+
+ Args:
+ lexicon (str or None): lexicon file containing the possible words and corresponding spellings.
+ Each line consists of a word and its space separated spelling. If `None`, uses lexicon-free
+ decoding.
+ tokens (str or List[str]): file or list containing valid tokens. If using a file, the expected
+ format is for tokens mapping to the same index to be on the same line
+ lm (str, CTCDecoderLM, or None, optional): either a path containing KenLM language model,
+ custom language model of type `CTCDecoderLM`, or `None` if not using a language model
+ lm_dict (str or None, optional): file consisting of the dictionary used for the LM, with a word
+ per line sorted by LM index. If decoding with a lexicon, entries in lm_dict must also occur
+ in the lexicon file. If `None`, dictionary for LM is constructed using the lexicon file.
+ (Default: None)
+ nbest (int, optional): number of best decodings to return (Default: 1)
+ beam_size (int, optional): max number of hypos to hold after each decode step (Default: 50)
+ beam_size_token (int, optional): max number of tokens to consider at each decode step.
+ If `None`, it is set to the total number of tokens (Default: None)
+ beam_threshold (float, optional): threshold for pruning hypothesis (Default: 50)
+ lm_weight (float, optional): weight of language model (Default: 2)
+ word_score (float, optional): word insertion score (Default: 0)
+ unk_score (float, optional): unknown word insertion score (Default: -inf)
+ sil_score (float, optional): silence insertion score (Default: 0)
+ log_add (bool, optional): whether or not to use logadd when merging hypotheses (Default: False)
+ blank_token (str, optional): token corresponding to blank (Default: "-")
+ sil_token (str, optional): token corresponding to silence (Default: "|")
+ unk_word (str, optional): word corresponding to unknown (Default: "")
+
+ Returns:
+ CTCDecoder: decoder
+
+ Example
+ >>> decoder = ctc_decoder(
+ >>> lexicon="lexicon.txt",
+ >>> tokens="tokens.txt",
+ >>> lm="kenlm.bin",
+ >>> )
+ >>> results = decoder(emissions) # List of shape (B, nbest) of Hypotheses
+ """
+ if lm_dict is not None and type(lm_dict) is not str:
+ raise ValueError("lm_dict must be None or str type.")
+
+ tokens_dict = _Dictionary(tokens)
+
+ # decoder options
+ if lexicon:
+ lexicon = _load_words(lexicon)
+ decoder_options = _LexiconDecoderOptions(
+ beam_size=beam_size,
+ beam_size_token=beam_size_token or tokens_dict.index_size(),
+ beam_threshold=beam_threshold,
+ lm_weight=lm_weight,
+ word_score=word_score,
+ unk_score=unk_score,
+ sil_score=sil_score,
+ log_add=log_add,
+ criterion_type=_CriterionType.CTC,
+ )
+ else:
+ decoder_options = _LexiconFreeDecoderOptions(
+ beam_size=beam_size,
+ beam_size_token=beam_size_token or tokens_dict.index_size(),
+ beam_threshold=beam_threshold,
+ lm_weight=lm_weight,
+ sil_score=sil_score,
+ log_add=log_add,
+ criterion_type=_CriterionType.CTC,
+ )
+
+ # construct word dict and language model
+ word_dict = _get_word_dict(lexicon, lm, lm_dict, tokens_dict, unk_word)
+
+ if type(lm) == str:
+ if _KenLM is None:
+ raise RuntimeError(
+ "flashlight-text is installed, but KenLM is not installed. "
+ "Please refer to https://github.com/kpu/kenlm#python-module for how to install it."
+ )
+ lm = _KenLM(lm, word_dict)
+ elif lm is None:
+ lm = _ZeroLM()
+
+ return CTCDecoder(
+ nbest=nbest,
+ lexicon=lexicon,
+ word_dict=word_dict,
+ tokens_dict=tokens_dict,
+ lm=lm,
+ decoder_options=decoder_options,
+ blank_token=blank_token,
+ sil_token=sil_token,
+ unk_word=unk_word,
+ )
+
+
+def _get_filenames(model: str) -> _PretrainedFiles:
+ if model not in ["librispeech", "librispeech-3-gram", "librispeech-4-gram"]:
+ raise ValueError(
+ f"{model} not supported. Must be one of ['librispeech-3-gram', 'librispeech-4-gram', 'librispeech']"
+ )
+
+ prefix = f"decoder-assets/{model}"
+ return _PretrainedFiles(
+ lexicon=f"{prefix}/lexicon.txt",
+ tokens=f"{prefix}/tokens.txt",
+ lm=f"{prefix}/lm.bin" if model != "librispeech" else None,
+ )
+
+
+def download_pretrained_files(model: str) -> _PretrainedFiles:
+ """
+ Retrieves pretrained data files used for :func:`ctc_decoder`.
+
+ Args:
+ model (str): pretrained language model to download.
+ Valid values are: ``"librispeech-3-gram"``, ``"librispeech-4-gram"`` and ``"librispeech"``.
+
+ Returns:
+ Object with the following attributes
+
+ * ``lm``: path corresponding to downloaded language model,
+ or ``None`` if the model is not associated with an lm
+ * ``lexicon``: path corresponding to downloaded lexicon file
+ * ``tokens``: path corresponding to downloaded tokens file
+ """
+
+ files = _get_filenames(model)
+ lexicon_file = download_asset(files.lexicon)
+ tokens_file = download_asset(files.tokens)
+ if files.lm is not None:
+ lm_file = download_asset(files.lm)
+ else:
+ lm_file = None
+
+ return _PretrainedFiles(
+ lexicon=lexicon_file,
+ tokens=tokens_file,
+ lm=lm_file,
+ )
diff --git a/lib/python3.12/site-packages/torchaudio/models/decoder/_cuda_ctc_decoder.py b/lib/python3.12/site-packages/torchaudio/models/decoder/_cuda_ctc_decoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..9b1f509644091e04ea3bdc4301a74c546044f31d
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/decoder/_cuda_ctc_decoder.py
@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+import math
+
+from typing import List, NamedTuple, Union
+
+import torch
+import torchaudio
+
+torchaudio._extension._load_lib("libctc_prefix_decoder")
+import torchaudio.lib.pybind11_prefixctc as cuctc
+
+
+__all__ = ["CUCTCHypothesis", "CUCTCDecoder", "cuda_ctc_decoder"]
+
+
+def _get_vocab_list(vocab_file):
+ vocab = []
+ with open(vocab_file, "r", encoding="utf-8") as f:
+ for line in f:
+ line = line.strip().split()
+ vocab.append(line[0])
+ return vocab
+
+
+class CUCTCHypothesis(NamedTuple):
+ r"""Represents hypothesis generated by CUCTC beam search decoder :class:`CUCTCDecoder`."""
+ tokens: List[int]
+ """Predicted sequence of token IDs. Shape `(L, )`, where `L` is the length of the output sequence"""
+
+ words: List[str]
+ """List of predicted tokens. Algin with modeling unit.
+ """
+
+ score: float
+ """Score corresponding to hypothesis"""
+
+
+_DEFAULT_BLANK_SKIP_THREASHOLD = 0.95
+
+
+class CUCTCDecoder:
+ """CUDA CTC beam search decoder.
+
+ .. devices:: CUDA
+
+ Note:
+ To build the decoder, please use the factory function :func:`cuda_ctc_decoder`.
+ """
+
+ def __init__(
+ self,
+ vocab_list: List[str],
+ blank_id: int = 0,
+ beam_size: int = 10,
+ nbest: int = 1,
+ blank_skip_threshold: float = _DEFAULT_BLANK_SKIP_THREASHOLD,
+ cuda_stream: torch.cuda.streams.Stream = None,
+ ):
+ """
+ Args:
+ blank_id (int): token id corresopnding to blank, only support 0 for now. (Default: 0)
+ vocab_list (List[str]): list of vocabulary tokens
+ beam_size (int, optional): max number of hypos to hold after each decode step (Default: 10)
+ nbest (int): number of best decodings to return
+ blank_skip_threshold (float):
+ skip frames if log_prob(blank) > log(blank_skip_threshold), to speed up decoding.
+ (Default: 0.95).
+ cuda_stream (torch.cuda.streams.Stream): using assigned cuda stream (Default: using default stream)
+
+ """
+ if cuda_stream:
+ if not isinstance(cuda_stream, torch.cuda.streams.Stream):
+ raise AssertionError("cuda_stream must be torch.cuda.streams.Stream")
+ cuda_stream_ = cuda_stream.cuda_stream if cuda_stream else torch.cuda.current_stream().cuda_stream
+ self.internal_data = cuctc.prefixCTC_alloc(cuda_stream_)
+ self.memory = torch.empty(0, dtype=torch.int8, device=torch.device("cuda"))
+ if blank_id != 0:
+ raise AssertionError("blank_id must be 0")
+ self.blank_id = blank_id
+ self.vocab_list = vocab_list
+ self.space_id = 0
+ self.nbest = nbest
+ if not (blank_skip_threshold >= 0 and blank_skip_threshold <= 1):
+ raise AssertionError("blank_skip_threshold must be between 0 and 1")
+ self.blank_skip_threshold = math.log(blank_skip_threshold)
+ self.beam_size = min(beam_size, len(vocab_list)) # beam size must be smaller than vocab size
+
+ def __del__(self):
+ if cuctc is not None:
+ cuctc.prefixCTC_free(self.internal_data)
+
+ def __call__(self, log_prob: torch.Tensor, encoder_out_lens: torch.Tensor):
+ """
+ Args:
+ log_prob (torch.FloatTensor): GPU tensor of shape `(batch, frame, num_tokens)` storing sequences of
+ probability distribution over labels; log_softmax(output of acoustic model).
+ lengths (dtype torch.int32): GPU tensor of shape `(batch, )` storing the valid length of
+ in time axis of the output Tensor in each batch.
+
+ Returns:
+ List[List[CUCTCHypothesis]]:
+ List of sorted best hypotheses for each audio sequence in the batch.
+ """
+ if not encoder_out_lens.dtype == torch.int32:
+ raise AssertionError("encoder_out_lens must be torch.int32")
+ if not log_prob.dtype == torch.float32:
+ raise AssertionError("log_prob must be torch.float32")
+ if not (log_prob.is_cuda and encoder_out_lens.is_cuda):
+ raise AssertionError("inputs must be cuda tensors")
+ if not (log_prob.is_contiguous() and encoder_out_lens.is_contiguous()):
+ raise AssertionError("input tensors must be contiguous")
+ required_size, score_hyps = cuctc.ctc_beam_search_decoder_batch_gpu_v2(
+ self.internal_data,
+ self.memory.data_ptr(),
+ self.memory.size(0),
+ log_prob.data_ptr(),
+ encoder_out_lens.data_ptr(),
+ log_prob.size(),
+ log_prob.stride(),
+ self.beam_size,
+ self.blank_id,
+ self.space_id,
+ self.blank_skip_threshold,
+ )
+ if required_size > 0:
+ self.memory = torch.empty(required_size, dtype=torch.int8, device=log_prob.device).contiguous()
+ _, score_hyps = cuctc.ctc_beam_search_decoder_batch_gpu_v2(
+ self.internal_data,
+ self.memory.data_ptr(),
+ self.memory.size(0),
+ log_prob.data_ptr(),
+ encoder_out_lens.data_ptr(),
+ log_prob.size(),
+ log_prob.stride(),
+ self.beam_size,
+ self.blank_id,
+ self.space_id,
+ self.blank_skip_threshold,
+ )
+ batch_size = len(score_hyps)
+ hypos = []
+ for i in range(batch_size):
+ hypos.append(
+ [
+ CUCTCHypothesis(
+ tokens=score_hyps[i][j][1],
+ words=[self.vocab_list[word_id] for word_id in score_hyps[i][j][1]],
+ score=score_hyps[i][j][0],
+ )
+ for j in range(self.nbest)
+ ]
+ )
+ return hypos
+
+
+def cuda_ctc_decoder(
+ tokens: Union[str, List[str]],
+ nbest: int = 1,
+ beam_size: int = 10,
+ blank_skip_threshold: float = _DEFAULT_BLANK_SKIP_THREASHOLD,
+) -> CUCTCDecoder:
+ """Builds an instance of :class:`CUCTCDecoder`.
+
+ Args:
+ tokens (str or List[str]): File or list containing valid tokens.
+ If using a file, the expected format is for tokens mapping to the same index to be on the same line
+ beam_size (int, optional): The maximum number of hypos to hold after each decode step (Default: 10)
+ nbest (int): The number of best decodings to return
+ blank_id (int): The token ID corresopnding to the blank symbol.
+ blank_skip_threshold (float): skip frames if log_prob(blank) > log(blank_skip_threshold), to speed up decoding
+ (Default: 0.95).
+
+ Returns:
+ CUCTCDecoder: decoder
+
+ Example
+ >>> decoder = cuda_ctc_decoder(
+ >>> vocab_file="tokens.txt",
+ >>> blank_skip_threshold=0.95,
+ >>> )
+ >>> results = decoder(log_probs, encoder_out_lens) # List of shape (B, nbest) of Hypotheses
+ """
+ if type(tokens) == str:
+ tokens = _get_vocab_list(tokens)
+
+ return CUCTCDecoder(vocab_list=tokens, beam_size=beam_size, nbest=nbest, blank_skip_threshold=blank_skip_threshold)
diff --git a/lib/python3.12/site-packages/torchaudio/models/deepspeech.py b/lib/python3.12/site-packages/torchaudio/models/deepspeech.py
new file mode 100644
index 0000000000000000000000000000000000000000..ef23d1d351bde615cb2b1b38ffdd7782fbb5b627
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/deepspeech.py
@@ -0,0 +1,84 @@
+import torch
+
+__all__ = ["DeepSpeech"]
+
+
+class FullyConnected(torch.nn.Module):
+ """
+ Args:
+ n_feature: Number of input features
+ n_hidden: Internal hidden unit size.
+ """
+
+ def __init__(self, n_feature: int, n_hidden: int, dropout: float, relu_max_clip: int = 20) -> None:
+ super(FullyConnected, self).__init__()
+ self.fc = torch.nn.Linear(n_feature, n_hidden, bias=True)
+ self.relu_max_clip = relu_max_clip
+ self.dropout = dropout
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.fc(x)
+ x = torch.nn.functional.relu(x)
+ x = torch.nn.functional.hardtanh(x, 0, self.relu_max_clip)
+ if self.dropout:
+ x = torch.nn.functional.dropout(x, self.dropout, self.training)
+ return x
+
+
+class DeepSpeech(torch.nn.Module):
+ """DeepSpeech architecture introduced in
+ *Deep Speech: Scaling up end-to-end speech recognition* :cite:`hannun2014deep`.
+
+ Args:
+ n_feature: Number of input features
+ n_hidden: Internal hidden unit size.
+ n_class: Number of output classes
+ """
+
+ def __init__(
+ self,
+ n_feature: int,
+ n_hidden: int = 2048,
+ n_class: int = 40,
+ dropout: float = 0.0,
+ ) -> None:
+ super(DeepSpeech, self).__init__()
+ self.n_hidden = n_hidden
+ self.fc1 = FullyConnected(n_feature, n_hidden, dropout)
+ self.fc2 = FullyConnected(n_hidden, n_hidden, dropout)
+ self.fc3 = FullyConnected(n_hidden, n_hidden, dropout)
+ self.bi_rnn = torch.nn.RNN(n_hidden, n_hidden, num_layers=1, nonlinearity="relu", bidirectional=True)
+ self.fc4 = FullyConnected(n_hidden, n_hidden, dropout)
+ self.out = torch.nn.Linear(n_hidden, n_class)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Args:
+ x (torch.Tensor): Tensor of dimension (batch, channel, time, feature).
+ Returns:
+ Tensor: Predictor tensor of dimension (batch, time, class).
+ """
+ # N x C x T x F
+ x = self.fc1(x)
+ # N x C x T x H
+ x = self.fc2(x)
+ # N x C x T x H
+ x = self.fc3(x)
+ # N x C x T x H
+ x = x.squeeze(1)
+ # N x T x H
+ x = x.transpose(0, 1)
+ # T x N x H
+ x, _ = self.bi_rnn(x)
+ # The fifth (non-recurrent) layer takes both the forward and backward units as inputs
+ x = x[:, :, : self.n_hidden] + x[:, :, self.n_hidden :]
+ # T x N x H
+ x = self.fc4(x)
+ # T x N x H
+ x = self.out(x)
+ # T x N x n_class
+ x = x.permute(1, 0, 2)
+ # N x T x n_class
+ x = torch.nn.functional.log_softmax(x, dim=2)
+ # N x T x n_class
+ return x
diff --git a/lib/python3.12/site-packages/torchaudio/models/emformer.py b/lib/python3.12/site-packages/torchaudio/models/emformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..9ddd257552ecda94cb55bbc1eed1dae8a5382380
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/emformer.py
@@ -0,0 +1,884 @@
+import math
+from typing import List, Optional, Tuple
+
+import torch
+
+
+__all__ = ["Emformer"]
+
+
+def _lengths_to_padding_mask(lengths: torch.Tensor) -> torch.Tensor:
+ batch_size = lengths.shape[0]
+ max_length = int(torch.max(lengths).item())
+ padding_mask = torch.arange(max_length, device=lengths.device, dtype=lengths.dtype).expand(
+ batch_size, max_length
+ ) >= lengths.unsqueeze(1)
+ return padding_mask
+
+
+def _gen_padding_mask(
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ summary: torch.Tensor,
+ lengths: torch.Tensor,
+ mems: torch.Tensor,
+ left_context_key: Optional[torch.Tensor] = None,
+) -> Optional[torch.Tensor]:
+ T = right_context.size(0) + utterance.size(0) + summary.size(0)
+ B = right_context.size(1)
+ if B == 1:
+ padding_mask = None
+ else:
+ right_context_blocks_length = T - torch.max(lengths).int() - summary.size(0)
+ left_context_blocks_length = left_context_key.size(0) if left_context_key is not None else 0
+ klengths = lengths + mems.size(0) + right_context_blocks_length + left_context_blocks_length
+ padding_mask = _lengths_to_padding_mask(lengths=klengths)
+ return padding_mask
+
+
+def _get_activation_module(activation: str) -> torch.nn.Module:
+ if activation == "relu":
+ return torch.nn.ReLU()
+ elif activation == "gelu":
+ return torch.nn.GELU()
+ elif activation == "silu":
+ return torch.nn.SiLU()
+ else:
+ raise ValueError(f"Unsupported activation {activation}")
+
+
+def _get_weight_init_gains(weight_init_scale_strategy: Optional[str], num_layers: int) -> List[Optional[float]]:
+ if weight_init_scale_strategy is None:
+ return [None for _ in range(num_layers)]
+ elif weight_init_scale_strategy == "depthwise":
+ return [1.0 / math.sqrt(layer_idx + 1) for layer_idx in range(num_layers)]
+ elif weight_init_scale_strategy == "constant":
+ return [1.0 / math.sqrt(2) for layer_idx in range(num_layers)]
+ else:
+ raise ValueError(f"Unsupported weight_init_scale_strategy value {weight_init_scale_strategy}")
+
+
+def _gen_attention_mask_block(
+ col_widths: List[int], col_mask: List[bool], num_rows: int, device: torch.device
+) -> torch.Tensor:
+ if len(col_widths) != len(col_mask):
+ raise ValueError("Length of col_widths must match that of col_mask")
+
+ mask_block = [
+ torch.ones(num_rows, col_width, device=device)
+ if is_ones_col
+ else torch.zeros(num_rows, col_width, device=device)
+ for col_width, is_ones_col in zip(col_widths, col_mask)
+ ]
+ return torch.cat(mask_block, dim=1)
+
+
+class _EmformerAttention(torch.nn.Module):
+ r"""Emformer layer attention module.
+
+ Args:
+ input_dim (int): input dimension.
+ num_heads (int): number of attention heads in each Emformer layer.
+ dropout (float, optional): dropout probability. (Default: 0.0)
+ weight_init_gain (float or None, optional): scale factor to apply when initializing
+ attention module parameters. (Default: ``None``)
+ tanh_on_mem (bool, optional): if ``True``, applies tanh to memory elements. (Default: ``False``)
+ negative_inf (float, optional): value to use for negative infinity in attention weights. (Default: -1e8)
+ """
+
+ def __init__(
+ self,
+ input_dim: int,
+ num_heads: int,
+ dropout: float = 0.0,
+ weight_init_gain: Optional[float] = None,
+ tanh_on_mem: bool = False,
+ negative_inf: float = -1e8,
+ ):
+ super().__init__()
+
+ if input_dim % num_heads != 0:
+ raise ValueError(f"input_dim ({input_dim}) is not a multiple of num_heads ({num_heads}).")
+
+ self.input_dim = input_dim
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.tanh_on_mem = tanh_on_mem
+ self.negative_inf = negative_inf
+
+ self.scaling = (self.input_dim // self.num_heads) ** -0.5
+
+ self.emb_to_key_value = torch.nn.Linear(input_dim, 2 * input_dim, bias=True)
+ self.emb_to_query = torch.nn.Linear(input_dim, input_dim, bias=True)
+ self.out_proj = torch.nn.Linear(input_dim, input_dim, bias=True)
+
+ if weight_init_gain:
+ torch.nn.init.xavier_uniform_(self.emb_to_key_value.weight, gain=weight_init_gain)
+ torch.nn.init.xavier_uniform_(self.emb_to_query.weight, gain=weight_init_gain)
+
+ def _gen_key_value(self, input: torch.Tensor, mems: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ T, _, _ = input.shape
+ summary_length = mems.size(0) + 1
+ right_ctx_utterance_block = input[: T - summary_length]
+ mems_right_ctx_utterance_block = torch.cat([mems, right_ctx_utterance_block])
+ key, value = self.emb_to_key_value(mems_right_ctx_utterance_block).chunk(chunks=2, dim=2)
+ return key, value
+
+ def _gen_attention_probs(
+ self,
+ attention_weights: torch.Tensor,
+ attention_mask: torch.Tensor,
+ padding_mask: Optional[torch.Tensor],
+ ) -> torch.Tensor:
+ attention_weights_float = attention_weights.float()
+ attention_weights_float = attention_weights_float.masked_fill(attention_mask.unsqueeze(0), self.negative_inf)
+ T = attention_weights.size(1)
+ B = attention_weights.size(0) // self.num_heads
+ if padding_mask is not None:
+ attention_weights_float = attention_weights_float.view(B, self.num_heads, T, -1)
+ attention_weights_float = attention_weights_float.masked_fill(
+ padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), self.negative_inf
+ )
+ attention_weights_float = attention_weights_float.view(B * self.num_heads, T, -1)
+ attention_probs = torch.nn.functional.softmax(attention_weights_float, dim=-1).type_as(attention_weights)
+ return torch.nn.functional.dropout(attention_probs, p=float(self.dropout), training=self.training)
+
+ def _forward_impl(
+ self,
+ utterance: torch.Tensor,
+ lengths: torch.Tensor,
+ right_context: torch.Tensor,
+ summary: torch.Tensor,
+ mems: torch.Tensor,
+ attention_mask: torch.Tensor,
+ left_context_key: Optional[torch.Tensor] = None,
+ left_context_val: Optional[torch.Tensor] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ B = utterance.size(1)
+ T = right_context.size(0) + utterance.size(0) + summary.size(0)
+
+ # Compute query with [right context, utterance, summary].
+ query = self.emb_to_query(torch.cat([right_context, utterance, summary]))
+
+ # Compute key and value with [mems, right context, utterance].
+ key, value = self.emb_to_key_value(torch.cat([mems, right_context, utterance])).chunk(chunks=2, dim=2)
+
+ if left_context_key is not None and left_context_val is not None:
+ right_context_blocks_length = T - torch.max(lengths).int() - summary.size(0)
+ key = torch.cat(
+ [
+ key[: mems.size(0) + right_context_blocks_length],
+ left_context_key,
+ key[mems.size(0) + right_context_blocks_length :],
+ ],
+ )
+ value = torch.cat(
+ [
+ value[: mems.size(0) + right_context_blocks_length],
+ left_context_val,
+ value[mems.size(0) + right_context_blocks_length :],
+ ],
+ )
+
+ # Compute attention weights from query, key, and value.
+ reshaped_query, reshaped_key, reshaped_value = [
+ tensor.contiguous().view(-1, B * self.num_heads, self.input_dim // self.num_heads).transpose(0, 1)
+ for tensor in [query, key, value]
+ ]
+ attention_weights = torch.bmm(reshaped_query * self.scaling, reshaped_key.transpose(1, 2))
+
+ # Compute padding mask.
+ padding_mask = _gen_padding_mask(utterance, right_context, summary, lengths, mems, left_context_key)
+
+ # Compute attention probabilities.
+ attention_probs = self._gen_attention_probs(attention_weights, attention_mask, padding_mask)
+
+ # Compute attention.
+ attention = torch.bmm(attention_probs, reshaped_value)
+ if attention.shape != (
+ B * self.num_heads,
+ T,
+ self.input_dim // self.num_heads,
+ ):
+ raise AssertionError("Computed attention has incorrect dimensions")
+ attention = attention.transpose(0, 1).contiguous().view(T, B, self.input_dim)
+
+ # Apply output projection.
+ output_right_context_mems = self.out_proj(attention)
+
+ summary_length = summary.size(0)
+ output_right_context = output_right_context_mems[: T - summary_length]
+ output_mems = output_right_context_mems[T - summary_length :]
+ if self.tanh_on_mem:
+ output_mems = torch.tanh(output_mems)
+ else:
+ output_mems = torch.clamp(output_mems, min=-10, max=10)
+
+ return output_right_context, output_mems, key, value
+
+ def forward(
+ self,
+ utterance: torch.Tensor,
+ lengths: torch.Tensor,
+ right_context: torch.Tensor,
+ summary: torch.Tensor,
+ mems: torch.Tensor,
+ attention_mask: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ r"""Forward pass for training.
+
+ B: batch size;
+ D: feature dimension of each frame;
+ T: number of utterance frames;
+ R: number of right context frames;
+ S: number of summary elements;
+ M: number of memory elements.
+
+ Args:
+ utterance (torch.Tensor): utterance frames, with shape `(T, B, D)`.
+ lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``utterance``.
+ right_context (torch.Tensor): right context frames, with shape `(R, B, D)`.
+ summary (torch.Tensor): summary elements, with shape `(S, B, D)`.
+ mems (torch.Tensor): memory elements, with shape `(M, B, D)`.
+ attention_mask (torch.Tensor): attention mask for underlying attention module.
+
+ Returns:
+ (Tensor, Tensor):
+ Tensor
+ output frames corresponding to utterance and right_context, with shape `(T + R, B, D)`.
+ Tensor
+ updated memory elements, with shape `(M, B, D)`.
+ """
+ output, output_mems, _, _ = self._forward_impl(utterance, lengths, right_context, summary, mems, attention_mask)
+ return output, output_mems[:-1]
+
+ @torch.jit.export
+ def infer(
+ self,
+ utterance: torch.Tensor,
+ lengths: torch.Tensor,
+ right_context: torch.Tensor,
+ summary: torch.Tensor,
+ mems: torch.Tensor,
+ left_context_key: torch.Tensor,
+ left_context_val: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ r"""Forward pass for inference.
+
+ B: batch size;
+ D: feature dimension of each frame;
+ T: number of utterance frames;
+ R: number of right context frames;
+ S: number of summary elements;
+ M: number of memory elements.
+
+ Args:
+ utterance (torch.Tensor): utterance frames, with shape `(T, B, D)`.
+ lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``utterance``.
+ right_context (torch.Tensor): right context frames, with shape `(R, B, D)`.
+ summary (torch.Tensor): summary elements, with shape `(S, B, D)`.
+ mems (torch.Tensor): memory elements, with shape `(M, B, D)`.
+ left_context_key (torch.Tensor): left context attention key computed from preceding invocation.
+ left_context_val (torch.Tensor): left context attention value computed from preceding invocation.
+
+ Returns:
+ (Tensor, Tensor, Tensor, and Tensor):
+ Tensor
+ output frames corresponding to utterance and right_context, with shape `(T + R, B, D)`.
+ Tensor
+ updated memory elements, with shape `(M, B, D)`.
+ Tensor
+ attention key computed for left context and utterance.
+ Tensor
+ attention value computed for left context and utterance.
+ """
+ query_dim = right_context.size(0) + utterance.size(0) + summary.size(0)
+ key_dim = right_context.size(0) + utterance.size(0) + mems.size(0) + left_context_key.size(0)
+ attention_mask = torch.zeros(query_dim, key_dim).to(dtype=torch.bool, device=utterance.device)
+ attention_mask[-1, : mems.size(0)] = True
+ output, output_mems, key, value = self._forward_impl(
+ utterance,
+ lengths,
+ right_context,
+ summary,
+ mems,
+ attention_mask,
+ left_context_key=left_context_key,
+ left_context_val=left_context_val,
+ )
+ return (
+ output,
+ output_mems,
+ key[mems.size(0) + right_context.size(0) :],
+ value[mems.size(0) + right_context.size(0) :],
+ )
+
+
+class _EmformerLayer(torch.nn.Module):
+ r"""Emformer layer that constitutes Emformer.
+
+ Args:
+ input_dim (int): input dimension.
+ num_heads (int): number of attention heads.
+ ffn_dim: (int): hidden layer dimension of feedforward network.
+ segment_length (int): length of each input segment.
+ dropout (float, optional): dropout probability. (Default: 0.0)
+ activation (str, optional): activation function to use in feedforward network.
+ Must be one of ("relu", "gelu", "silu"). (Default: "relu")
+ left_context_length (int, optional): length of left context. (Default: 0)
+ max_memory_size (int, optional): maximum number of memory elements to use. (Default: 0)
+ weight_init_gain (float or None, optional): scale factor to apply when initializing
+ attention module parameters. (Default: ``None``)
+ tanh_on_mem (bool, optional): if ``True``, applies tanh to memory elements. (Default: ``False``)
+ negative_inf (float, optional): value to use for negative infinity in attention weights. (Default: -1e8)
+ """
+
+ def __init__(
+ self,
+ input_dim: int,
+ num_heads: int,
+ ffn_dim: int,
+ segment_length: int,
+ dropout: float = 0.0,
+ activation: str = "relu",
+ left_context_length: int = 0,
+ max_memory_size: int = 0,
+ weight_init_gain: Optional[float] = None,
+ tanh_on_mem: bool = False,
+ negative_inf: float = -1e8,
+ ):
+ super().__init__()
+
+ self.attention = _EmformerAttention(
+ input_dim=input_dim,
+ num_heads=num_heads,
+ dropout=dropout,
+ weight_init_gain=weight_init_gain,
+ tanh_on_mem=tanh_on_mem,
+ negative_inf=negative_inf,
+ )
+ self.dropout = torch.nn.Dropout(dropout)
+ self.memory_op = torch.nn.AvgPool1d(kernel_size=segment_length, stride=segment_length, ceil_mode=True)
+
+ activation_module = _get_activation_module(activation)
+ self.pos_ff = torch.nn.Sequential(
+ torch.nn.LayerNorm(input_dim),
+ torch.nn.Linear(input_dim, ffn_dim),
+ activation_module,
+ torch.nn.Dropout(dropout),
+ torch.nn.Linear(ffn_dim, input_dim),
+ torch.nn.Dropout(dropout),
+ )
+ self.layer_norm_input = torch.nn.LayerNorm(input_dim)
+ self.layer_norm_output = torch.nn.LayerNorm(input_dim)
+
+ self.left_context_length = left_context_length
+ self.segment_length = segment_length
+ self.max_memory_size = max_memory_size
+ self.input_dim = input_dim
+
+ self.use_mem = max_memory_size > 0
+
+ def _init_state(self, batch_size: int, device: Optional[torch.device]) -> List[torch.Tensor]:
+ empty_memory = torch.zeros(self.max_memory_size, batch_size, self.input_dim, device=device)
+ left_context_key = torch.zeros(self.left_context_length, batch_size, self.input_dim, device=device)
+ left_context_val = torch.zeros(self.left_context_length, batch_size, self.input_dim, device=device)
+ past_length = torch.zeros(1, batch_size, dtype=torch.int32, device=device)
+ return [empty_memory, left_context_key, left_context_val, past_length]
+
+ def _unpack_state(self, state: List[torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ past_length = state[3][0][0].item()
+ past_left_context_length = min(self.left_context_length, past_length)
+ past_mem_length = min(self.max_memory_size, math.ceil(past_length / self.segment_length))
+ pre_mems = state[0][self.max_memory_size - past_mem_length :]
+ lc_key = state[1][self.left_context_length - past_left_context_length :]
+ lc_val = state[2][self.left_context_length - past_left_context_length :]
+ return pre_mems, lc_key, lc_val
+
+ def _pack_state(
+ self,
+ next_k: torch.Tensor,
+ next_v: torch.Tensor,
+ update_length: int,
+ mems: torch.Tensor,
+ state: List[torch.Tensor],
+ ) -> List[torch.Tensor]:
+ new_k = torch.cat([state[1], next_k])
+ new_v = torch.cat([state[2], next_v])
+ state[0] = torch.cat([state[0], mems])[-self.max_memory_size :]
+ state[1] = new_k[new_k.shape[0] - self.left_context_length :]
+ state[2] = new_v[new_v.shape[0] - self.left_context_length :]
+ state[3] = state[3] + update_length
+ return state
+
+ def _process_attention_output(
+ self,
+ rc_output: torch.Tensor,
+ utterance: torch.Tensor,
+ right_context: torch.Tensor,
+ ) -> torch.Tensor:
+ result = self.dropout(rc_output) + torch.cat([right_context, utterance])
+ result = self.pos_ff(result) + result
+ result = self.layer_norm_output(result)
+ return result
+
+ def _apply_pre_attention_layer_norm(
+ self, utterance: torch.Tensor, right_context: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ layer_norm_input = self.layer_norm_input(torch.cat([right_context, utterance]))
+ return (
+ layer_norm_input[right_context.size(0) :],
+ layer_norm_input[: right_context.size(0)],
+ )
+
+ def _apply_post_attention_ffn(
+ self, rc_output: torch.Tensor, utterance: torch.Tensor, right_context: torch.Tensor
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ rc_output = self._process_attention_output(rc_output, utterance, right_context)
+ return rc_output[right_context.size(0) :], rc_output[: right_context.size(0)]
+
+ def _apply_attention_forward(
+ self,
+ utterance: torch.Tensor,
+ lengths: torch.Tensor,
+ right_context: torch.Tensor,
+ mems: torch.Tensor,
+ attention_mask: Optional[torch.Tensor],
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ if attention_mask is None:
+ raise ValueError("attention_mask must be not None when for_inference is False")
+
+ if self.use_mem:
+ summary = self.memory_op(utterance.permute(1, 2, 0)).permute(2, 0, 1)
+ else:
+ summary = torch.empty(0).to(dtype=utterance.dtype, device=utterance.device)
+ rc_output, next_m = self.attention(
+ utterance=utterance,
+ lengths=lengths,
+ right_context=right_context,
+ summary=summary,
+ mems=mems,
+ attention_mask=attention_mask,
+ )
+ return rc_output, next_m
+
+ def _apply_attention_infer(
+ self,
+ utterance: torch.Tensor,
+ lengths: torch.Tensor,
+ right_context: torch.Tensor,
+ mems: torch.Tensor,
+ state: Optional[List[torch.Tensor]],
+ ) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor]]:
+ if state is None:
+ state = self._init_state(utterance.size(1), device=utterance.device)
+ pre_mems, lc_key, lc_val = self._unpack_state(state)
+ if self.use_mem:
+ summary = self.memory_op(utterance.permute(1, 2, 0)).permute(2, 0, 1)
+ summary = summary[:1]
+ else:
+ summary = torch.empty(0).to(dtype=utterance.dtype, device=utterance.device)
+ rc_output, next_m, next_k, next_v = self.attention.infer(
+ utterance=utterance,
+ lengths=lengths,
+ right_context=right_context,
+ summary=summary,
+ mems=pre_mems,
+ left_context_key=lc_key,
+ left_context_val=lc_val,
+ )
+ state = self._pack_state(next_k, next_v, utterance.size(0), mems, state)
+ return rc_output, next_m, state
+
+ def forward(
+ self,
+ utterance: torch.Tensor,
+ lengths: torch.Tensor,
+ right_context: torch.Tensor,
+ mems: torch.Tensor,
+ attention_mask: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ r"""Forward pass for training.
+
+ B: batch size;
+ D: feature dimension of each frame;
+ T: number of utterance frames;
+ R: number of right context frames;
+ M: number of memory elements.
+
+ Args:
+ utterance (torch.Tensor): utterance frames, with shape `(T, B, D)`.
+ lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``utterance``.
+ right_context (torch.Tensor): right context frames, with shape `(R, B, D)`.
+ mems (torch.Tensor): memory elements, with shape `(M, B, D)`.
+ attention_mask (torch.Tensor): attention mask for underlying attention module.
+
+ Returns:
+ (Tensor, Tensor, Tensor):
+ Tensor
+ encoded utterance frames, with shape `(T, B, D)`.
+ Tensor
+ updated right context frames, with shape `(R, B, D)`.
+ Tensor
+ updated memory elements, with shape `(M, B, D)`.
+ """
+ (
+ layer_norm_utterance,
+ layer_norm_right_context,
+ ) = self._apply_pre_attention_layer_norm(utterance, right_context)
+ rc_output, output_mems = self._apply_attention_forward(
+ layer_norm_utterance,
+ lengths,
+ layer_norm_right_context,
+ mems,
+ attention_mask,
+ )
+ output_utterance, output_right_context = self._apply_post_attention_ffn(rc_output, utterance, right_context)
+ return output_utterance, output_right_context, output_mems
+
+ @torch.jit.export
+ def infer(
+ self,
+ utterance: torch.Tensor,
+ lengths: torch.Tensor,
+ right_context: torch.Tensor,
+ state: Optional[List[torch.Tensor]],
+ mems: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor], torch.Tensor]:
+ r"""Forward pass for inference.
+
+ B: batch size;
+ D: feature dimension of each frame;
+ T: number of utterance frames;
+ R: number of right context frames;
+ M: number of memory elements.
+
+ Args:
+ utterance (torch.Tensor): utterance frames, with shape `(T, B, D)`.
+ lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``utterance``.
+ right_context (torch.Tensor): right context frames, with shape `(R, B, D)`.
+ state (List[torch.Tensor] or None): list of tensors representing layer internal state
+ generated in preceding invocation of ``infer``.
+ mems (torch.Tensor): memory elements, with shape `(M, B, D)`.
+
+ Returns:
+ (Tensor, Tensor, List[torch.Tensor], Tensor):
+ Tensor
+ encoded utterance frames, with shape `(T, B, D)`.
+ Tensor
+ updated right context frames, with shape `(R, B, D)`.
+ List[Tensor]
+ list of tensors representing layer internal state
+ generated in current invocation of ``infer``.
+ Tensor
+ updated memory elements, with shape `(M, B, D)`.
+ """
+ (
+ layer_norm_utterance,
+ layer_norm_right_context,
+ ) = self._apply_pre_attention_layer_norm(utterance, right_context)
+ rc_output, output_mems, output_state = self._apply_attention_infer(
+ layer_norm_utterance, lengths, layer_norm_right_context, mems, state
+ )
+ output_utterance, output_right_context = self._apply_post_attention_ffn(rc_output, utterance, right_context)
+ return output_utterance, output_right_context, output_state, output_mems
+
+
+class _EmformerImpl(torch.nn.Module):
+ def __init__(
+ self,
+ emformer_layers: torch.nn.ModuleList,
+ segment_length: int,
+ left_context_length: int = 0,
+ right_context_length: int = 0,
+ max_memory_size: int = 0,
+ ):
+ super().__init__()
+
+ self.use_mem = max_memory_size > 0
+ self.memory_op = torch.nn.AvgPool1d(
+ kernel_size=segment_length,
+ stride=segment_length,
+ ceil_mode=True,
+ )
+ self.emformer_layers = emformer_layers
+ self.left_context_length = left_context_length
+ self.right_context_length = right_context_length
+ self.segment_length = segment_length
+ self.max_memory_size = max_memory_size
+
+ def _gen_right_context(self, input: torch.Tensor) -> torch.Tensor:
+ T = input.shape[0]
+ num_segs = math.ceil((T - self.right_context_length) / self.segment_length)
+ right_context_blocks = []
+ for seg_idx in range(num_segs - 1):
+ start = (seg_idx + 1) * self.segment_length
+ end = start + self.right_context_length
+ right_context_blocks.append(input[start:end])
+ right_context_blocks.append(input[T - self.right_context_length :])
+ return torch.cat(right_context_blocks)
+
+ def _gen_attention_mask_col_widths(self, seg_idx: int, utterance_length: int) -> List[int]:
+ num_segs = math.ceil(utterance_length / self.segment_length)
+ rc = self.right_context_length
+ lc = self.left_context_length
+ rc_start = seg_idx * rc
+ rc_end = rc_start + rc
+ seg_start = max(seg_idx * self.segment_length - lc, 0)
+ seg_end = min((seg_idx + 1) * self.segment_length, utterance_length)
+ rc_length = self.right_context_length * num_segs
+
+ if self.use_mem:
+ m_start = max(seg_idx - self.max_memory_size, 0)
+ mem_length = num_segs - 1
+ col_widths = [
+ m_start, # before memory
+ seg_idx - m_start, # memory
+ mem_length - seg_idx, # after memory
+ rc_start, # before right context
+ rc, # right context
+ rc_length - rc_end, # after right context
+ seg_start, # before query segment
+ seg_end - seg_start, # query segment
+ utterance_length - seg_end, # after query segment
+ ]
+ else:
+ col_widths = [
+ rc_start, # before right context
+ rc, # right context
+ rc_length - rc_end, # after right context
+ seg_start, # before query segment
+ seg_end - seg_start, # query segment
+ utterance_length - seg_end, # after query segment
+ ]
+
+ return col_widths
+
+ def _gen_attention_mask(self, input: torch.Tensor) -> torch.Tensor:
+ utterance_length = input.size(0)
+ num_segs = math.ceil(utterance_length / self.segment_length)
+
+ rc_mask = []
+ query_mask = []
+ summary_mask = []
+
+ if self.use_mem:
+ num_cols = 9
+ # memory, right context, query segment
+ rc_q_cols_mask = [idx in [1, 4, 7] for idx in range(num_cols)]
+ # right context, query segment
+ s_cols_mask = [idx in [4, 7] for idx in range(num_cols)]
+ masks_to_concat = [rc_mask, query_mask, summary_mask]
+ else:
+ num_cols = 6
+ # right context, query segment
+ rc_q_cols_mask = [idx in [1, 4] for idx in range(num_cols)]
+ s_cols_mask = None
+ masks_to_concat = [rc_mask, query_mask]
+
+ for seg_idx in range(num_segs):
+ col_widths = self._gen_attention_mask_col_widths(seg_idx, utterance_length)
+
+ rc_mask_block = _gen_attention_mask_block(
+ col_widths, rc_q_cols_mask, self.right_context_length, input.device
+ )
+ rc_mask.append(rc_mask_block)
+
+ query_mask_block = _gen_attention_mask_block(
+ col_widths,
+ rc_q_cols_mask,
+ min(
+ self.segment_length,
+ utterance_length - seg_idx * self.segment_length,
+ ),
+ input.device,
+ )
+ query_mask.append(query_mask_block)
+
+ if s_cols_mask is not None:
+ summary_mask_block = _gen_attention_mask_block(col_widths, s_cols_mask, 1, input.device)
+ summary_mask.append(summary_mask_block)
+
+ attention_mask = (1 - torch.cat([torch.cat(mask) for mask in masks_to_concat])).to(torch.bool)
+ return attention_mask
+
+ def forward(self, input: torch.Tensor, lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ r"""Forward pass for training and non-streaming inference.
+
+ B: batch size;
+ T: max number of input frames in batch;
+ D: feature dimension of each frame.
+
+ Args:
+ input (torch.Tensor): utterance frames right-padded with right context frames, with
+ shape `(B, T + right_context_length, D)`.
+ lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid utterance frames for i-th batch element in ``input``.
+
+ Returns:
+ (Tensor, Tensor):
+ Tensor
+ output frames, with shape `(B, T, D)`.
+ Tensor
+ output lengths, with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in output frames.
+ """
+ input = input.permute(1, 0, 2)
+ right_context = self._gen_right_context(input)
+ utterance = input[: input.size(0) - self.right_context_length]
+ attention_mask = self._gen_attention_mask(utterance)
+ mems = (
+ self.memory_op(utterance.permute(1, 2, 0)).permute(2, 0, 1)[:-1]
+ if self.use_mem
+ else torch.empty(0).to(dtype=input.dtype, device=input.device)
+ )
+ output = utterance
+ for layer in self.emformer_layers:
+ output, right_context, mems = layer(output, lengths, right_context, mems, attention_mask)
+ return output.permute(1, 0, 2), lengths
+
+ @torch.jit.export
+ def infer(
+ self,
+ input: torch.Tensor,
+ lengths: torch.Tensor,
+ states: Optional[List[List[torch.Tensor]]] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, List[List[torch.Tensor]]]:
+ r"""Forward pass for streaming inference.
+
+ B: batch size;
+ D: feature dimension of each frame.
+
+ Args:
+ input (torch.Tensor): utterance frames right-padded with right context frames, with
+ shape `(B, segment_length + right_context_length, D)`.
+ lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``input``.
+ states (List[List[torch.Tensor]] or None, optional): list of lists of tensors
+ representing internal state generated in preceding invocation of ``infer``. (Default: ``None``)
+
+ Returns:
+ (Tensor, Tensor, List[List[Tensor]]):
+ Tensor
+ output frames, with shape `(B, segment_length, D)`.
+ Tensor
+ output lengths, with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in output frames.
+ List[List[Tensor]]
+ output states; list of lists of tensors representing internal state
+ generated in current invocation of ``infer``.
+ """
+ if input.size(1) != self.segment_length + self.right_context_length:
+ raise ValueError(
+ "Per configured segment_length and right_context_length"
+ f", expected size of {self.segment_length + self.right_context_length} for dimension 1 of input"
+ f", but got {input.size(1)}."
+ )
+ input = input.permute(1, 0, 2)
+ right_context_start_idx = input.size(0) - self.right_context_length
+ right_context = input[right_context_start_idx:]
+ utterance = input[:right_context_start_idx]
+ output_lengths = torch.clamp(lengths - self.right_context_length, min=0)
+ mems = (
+ self.memory_op(utterance.permute(1, 2, 0)).permute(2, 0, 1)
+ if self.use_mem
+ else torch.empty(0).to(dtype=input.dtype, device=input.device)
+ )
+ output = utterance
+ output_states: List[List[torch.Tensor]] = []
+ for layer_idx, layer in enumerate(self.emformer_layers):
+ output, right_context, output_state, mems = layer.infer(
+ output,
+ output_lengths,
+ right_context,
+ None if states is None else states[layer_idx],
+ mems,
+ )
+ output_states.append(output_state)
+
+ return output.permute(1, 0, 2), output_lengths, output_states
+
+
+class Emformer(_EmformerImpl):
+ r"""Emformer architecture introduced in
+ *Emformer: Efficient Memory Transformer Based Acoustic Model for Low Latency Streaming Speech Recognition*
+ :cite:`shi2021emformer`.
+
+ See Also:
+ * :func:`~torchaudio.models.emformer_rnnt_model`,
+ :func:`~torchaudio.models.emformer_rnnt_base`: factory functions.
+ * :class:`torchaudio.pipelines.RNNTBundle`: ASR pipelines with pretrained model.
+
+ Args:
+ input_dim (int): input dimension.
+ num_heads (int): number of attention heads in each Emformer layer.
+ ffn_dim (int): hidden layer dimension of each Emformer layer's feedforward network.
+ num_layers (int): number of Emformer layers to instantiate.
+ segment_length (int): length of each input segment.
+ dropout (float, optional): dropout probability. (Default: 0.0)
+ activation (str, optional): activation function to use in each Emformer layer's
+ feedforward network. Must be one of ("relu", "gelu", "silu"). (Default: "relu")
+ left_context_length (int, optional): length of left context. (Default: 0)
+ right_context_length (int, optional): length of right context. (Default: 0)
+ max_memory_size (int, optional): maximum number of memory elements to use. (Default: 0)
+ weight_init_scale_strategy (str or None, optional): per-layer weight initialization scaling
+ strategy. Must be one of ("depthwise", "constant", ``None``). (Default: "depthwise")
+ tanh_on_mem (bool, optional): if ``True``, applies tanh to memory elements. (Default: ``False``)
+ negative_inf (float, optional): value to use for negative infinity in attention weights. (Default: -1e8)
+
+ Examples:
+ >>> emformer = Emformer(512, 8, 2048, 20, 4, right_context_length=1)
+ >>> input = torch.rand(128, 400, 512) # batch, num_frames, feature_dim
+ >>> lengths = torch.randint(1, 200, (128,)) # batch
+ >>> output, lengths = emformer(input, lengths)
+ >>> input = torch.rand(128, 5, 512)
+ >>> lengths = torch.ones(128) * 5
+ >>> output, lengths, states = emformer.infer(input, lengths, None)
+ """
+
+ def __init__(
+ self,
+ input_dim: int,
+ num_heads: int,
+ ffn_dim: int,
+ num_layers: int,
+ segment_length: int,
+ dropout: float = 0.0,
+ activation: str = "relu",
+ left_context_length: int = 0,
+ right_context_length: int = 0,
+ max_memory_size: int = 0,
+ weight_init_scale_strategy: Optional[str] = "depthwise",
+ tanh_on_mem: bool = False,
+ negative_inf: float = -1e8,
+ ):
+ weight_init_gains = _get_weight_init_gains(weight_init_scale_strategy, num_layers)
+ emformer_layers = torch.nn.ModuleList(
+ [
+ _EmformerLayer(
+ input_dim,
+ num_heads,
+ ffn_dim,
+ segment_length,
+ dropout=dropout,
+ activation=activation,
+ left_context_length=left_context_length,
+ max_memory_size=max_memory_size,
+ weight_init_gain=weight_init_gains[layer_idx],
+ tanh_on_mem=tanh_on_mem,
+ negative_inf=negative_inf,
+ )
+ for layer_idx in range(num_layers)
+ ]
+ )
+ super().__init__(
+ emformer_layers,
+ segment_length,
+ left_context_length=left_context_length,
+ right_context_length=right_context_length,
+ max_memory_size=max_memory_size,
+ )
diff --git a/lib/python3.12/site-packages/torchaudio/models/rnnt.py b/lib/python3.12/site-packages/torchaudio/models/rnnt.py
new file mode 100644
index 0000000000000000000000000000000000000000..f9dbe22c9fb4a97cf7f8779a953b5bd7b5bbffd9
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/rnnt.py
@@ -0,0 +1,816 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional, Tuple
+
+import torch
+from torchaudio.models import Emformer
+
+
+__all__ = ["RNNT", "emformer_rnnt_base", "emformer_rnnt_model"]
+
+
+class _TimeReduction(torch.nn.Module):
+ r"""Coalesces frames along time dimension into a
+ fewer number of frames with higher feature dimensionality.
+
+ Args:
+ stride (int): number of frames to merge for each output frame.
+ """
+
+ def __init__(self, stride: int) -> None:
+ super().__init__()
+ self.stride = stride
+
+ def forward(self, input: torch.Tensor, lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ r"""Forward pass.
+
+ B: batch size;
+ T: maximum input sequence length in batch;
+ D: feature dimension of each input sequence frame.
+
+ Args:
+ input (torch.Tensor): input sequences, with shape `(B, T, D)`.
+ lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``input``.
+
+ Returns:
+ (torch.Tensor, torch.Tensor):
+ torch.Tensor
+ output sequences, with shape
+ `(B, T // stride, D * stride)`
+ torch.Tensor
+ output lengths, with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in output sequences.
+ """
+ B, T, D = input.shape
+ num_frames = T - (T % self.stride)
+ input = input[:, :num_frames, :]
+ lengths = lengths.div(self.stride, rounding_mode="trunc")
+ T_max = num_frames // self.stride
+
+ output = input.reshape(B, T_max, D * self.stride)
+ output = output.contiguous()
+ return output, lengths
+
+
+class _CustomLSTM(torch.nn.Module):
+ r"""Custom long-short-term memory (LSTM) block that applies layer normalization
+ to internal nodes.
+
+ Args:
+ input_dim (int): input dimension.
+ hidden_dim (int): hidden dimension.
+ layer_norm (bool, optional): if ``True``, enables layer normalization. (Default: ``False``)
+ layer_norm_epsilon (float, optional): value of epsilon to use in
+ layer normalization layers (Default: 1e-5)
+ """
+
+ def __init__(
+ self,
+ input_dim: int,
+ hidden_dim: int,
+ layer_norm: bool = False,
+ layer_norm_epsilon: float = 1e-5,
+ ) -> None:
+ super().__init__()
+ self.x2g = torch.nn.Linear(input_dim, 4 * hidden_dim, bias=(not layer_norm))
+ self.p2g = torch.nn.Linear(hidden_dim, 4 * hidden_dim, bias=False)
+ if layer_norm:
+ self.c_norm = torch.nn.LayerNorm(hidden_dim, eps=layer_norm_epsilon)
+ self.g_norm = torch.nn.LayerNorm(4 * hidden_dim, eps=layer_norm_epsilon)
+ else:
+ self.c_norm = torch.nn.Identity()
+ self.g_norm = torch.nn.Identity()
+
+ self.hidden_dim = hidden_dim
+
+ def forward(
+ self, input: torch.Tensor, state: Optional[List[torch.Tensor]]
+ ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
+ r"""Forward pass.
+
+ B: batch size;
+ T: maximum sequence length in batch;
+ D: feature dimension of each input sequence element.
+
+ Args:
+ input (torch.Tensor): with shape `(T, B, D)`.
+ state (List[torch.Tensor] or None): list of tensors
+ representing internal state generated in preceding invocation
+ of ``forward``.
+
+ Returns:
+ (torch.Tensor, List[torch.Tensor]):
+ torch.Tensor
+ output, with shape `(T, B, hidden_dim)`.
+ List[torch.Tensor]
+ list of tensors representing internal state generated
+ in current invocation of ``forward``.
+ """
+ if state is None:
+ B = input.size(1)
+ h = torch.zeros(B, self.hidden_dim, device=input.device, dtype=input.dtype)
+ c = torch.zeros(B, self.hidden_dim, device=input.device, dtype=input.dtype)
+ else:
+ h, c = state
+
+ gated_input = self.x2g(input)
+ outputs = []
+ for gates in gated_input.unbind(0):
+ gates = gates + self.p2g(h)
+ gates = self.g_norm(gates)
+ input_gate, forget_gate, cell_gate, output_gate = gates.chunk(4, 1)
+ input_gate = input_gate.sigmoid()
+ forget_gate = forget_gate.sigmoid()
+ cell_gate = cell_gate.tanh()
+ output_gate = output_gate.sigmoid()
+ c = forget_gate * c + input_gate * cell_gate
+ c = self.c_norm(c)
+ h = output_gate * c.tanh()
+ outputs.append(h)
+
+ output = torch.stack(outputs, dim=0)
+ state = [h, c]
+
+ return output, state
+
+
+class _Transcriber(ABC):
+ @abstractmethod
+ def forward(self, input: torch.Tensor, lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ pass
+
+ @abstractmethod
+ def infer(
+ self,
+ input: torch.Tensor,
+ lengths: torch.Tensor,
+ states: Optional[List[List[torch.Tensor]]],
+ ) -> Tuple[torch.Tensor, torch.Tensor, List[List[torch.Tensor]]]:
+ pass
+
+
+class _EmformerEncoder(torch.nn.Module, _Transcriber):
+ r"""Emformer-based recurrent neural network transducer (RNN-T) encoder (transcription network).
+
+ Args:
+ input_dim (int): feature dimension of each input sequence element.
+ output_dim (int): feature dimension of each output sequence element.
+ segment_length (int): length of input segment expressed as number of frames.
+ right_context_length (int): length of right context expressed as number of frames.
+ time_reduction_input_dim (int): dimension to scale each element in input sequences to
+ prior to applying time reduction block.
+ time_reduction_stride (int): factor by which to reduce length of input sequence.
+ transformer_num_heads (int): number of attention heads in each Emformer layer.
+ transformer_ffn_dim (int): hidden layer dimension of each Emformer layer's feedforward network.
+ transformer_num_layers (int): number of Emformer layers to instantiate.
+ transformer_left_context_length (int): length of left context.
+ transformer_dropout (float, optional): transformer dropout probability. (Default: 0.0)
+ transformer_activation (str, optional): activation function to use in each Emformer layer's
+ feedforward network. Must be one of ("relu", "gelu", "silu"). (Default: "relu")
+ transformer_max_memory_size (int, optional): maximum number of memory elements to use. (Default: 0)
+ transformer_weight_init_scale_strategy (str, optional): per-layer weight initialization scaling
+ strategy. Must be one of ("depthwise", "constant", ``None``). (Default: "depthwise")
+ transformer_tanh_on_mem (bool, optional): if ``True``, applies tanh to memory elements. (Default: ``False``)
+ """
+
+ def __init__(
+ self,
+ *,
+ input_dim: int,
+ output_dim: int,
+ segment_length: int,
+ right_context_length: int,
+ time_reduction_input_dim: int,
+ time_reduction_stride: int,
+ transformer_num_heads: int,
+ transformer_ffn_dim: int,
+ transformer_num_layers: int,
+ transformer_left_context_length: int,
+ transformer_dropout: float = 0.0,
+ transformer_activation: str = "relu",
+ transformer_max_memory_size: int = 0,
+ transformer_weight_init_scale_strategy: str = "depthwise",
+ transformer_tanh_on_mem: bool = False,
+ ) -> None:
+ super().__init__()
+ self.input_linear = torch.nn.Linear(
+ input_dim,
+ time_reduction_input_dim,
+ bias=False,
+ )
+ self.time_reduction = _TimeReduction(time_reduction_stride)
+ transformer_input_dim = time_reduction_input_dim * time_reduction_stride
+ self.transformer = Emformer(
+ transformer_input_dim,
+ transformer_num_heads,
+ transformer_ffn_dim,
+ transformer_num_layers,
+ segment_length // time_reduction_stride,
+ dropout=transformer_dropout,
+ activation=transformer_activation,
+ left_context_length=transformer_left_context_length,
+ right_context_length=right_context_length // time_reduction_stride,
+ max_memory_size=transformer_max_memory_size,
+ weight_init_scale_strategy=transformer_weight_init_scale_strategy,
+ tanh_on_mem=transformer_tanh_on_mem,
+ )
+ self.output_linear = torch.nn.Linear(transformer_input_dim, output_dim)
+ self.layer_norm = torch.nn.LayerNorm(output_dim)
+
+ def forward(self, input: torch.Tensor, lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ r"""Forward pass for training.
+
+ B: batch size;
+ T: maximum input sequence length in batch;
+ D: feature dimension of each input sequence frame (input_dim).
+
+ Args:
+ input (torch.Tensor): input frame sequences right-padded with right context, with
+ shape `(B, T + right context length, D)`.
+ lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``input``.
+
+ Returns:
+ (torch.Tensor, torch.Tensor):
+ torch.Tensor
+ output frame sequences, with
+ shape `(B, T // time_reduction_stride, output_dim)`.
+ torch.Tensor
+ output input lengths, with shape `(B,)` and i-th element representing
+ number of valid elements for i-th batch element in output frame sequences.
+ """
+ input_linear_out = self.input_linear(input)
+ time_reduction_out, time_reduction_lengths = self.time_reduction(input_linear_out, lengths)
+ transformer_out, transformer_lengths = self.transformer(time_reduction_out, time_reduction_lengths)
+ output_linear_out = self.output_linear(transformer_out)
+ layer_norm_out = self.layer_norm(output_linear_out)
+ return layer_norm_out, transformer_lengths
+
+ @torch.jit.export
+ def infer(
+ self,
+ input: torch.Tensor,
+ lengths: torch.Tensor,
+ states: Optional[List[List[torch.Tensor]]],
+ ) -> Tuple[torch.Tensor, torch.Tensor, List[List[torch.Tensor]]]:
+ r"""Forward pass for inference.
+
+ B: batch size;
+ T: maximum input sequence segment length in batch;
+ D: feature dimension of each input sequence frame (input_dim).
+
+ Args:
+ input (torch.Tensor): input frame sequence segments right-padded with right context, with
+ shape `(B, T + right context length, D)`.
+ lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``input``.
+ state (List[List[torch.Tensor]] or None): list of lists of tensors
+ representing internal state generated in preceding invocation
+ of ``infer``.
+
+ Returns:
+ (torch.Tensor, torch.Tensor, List[List[torch.Tensor]]):
+ torch.Tensor
+ output frame sequences, with
+ shape `(B, T // time_reduction_stride, output_dim)`.
+ torch.Tensor
+ output input lengths, with shape `(B,)` and i-th element representing
+ number of valid elements for i-th batch element in output.
+ List[List[torch.Tensor]]
+ output states; list of lists of tensors
+ representing internal state generated in current invocation
+ of ``infer``.
+ """
+ input_linear_out = self.input_linear(input)
+ time_reduction_out, time_reduction_lengths = self.time_reduction(input_linear_out, lengths)
+ (
+ transformer_out,
+ transformer_lengths,
+ transformer_states,
+ ) = self.transformer.infer(time_reduction_out, time_reduction_lengths, states)
+ output_linear_out = self.output_linear(transformer_out)
+ layer_norm_out = self.layer_norm(output_linear_out)
+ return layer_norm_out, transformer_lengths, transformer_states
+
+
+class _Predictor(torch.nn.Module):
+ r"""Recurrent neural network transducer (RNN-T) prediction network.
+
+ Args:
+ num_symbols (int): size of target token lexicon.
+ output_dim (int): feature dimension of each output sequence element.
+ symbol_embedding_dim (int): dimension of each target token embedding.
+ num_lstm_layers (int): number of LSTM layers to instantiate.
+ lstm_hidden_dim (int): output dimension of each LSTM layer.
+ lstm_layer_norm (bool, optional): if ``True``, enables layer normalization
+ for LSTM layers. (Default: ``False``)
+ lstm_layer_norm_epsilon (float, optional): value of epsilon to use in
+ LSTM layer normalization layers. (Default: 1e-5)
+ lstm_dropout (float, optional): LSTM dropout probability. (Default: 0.0)
+
+ """
+
+ def __init__(
+ self,
+ num_symbols: int,
+ output_dim: int,
+ symbol_embedding_dim: int,
+ num_lstm_layers: int,
+ lstm_hidden_dim: int,
+ lstm_layer_norm: bool = False,
+ lstm_layer_norm_epsilon: float = 1e-5,
+ lstm_dropout: float = 0.0,
+ ) -> None:
+ super().__init__()
+ self.embedding = torch.nn.Embedding(num_symbols, symbol_embedding_dim)
+ self.input_layer_norm = torch.nn.LayerNorm(symbol_embedding_dim)
+ self.lstm_layers = torch.nn.ModuleList(
+ [
+ _CustomLSTM(
+ symbol_embedding_dim if idx == 0 else lstm_hidden_dim,
+ lstm_hidden_dim,
+ layer_norm=lstm_layer_norm,
+ layer_norm_epsilon=lstm_layer_norm_epsilon,
+ )
+ for idx in range(num_lstm_layers)
+ ]
+ )
+ self.dropout = torch.nn.Dropout(p=lstm_dropout)
+ self.linear = torch.nn.Linear(lstm_hidden_dim, output_dim)
+ self.output_layer_norm = torch.nn.LayerNorm(output_dim)
+
+ self.lstm_dropout = lstm_dropout
+
+ def forward(
+ self,
+ input: torch.Tensor,
+ lengths: torch.Tensor,
+ state: Optional[List[List[torch.Tensor]]] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, List[List[torch.Tensor]]]:
+ r"""Forward pass.
+
+ B: batch size;
+ U: maximum sequence length in batch;
+ D: feature dimension of each input sequence element.
+
+ Args:
+ input (torch.Tensor): target sequences, with shape `(B, U)` and each element
+ mapping to a target symbol, i.e. in range `[0, num_symbols)`.
+ lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``input``.
+ state (List[List[torch.Tensor]] or None, optional): list of lists of tensors
+ representing internal state generated in preceding invocation
+ of ``forward``. (Default: ``None``)
+
+ Returns:
+ (torch.Tensor, torch.Tensor, List[List[torch.Tensor]]):
+ torch.Tensor
+ output encoding sequences, with shape `(B, U, output_dim)`
+ torch.Tensor
+ output lengths, with shape `(B,)` and i-th element representing
+ number of valid elements for i-th batch element in output encoding sequences.
+ List[List[torch.Tensor]]
+ output states; list of lists of tensors
+ representing internal state generated in current invocation of ``forward``.
+ """
+ input_tb = input.permute(1, 0)
+ embedding_out = self.embedding(input_tb)
+ input_layer_norm_out = self.input_layer_norm(embedding_out)
+
+ lstm_out = input_layer_norm_out
+ state_out: List[List[torch.Tensor]] = []
+ for layer_idx, lstm in enumerate(self.lstm_layers):
+ lstm_out, lstm_state_out = lstm(lstm_out, None if state is None else state[layer_idx])
+ lstm_out = self.dropout(lstm_out)
+ state_out.append(lstm_state_out)
+
+ linear_out = self.linear(lstm_out)
+ output_layer_norm_out = self.output_layer_norm(linear_out)
+ return output_layer_norm_out.permute(1, 0, 2), lengths, state_out
+
+
+class _Joiner(torch.nn.Module):
+ r"""Recurrent neural network transducer (RNN-T) joint network.
+
+ Args:
+ input_dim (int): source and target input dimension.
+ output_dim (int): output dimension.
+ activation (str, optional): activation function to use in the joiner.
+ Must be one of ("relu", "tanh"). (Default: "relu")
+
+ """
+
+ def __init__(self, input_dim: int, output_dim: int, activation: str = "relu") -> None:
+ super().__init__()
+ self.linear = torch.nn.Linear(input_dim, output_dim, bias=True)
+ if activation == "relu":
+ self.activation = torch.nn.ReLU()
+ elif activation == "tanh":
+ self.activation = torch.nn.Tanh()
+ else:
+ raise ValueError(f"Unsupported activation {activation}")
+
+ def forward(
+ self,
+ source_encodings: torch.Tensor,
+ source_lengths: torch.Tensor,
+ target_encodings: torch.Tensor,
+ target_lengths: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ r"""Forward pass for training.
+
+ B: batch size;
+ T: maximum source sequence length in batch;
+ U: maximum target sequence length in batch;
+ D: dimension of each source and target sequence encoding.
+
+ Args:
+ source_encodings (torch.Tensor): source encoding sequences, with
+ shape `(B, T, D)`.
+ source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ valid sequence length of i-th batch element in ``source_encodings``.
+ target_encodings (torch.Tensor): target encoding sequences, with shape `(B, U, D)`.
+ target_lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ valid sequence length of i-th batch element in ``target_encodings``.
+
+ Returns:
+ (torch.Tensor, torch.Tensor, torch.Tensor):
+ torch.Tensor
+ joint network output, with shape `(B, T, U, output_dim)`.
+ torch.Tensor
+ output source lengths, with shape `(B,)` and i-th element representing
+ number of valid elements along dim 1 for i-th batch element in joint network output.
+ torch.Tensor
+ output target lengths, with shape `(B,)` and i-th element representing
+ number of valid elements along dim 2 for i-th batch element in joint network output.
+ """
+ joint_encodings = source_encodings.unsqueeze(2).contiguous() + target_encodings.unsqueeze(1).contiguous()
+ activation_out = self.activation(joint_encodings)
+ output = self.linear(activation_out)
+ return output, source_lengths, target_lengths
+
+
+class RNNT(torch.nn.Module):
+ r"""torchaudio.models.RNNT()
+
+ Recurrent neural network transducer (RNN-T) model.
+
+ Note:
+ To build the model, please use one of the factory functions.
+
+ See Also:
+ :class:`torchaudio.pipelines.RNNTBundle`: ASR pipeline with pre-trained models.
+
+ Args:
+ transcriber (torch.nn.Module): transcription network.
+ predictor (torch.nn.Module): prediction network.
+ joiner (torch.nn.Module): joint network.
+ """
+
+ def __init__(self, transcriber: _Transcriber, predictor: _Predictor, joiner: _Joiner) -> None:
+ super().__init__()
+ self.transcriber = transcriber
+ self.predictor = predictor
+ self.joiner = joiner
+
+ def forward(
+ self,
+ sources: torch.Tensor,
+ source_lengths: torch.Tensor,
+ targets: torch.Tensor,
+ target_lengths: torch.Tensor,
+ predictor_state: Optional[List[List[torch.Tensor]]] = None,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[List[torch.Tensor]]]:
+ r"""Forward pass for training.
+
+ B: batch size;
+ T: maximum source sequence length in batch;
+ U: maximum target sequence length in batch;
+ D: feature dimension of each source sequence element.
+
+ Args:
+ sources (torch.Tensor): source frame sequences right-padded with right context, with
+ shape `(B, T, D)`.
+ source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``sources``.
+ targets (torch.Tensor): target sequences, with shape `(B, U)` and each element
+ mapping to a target symbol.
+ target_lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``targets``.
+ predictor_state (List[List[torch.Tensor]] or None, optional): list of lists of tensors
+ representing prediction network internal state generated in preceding invocation
+ of ``forward``. (Default: ``None``)
+
+ Returns:
+ (torch.Tensor, torch.Tensor, torch.Tensor, List[List[torch.Tensor]]):
+ torch.Tensor
+ joint network output, with shape
+ `(B, max output source length, max output target length, output_dim (number of target symbols))`.
+ torch.Tensor
+ output source lengths, with shape `(B,)` and i-th element representing
+ number of valid elements along dim 1 for i-th batch element in joint network output.
+ torch.Tensor
+ output target lengths, with shape `(B,)` and i-th element representing
+ number of valid elements along dim 2 for i-th batch element in joint network output.
+ List[List[torch.Tensor]]
+ output states; list of lists of tensors
+ representing prediction network internal state generated in current invocation
+ of ``forward``.
+ """
+ source_encodings, source_lengths = self.transcriber(
+ input=sources,
+ lengths=source_lengths,
+ )
+ target_encodings, target_lengths, predictor_state = self.predictor(
+ input=targets,
+ lengths=target_lengths,
+ state=predictor_state,
+ )
+ output, source_lengths, target_lengths = self.joiner(
+ source_encodings=source_encodings,
+ source_lengths=source_lengths,
+ target_encodings=target_encodings,
+ target_lengths=target_lengths,
+ )
+
+ return (
+ output,
+ source_lengths,
+ target_lengths,
+ predictor_state,
+ )
+
+ @torch.jit.export
+ def transcribe_streaming(
+ self,
+ sources: torch.Tensor,
+ source_lengths: torch.Tensor,
+ state: Optional[List[List[torch.Tensor]]],
+ ) -> Tuple[torch.Tensor, torch.Tensor, List[List[torch.Tensor]]]:
+ r"""Applies transcription network to sources in streaming mode.
+
+ B: batch size;
+ T: maximum source sequence segment length in batch;
+ D: feature dimension of each source sequence frame.
+
+ Args:
+ sources (torch.Tensor): source frame sequence segments right-padded with right context, with
+ shape `(B, T + right context length, D)`.
+ source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``sources``.
+ state (List[List[torch.Tensor]] or None): list of lists of tensors
+ representing transcription network internal state generated in preceding invocation
+ of ``transcribe_streaming``.
+
+ Returns:
+ (torch.Tensor, torch.Tensor, List[List[torch.Tensor]]):
+ torch.Tensor
+ output frame sequences, with
+ shape `(B, T // time_reduction_stride, output_dim)`.
+ torch.Tensor
+ output lengths, with shape `(B,)` and i-th element representing
+ number of valid elements for i-th batch element in output.
+ List[List[torch.Tensor]]
+ output states; list of lists of tensors
+ representing transcription network internal state generated in current invocation
+ of ``transcribe_streaming``.
+ """
+ return self.transcriber.infer(sources, source_lengths, state)
+
+ @torch.jit.export
+ def transcribe(
+ self,
+ sources: torch.Tensor,
+ source_lengths: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ r"""Applies transcription network to sources in non-streaming mode.
+
+ B: batch size;
+ T: maximum source sequence length in batch;
+ D: feature dimension of each source sequence frame.
+
+ Args:
+ sources (torch.Tensor): source frame sequences right-padded with right context, with
+ shape `(B, T + right context length, D)`.
+ source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``sources``.
+
+ Returns:
+ (torch.Tensor, torch.Tensor):
+ torch.Tensor
+ output frame sequences, with
+ shape `(B, T // time_reduction_stride, output_dim)`.
+ torch.Tensor
+ output lengths, with shape `(B,)` and i-th element representing
+ number of valid elements for i-th batch element in output frame sequences.
+ """
+ return self.transcriber(sources, source_lengths)
+
+ @torch.jit.export
+ def predict(
+ self,
+ targets: torch.Tensor,
+ target_lengths: torch.Tensor,
+ state: Optional[List[List[torch.Tensor]]],
+ ) -> Tuple[torch.Tensor, torch.Tensor, List[List[torch.Tensor]]]:
+ r"""Applies prediction network to targets.
+
+ B: batch size;
+ U: maximum target sequence length in batch;
+ D: feature dimension of each target sequence frame.
+
+ Args:
+ targets (torch.Tensor): target sequences, with shape `(B, U)` and each element
+ mapping to a target symbol, i.e. in range `[0, num_symbols)`.
+ target_lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ number of valid frames for i-th batch element in ``targets``.
+ state (List[List[torch.Tensor]] or None): list of lists of tensors
+ representing internal state generated in preceding invocation
+ of ``predict``.
+
+ Returns:
+ (torch.Tensor, torch.Tensor, List[List[torch.Tensor]]):
+ torch.Tensor
+ output frame sequences, with shape `(B, U, output_dim)`.
+ torch.Tensor
+ output lengths, with shape `(B,)` and i-th element representing
+ number of valid elements for i-th batch element in output.
+ List[List[torch.Tensor]]
+ output states; list of lists of tensors
+ representing internal state generated in current invocation of ``predict``.
+ """
+ return self.predictor(input=targets, lengths=target_lengths, state=state)
+
+ @torch.jit.export
+ def join(
+ self,
+ source_encodings: torch.Tensor,
+ source_lengths: torch.Tensor,
+ target_encodings: torch.Tensor,
+ target_lengths: torch.Tensor,
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ r"""Applies joint network to source and target encodings.
+
+ B: batch size;
+ T: maximum source sequence length in batch;
+ U: maximum target sequence length in batch;
+ D: dimension of each source and target sequence encoding.
+
+ Args:
+ source_encodings (torch.Tensor): source encoding sequences, with
+ shape `(B, T, D)`.
+ source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ valid sequence length of i-th batch element in ``source_encodings``.
+ target_encodings (torch.Tensor): target encoding sequences, with shape `(B, U, D)`.
+ target_lengths (torch.Tensor): with shape `(B,)` and i-th element representing
+ valid sequence length of i-th batch element in ``target_encodings``.
+
+ Returns:
+ (torch.Tensor, torch.Tensor, torch.Tensor):
+ torch.Tensor
+ joint network output, with shape `(B, T, U, output_dim)`.
+ torch.Tensor
+ output source lengths, with shape `(B,)` and i-th element representing
+ number of valid elements along dim 1 for i-th batch element in joint network output.
+ torch.Tensor
+ output target lengths, with shape `(B,)` and i-th element representing
+ number of valid elements along dim 2 for i-th batch element in joint network output.
+ """
+ output, source_lengths, target_lengths = self.joiner(
+ source_encodings=source_encodings,
+ source_lengths=source_lengths,
+ target_encodings=target_encodings,
+ target_lengths=target_lengths,
+ )
+ return output, source_lengths, target_lengths
+
+
+def emformer_rnnt_model(
+ *,
+ input_dim: int,
+ encoding_dim: int,
+ num_symbols: int,
+ segment_length: int,
+ right_context_length: int,
+ time_reduction_input_dim: int,
+ time_reduction_stride: int,
+ transformer_num_heads: int,
+ transformer_ffn_dim: int,
+ transformer_num_layers: int,
+ transformer_dropout: float,
+ transformer_activation: str,
+ transformer_left_context_length: int,
+ transformer_max_memory_size: int,
+ transformer_weight_init_scale_strategy: str,
+ transformer_tanh_on_mem: bool,
+ symbol_embedding_dim: int,
+ num_lstm_layers: int,
+ lstm_layer_norm: bool,
+ lstm_layer_norm_epsilon: float,
+ lstm_dropout: float,
+) -> RNNT:
+ r"""Builds Emformer-based :class:`~torchaudio.models.RNNT`.
+
+ Note:
+ For non-streaming inference, the expectation is for `transcribe` to be called on input
+ sequences right-concatenated with `right_context_length` frames.
+
+ For streaming inference, the expectation is for `transcribe_streaming` to be called
+ on input chunks comprising `segment_length` frames right-concatenated with `right_context_length`
+ frames.
+
+ Args:
+ input_dim (int): dimension of input sequence frames passed to transcription network.
+ encoding_dim (int): dimension of transcription- and prediction-network-generated encodings
+ passed to joint network.
+ num_symbols (int): cardinality of set of target tokens.
+ segment_length (int): length of input segment expressed as number of frames.
+ right_context_length (int): length of right context expressed as number of frames.
+ time_reduction_input_dim (int): dimension to scale each element in input sequences to
+ prior to applying time reduction block.
+ time_reduction_stride (int): factor by which to reduce length of input sequence.
+ transformer_num_heads (int): number of attention heads in each Emformer layer.
+ transformer_ffn_dim (int): hidden layer dimension of each Emformer layer's feedforward network.
+ transformer_num_layers (int): number of Emformer layers to instantiate.
+ transformer_left_context_length (int): length of left context considered by Emformer.
+ transformer_dropout (float): Emformer dropout probability.
+ transformer_activation (str): activation function to use in each Emformer layer's
+ feedforward network. Must be one of ("relu", "gelu", "silu").
+ transformer_max_memory_size (int): maximum number of memory elements to use.
+ transformer_weight_init_scale_strategy (str): per-layer weight initialization scaling
+ strategy. Must be one of ("depthwise", "constant", ``None``).
+ transformer_tanh_on_mem (bool): if ``True``, applies tanh to memory elements.
+ symbol_embedding_dim (int): dimension of each target token embedding.
+ num_lstm_layers (int): number of LSTM layers to instantiate.
+ lstm_layer_norm (bool): if ``True``, enables layer normalization for LSTM layers.
+ lstm_layer_norm_epsilon (float): value of epsilon to use in LSTM layer normalization layers.
+ lstm_dropout (float): LSTM dropout probability.
+
+ Returns:
+ RNNT:
+ Emformer RNN-T model.
+ """
+ encoder = _EmformerEncoder(
+ input_dim=input_dim,
+ output_dim=encoding_dim,
+ segment_length=segment_length,
+ right_context_length=right_context_length,
+ time_reduction_input_dim=time_reduction_input_dim,
+ time_reduction_stride=time_reduction_stride,
+ transformer_num_heads=transformer_num_heads,
+ transformer_ffn_dim=transformer_ffn_dim,
+ transformer_num_layers=transformer_num_layers,
+ transformer_dropout=transformer_dropout,
+ transformer_activation=transformer_activation,
+ transformer_left_context_length=transformer_left_context_length,
+ transformer_max_memory_size=transformer_max_memory_size,
+ transformer_weight_init_scale_strategy=transformer_weight_init_scale_strategy,
+ transformer_tanh_on_mem=transformer_tanh_on_mem,
+ )
+ predictor = _Predictor(
+ num_symbols,
+ encoding_dim,
+ symbol_embedding_dim=symbol_embedding_dim,
+ num_lstm_layers=num_lstm_layers,
+ lstm_hidden_dim=symbol_embedding_dim,
+ lstm_layer_norm=lstm_layer_norm,
+ lstm_layer_norm_epsilon=lstm_layer_norm_epsilon,
+ lstm_dropout=lstm_dropout,
+ )
+ joiner = _Joiner(encoding_dim, num_symbols)
+ return RNNT(encoder, predictor, joiner)
+
+
+def emformer_rnnt_base(num_symbols: int) -> RNNT:
+ r"""Builds basic version of Emformer-based :class:`~torchaudio.models.RNNT`.
+
+ Args:
+ num_symbols (int): The size of target token lexicon.
+
+ Returns:
+ RNNT:
+ Emformer RNN-T model.
+ """
+ return emformer_rnnt_model(
+ input_dim=80,
+ encoding_dim=1024,
+ num_symbols=num_symbols,
+ segment_length=16,
+ right_context_length=4,
+ time_reduction_input_dim=128,
+ time_reduction_stride=4,
+ transformer_num_heads=8,
+ transformer_ffn_dim=2048,
+ transformer_num_layers=20,
+ transformer_dropout=0.1,
+ transformer_activation="gelu",
+ transformer_left_context_length=30,
+ transformer_max_memory_size=0,
+ transformer_weight_init_scale_strategy="depthwise",
+ transformer_tanh_on_mem=True,
+ symbol_embedding_dim=512,
+ num_lstm_layers=3,
+ lstm_layer_norm=True,
+ lstm_layer_norm_epsilon=1e-3,
+ lstm_dropout=0.3,
+ )
diff --git a/lib/python3.12/site-packages/torchaudio/models/rnnt_decoder.py b/lib/python3.12/site-packages/torchaudio/models/rnnt_decoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..5a02b2ca907733a8e1ab404d1107bb702e977748
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/rnnt_decoder.py
@@ -0,0 +1,339 @@
+from typing import Callable, Dict, List, Optional, Tuple
+
+import torch
+from torchaudio.models import RNNT
+
+
+__all__ = ["Hypothesis", "RNNTBeamSearch"]
+
+
+Hypothesis = Tuple[List[int], torch.Tensor, List[List[torch.Tensor]], float]
+Hypothesis.__doc__ = """Hypothesis generated by RNN-T beam search decoder,
+ represented as tuple of (tokens, prediction network output, prediction network state, score).
+ """
+
+
+def _get_hypo_tokens(hypo: Hypothesis) -> List[int]:
+ return hypo[0]
+
+
+def _get_hypo_predictor_out(hypo: Hypothesis) -> torch.Tensor:
+ return hypo[1]
+
+
+def _get_hypo_state(hypo: Hypothesis) -> List[List[torch.Tensor]]:
+ return hypo[2]
+
+
+def _get_hypo_score(hypo: Hypothesis) -> float:
+ return hypo[3]
+
+
+def _get_hypo_key(hypo: Hypothesis) -> str:
+ return str(hypo[0])
+
+
+def _batch_state(hypos: List[Hypothesis]) -> List[List[torch.Tensor]]:
+ states: List[List[torch.Tensor]] = []
+ for i in range(len(_get_hypo_state(hypos[0]))):
+ batched_state_components: List[torch.Tensor] = []
+ for j in range(len(_get_hypo_state(hypos[0])[i])):
+ batched_state_components.append(torch.cat([_get_hypo_state(hypo)[i][j] for hypo in hypos]))
+ states.append(batched_state_components)
+ return states
+
+
+def _slice_state(states: List[List[torch.Tensor]], idx: int, device: torch.device) -> List[List[torch.Tensor]]:
+ idx_tensor = torch.tensor([idx], device=device)
+ return [[state.index_select(0, idx_tensor) for state in state_tuple] for state_tuple in states]
+
+
+def _default_hypo_sort_key(hypo: Hypothesis) -> float:
+ return _get_hypo_score(hypo) / (len(_get_hypo_tokens(hypo)) + 1)
+
+
+def _compute_updated_scores(
+ hypos: List[Hypothesis],
+ next_token_probs: torch.Tensor,
+ beam_width: int,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ hypo_scores = torch.tensor([_get_hypo_score(h) for h in hypos]).unsqueeze(1)
+ nonblank_scores = hypo_scores + next_token_probs[:, :-1] # [beam_width, num_tokens - 1]
+ nonblank_nbest_scores, nonblank_nbest_idx = nonblank_scores.reshape(-1).topk(beam_width)
+ nonblank_nbest_hypo_idx = nonblank_nbest_idx.div(nonblank_scores.shape[1], rounding_mode="trunc")
+ nonblank_nbest_token = nonblank_nbest_idx % nonblank_scores.shape[1]
+ return nonblank_nbest_scores, nonblank_nbest_hypo_idx, nonblank_nbest_token
+
+
+def _remove_hypo(hypo: Hypothesis, hypo_list: List[Hypothesis]) -> None:
+ for i, elem in enumerate(hypo_list):
+ if _get_hypo_key(hypo) == _get_hypo_key(elem):
+ del hypo_list[i]
+ break
+
+
+class RNNTBeamSearch(torch.nn.Module):
+ r"""Beam search decoder for RNN-T model.
+
+ See Also:
+ * :class:`torchaudio.pipelines.RNNTBundle`: ASR pipeline with pretrained model.
+
+ Args:
+ model (RNNT): RNN-T model to use.
+ blank (int): index of blank token in vocabulary.
+ temperature (float, optional): temperature to apply to joint network output.
+ Larger values yield more uniform samples. (Default: 1.0)
+ hypo_sort_key (Callable[[Hypothesis], float] or None, optional): callable that computes a score
+ for a given hypothesis to rank hypotheses by. If ``None``, defaults to callable that returns
+ hypothesis score normalized by token sequence length. (Default: None)
+ step_max_tokens (int, optional): maximum number of tokens to emit per input time step. (Default: 100)
+ """
+
+ def __init__(
+ self,
+ model: RNNT,
+ blank: int,
+ temperature: float = 1.0,
+ hypo_sort_key: Optional[Callable[[Hypothesis], float]] = None,
+ step_max_tokens: int = 100,
+ ) -> None:
+ super().__init__()
+ self.model = model
+ self.blank = blank
+ self.temperature = temperature
+
+ if hypo_sort_key is None:
+ self.hypo_sort_key = _default_hypo_sort_key
+ else:
+ self.hypo_sort_key = hypo_sort_key
+
+ self.step_max_tokens = step_max_tokens
+
+ def _init_b_hypos(self, device: torch.device) -> List[Hypothesis]:
+ token = self.blank
+ state = None
+
+ one_tensor = torch.tensor([1], device=device)
+ pred_out, _, pred_state = self.model.predict(torch.tensor([[token]], device=device), one_tensor, state)
+ init_hypo = (
+ [token],
+ pred_out[0].detach(),
+ pred_state,
+ 0.0,
+ )
+ return [init_hypo]
+
+ def _gen_next_token_probs(
+ self, enc_out: torch.Tensor, hypos: List[Hypothesis], device: torch.device
+ ) -> torch.Tensor:
+ one_tensor = torch.tensor([1], device=device)
+ predictor_out = torch.stack([_get_hypo_predictor_out(h) for h in hypos], dim=0)
+ joined_out, _, _ = self.model.join(
+ enc_out,
+ one_tensor,
+ predictor_out,
+ torch.tensor([1] * len(hypos), device=device),
+ ) # [beam_width, 1, 1, num_tokens]
+ joined_out = torch.nn.functional.log_softmax(joined_out / self.temperature, dim=3)
+ return joined_out[:, 0, 0]
+
+ def _gen_b_hypos(
+ self,
+ b_hypos: List[Hypothesis],
+ a_hypos: List[Hypothesis],
+ next_token_probs: torch.Tensor,
+ key_to_b_hypo: Dict[str, Hypothesis],
+ ) -> List[Hypothesis]:
+ for i in range(len(a_hypos)):
+ h_a = a_hypos[i]
+ append_blank_score = _get_hypo_score(h_a) + next_token_probs[i, -1]
+ if _get_hypo_key(h_a) in key_to_b_hypo:
+ h_b = key_to_b_hypo[_get_hypo_key(h_a)]
+ _remove_hypo(h_b, b_hypos)
+ score = float(torch.tensor(_get_hypo_score(h_b)).logaddexp(append_blank_score))
+ else:
+ score = float(append_blank_score)
+ h_b = (
+ _get_hypo_tokens(h_a),
+ _get_hypo_predictor_out(h_a),
+ _get_hypo_state(h_a),
+ score,
+ )
+ b_hypos.append(h_b)
+ key_to_b_hypo[_get_hypo_key(h_b)] = h_b
+ _, sorted_idx = torch.tensor([_get_hypo_score(hypo) for hypo in b_hypos]).sort()
+ return [b_hypos[idx] for idx in sorted_idx]
+
+ def _gen_a_hypos(
+ self,
+ a_hypos: List[Hypothesis],
+ b_hypos: List[Hypothesis],
+ next_token_probs: torch.Tensor,
+ t: int,
+ beam_width: int,
+ device: torch.device,
+ ) -> List[Hypothesis]:
+ (
+ nonblank_nbest_scores,
+ nonblank_nbest_hypo_idx,
+ nonblank_nbest_token,
+ ) = _compute_updated_scores(a_hypos, next_token_probs, beam_width)
+
+ if len(b_hypos) < beam_width:
+ b_nbest_score = -float("inf")
+ else:
+ b_nbest_score = _get_hypo_score(b_hypos[-beam_width])
+
+ base_hypos: List[Hypothesis] = []
+ new_tokens: List[int] = []
+ new_scores: List[float] = []
+ for i in range(beam_width):
+ score = float(nonblank_nbest_scores[i])
+ if score > b_nbest_score:
+ a_hypo_idx = int(nonblank_nbest_hypo_idx[i])
+ base_hypos.append(a_hypos[a_hypo_idx])
+ new_tokens.append(int(nonblank_nbest_token[i]))
+ new_scores.append(score)
+
+ if base_hypos:
+ new_hypos = self._gen_new_hypos(base_hypos, new_tokens, new_scores, t, device)
+ else:
+ new_hypos: List[Hypothesis] = []
+
+ return new_hypos
+
+ def _gen_new_hypos(
+ self,
+ base_hypos: List[Hypothesis],
+ tokens: List[int],
+ scores: List[float],
+ t: int,
+ device: torch.device,
+ ) -> List[Hypothesis]:
+ tgt_tokens = torch.tensor([[token] for token in tokens], device=device)
+ states = _batch_state(base_hypos)
+ pred_out, _, pred_states = self.model.predict(
+ tgt_tokens,
+ torch.tensor([1] * len(base_hypos), device=device),
+ states,
+ )
+ new_hypos: List[Hypothesis] = []
+ for i, h_a in enumerate(base_hypos):
+ new_tokens = _get_hypo_tokens(h_a) + [tokens[i]]
+ new_hypos.append((new_tokens, pred_out[i].detach(), _slice_state(pred_states, i, device), scores[i]))
+ return new_hypos
+
+ def _search(
+ self,
+ enc_out: torch.Tensor,
+ hypo: Optional[List[Hypothesis]],
+ beam_width: int,
+ ) -> List[Hypothesis]:
+ n_time_steps = enc_out.shape[1]
+ device = enc_out.device
+
+ a_hypos: List[Hypothesis] = []
+ b_hypos = self._init_b_hypos(device) if hypo is None else hypo
+ for t in range(n_time_steps):
+ a_hypos = b_hypos
+ b_hypos = torch.jit.annotate(List[Hypothesis], [])
+ key_to_b_hypo: Dict[str, Hypothesis] = {}
+ symbols_current_t = 0
+
+ while a_hypos:
+ next_token_probs = self._gen_next_token_probs(enc_out[:, t : t + 1], a_hypos, device)
+ next_token_probs = next_token_probs.cpu()
+ b_hypos = self._gen_b_hypos(b_hypos, a_hypos, next_token_probs, key_to_b_hypo)
+
+ if symbols_current_t == self.step_max_tokens:
+ break
+
+ a_hypos = self._gen_a_hypos(
+ a_hypos,
+ b_hypos,
+ next_token_probs,
+ t,
+ beam_width,
+ device,
+ )
+ if a_hypos:
+ symbols_current_t += 1
+
+ _, sorted_idx = torch.tensor([self.hypo_sort_key(hyp) for hyp in b_hypos]).topk(beam_width)
+ b_hypos = [b_hypos[idx] for idx in sorted_idx]
+
+ return b_hypos
+
+ def forward(self, input: torch.Tensor, length: torch.Tensor, beam_width: int) -> List[Hypothesis]:
+ r"""Performs beam search for the given input sequence.
+
+ T: number of frames;
+ D: feature dimension of each frame.
+
+ Args:
+ input (torch.Tensor): sequence of input frames, with shape (T, D) or (1, T, D).
+ length (torch.Tensor): number of valid frames in input
+ sequence, with shape () or (1,).
+ beam_width (int): beam size to use during search.
+
+ Returns:
+ List[Hypothesis]: top-``beam_width`` hypotheses found by beam search.
+ """
+ if input.dim() != 2 and not (input.dim() == 3 and input.shape[0] == 1):
+ raise ValueError("input must be of shape (T, D) or (1, T, D)")
+ if input.dim() == 2:
+ input = input.unsqueeze(0)
+
+ if length.shape != () and length.shape != (1,):
+ raise ValueError("length must be of shape () or (1,)")
+ if length.dim() == 0:
+ length = length.unsqueeze(0)
+
+ enc_out, _ = self.model.transcribe(input, length)
+ return self._search(enc_out, None, beam_width)
+
+ @torch.jit.export
+ def infer(
+ self,
+ input: torch.Tensor,
+ length: torch.Tensor,
+ beam_width: int,
+ state: Optional[List[List[torch.Tensor]]] = None,
+ hypothesis: Optional[List[Hypothesis]] = None,
+ ) -> Tuple[List[Hypothesis], List[List[torch.Tensor]]]:
+ r"""Performs beam search for the given input sequence in streaming mode.
+
+ T: number of frames;
+ D: feature dimension of each frame.
+
+ Args:
+ input (torch.Tensor): sequence of input frames, with shape (T, D) or (1, T, D).
+ length (torch.Tensor): number of valid frames in input
+ sequence, with shape () or (1,).
+ beam_width (int): beam size to use during search.
+ state (List[List[torch.Tensor]] or None, optional): list of lists of tensors
+ representing transcription network internal state generated in preceding
+ invocation. (Default: ``None``)
+ hypothesis (List[Hypothesis] or None): hypotheses from preceding invocation to seed
+ search with. (Default: ``None``)
+
+ Returns:
+ (List[Hypothesis], List[List[torch.Tensor]]):
+ List[Hypothesis]
+ top-``beam_width`` hypotheses found by beam search.
+ List[List[torch.Tensor]]
+ list of lists of tensors representing transcription network
+ internal state generated in current invocation.
+ """
+ if input.dim() != 2 and not (input.dim() == 3 and input.shape[0] == 1):
+ raise ValueError("input must be of shape (T, D) or (1, T, D)")
+ if input.dim() == 2:
+ input = input.unsqueeze(0)
+
+ if length.shape != () and length.shape != (1,):
+ raise ValueError("length must be of shape () or (1,)")
+ if length.dim() == 0:
+ length = length.unsqueeze(0)
+
+ enc_out, _, state = self.model.transcribe_streaming(input, length, state)
+ return self._search(enc_out, hypothesis, beam_width), state
diff --git a/lib/python3.12/site-packages/torchaudio/models/squim/__init__.py b/lib/python3.12/site-packages/torchaudio/models/squim/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..092d6eb8e36e2329c78d21bf609a8458818995e6
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/squim/__init__.py
@@ -0,0 +1,11 @@
+from .objective import squim_objective_base, squim_objective_model, SquimObjective
+from .subjective import squim_subjective_base, squim_subjective_model, SquimSubjective
+
+__all__ = [
+ "squim_objective_base",
+ "squim_objective_model",
+ "squim_subjective_base",
+ "squim_subjective_model",
+ "SquimObjective",
+ "SquimSubjective",
+]
diff --git a/lib/python3.12/site-packages/torchaudio/models/squim/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/squim/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f5da9332fff4be0a5e1ec62543cc69c6f3e24099
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/squim/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/squim/__pycache__/objective.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/squim/__pycache__/objective.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d81c3265506d20f39878a9f1ef41a78c9eb11dbb
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/squim/__pycache__/objective.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/squim/__pycache__/subjective.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/squim/__pycache__/subjective.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3268efddd4c476287614f62c8fb7ef2d8815b908
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/squim/__pycache__/subjective.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/squim/objective.py b/lib/python3.12/site-packages/torchaudio/models/squim/objective.py
new file mode 100644
index 0000000000000000000000000000000000000000..d0a0671a4ec3828fa5dbe1613c7378fbc55d34dd
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/squim/objective.py
@@ -0,0 +1,326 @@
+import math
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def transform_wb_pesq_range(x: float) -> float:
+ """The metric defined by ITU-T P.862 is often called 'PESQ score', which is defined
+ for narrow-band signals and has a value range of [-0.5, 4.5] exactly. Here, we use the metric
+ defined by ITU-T P.862.2, commonly known as 'wide-band PESQ' and will be referred to as "PESQ score".
+
+ Args:
+ x (float): Narrow-band PESQ score.
+
+ Returns:
+ (float): Wide-band PESQ score.
+ """
+ return 0.999 + (4.999 - 0.999) / (1 + math.exp(-1.3669 * x + 3.8224))
+
+
+PESQRange: Tuple[float, float] = (
+ 1.0, # P.862.2 uses a different input filter than P.862, and the lower bound of
+ # the raw score is not -0.5 anymore. It's hard to figure out the true lower bound.
+ # We are using 1.0 as a reasonable approximation.
+ transform_wb_pesq_range(4.5),
+)
+
+
+class RangeSigmoid(nn.Module):
+ def __init__(self, val_range: Tuple[float, float] = (0.0, 1.0)) -> None:
+ super(RangeSigmoid, self).__init__()
+ assert isinstance(val_range, tuple) and len(val_range) == 2
+ self.val_range: Tuple[float, float] = val_range
+ self.sigmoid: nn.modules.Module = nn.Sigmoid()
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ out = self.sigmoid(x) * (self.val_range[1] - self.val_range[0]) + self.val_range[0]
+ return out
+
+
+class Encoder(nn.Module):
+ """Encoder module that transform 1D waveform to 2D representations.
+
+ Args:
+ feat_dim (int, optional): The feature dimension after Encoder module. (Default: 512)
+ win_len (int, optional): kernel size in the Conv1D layer. (Default: 32)
+ """
+
+ def __init__(self, feat_dim: int = 512, win_len: int = 32) -> None:
+ super(Encoder, self).__init__()
+
+ self.conv1d = nn.Conv1d(1, feat_dim, win_len, stride=win_len // 2, bias=False)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """Apply waveforms to convolutional layer and ReLU layer.
+
+ Args:
+ x (torch.Tensor): Input waveforms. Tensor with dimensions `(batch, time)`.
+
+ Returns:
+ (torch,Tensor): Feature Tensor with dimensions `(batch, channel, frame)`.
+ """
+ out = x.unsqueeze(dim=1)
+ out = F.relu(self.conv1d(out))
+ return out
+
+
+class SingleRNN(nn.Module):
+ def __init__(self, rnn_type: str, input_size: int, hidden_size: int, dropout: float = 0.0) -> None:
+ super(SingleRNN, self).__init__()
+
+ self.rnn_type = rnn_type
+ self.input_size = input_size
+ self.hidden_size = hidden_size
+
+ self.rnn: nn.modules.Module = getattr(nn, rnn_type)(
+ input_size,
+ hidden_size,
+ 1,
+ dropout=dropout,
+ batch_first=True,
+ bidirectional=True,
+ )
+
+ self.proj = nn.Linear(hidden_size * 2, input_size)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ # input shape: batch, seq, dim
+ out, _ = self.rnn(x)
+ out = self.proj(out)
+ return out
+
+
+class DPRNN(nn.Module):
+ """*Dual-path recurrent neural networks (DPRNN)* :cite:`luo2020dual`.
+
+ Args:
+ feat_dim (int, optional): The feature dimension after Encoder module. (Default: 64)
+ hidden_dim (int, optional): Hidden dimension in the RNN layer of DPRNN. (Default: 128)
+ num_blocks (int, optional): Number of DPRNN layers. (Default: 6)
+ rnn_type (str, optional): Type of RNN in DPRNN. Valid options are ["RNN", "LSTM", "GRU"]. (Default: "LSTM")
+ d_model (int, optional): The number of expected features in the input. (Default: 256)
+ chunk_size (int, optional): Chunk size of input for DPRNN. (Default: 100)
+ chunk_stride (int, optional): Stride of chunk input for DPRNN. (Default: 50)
+ """
+
+ def __init__(
+ self,
+ feat_dim: int = 64,
+ hidden_dim: int = 128,
+ num_blocks: int = 6,
+ rnn_type: str = "LSTM",
+ d_model: int = 256,
+ chunk_size: int = 100,
+ chunk_stride: int = 50,
+ ) -> None:
+ super(DPRNN, self).__init__()
+
+ self.num_blocks = num_blocks
+
+ self.row_rnn = nn.ModuleList([])
+ self.col_rnn = nn.ModuleList([])
+ self.row_norm = nn.ModuleList([])
+ self.col_norm = nn.ModuleList([])
+ for _ in range(num_blocks):
+ self.row_rnn.append(SingleRNN(rnn_type, feat_dim, hidden_dim))
+ self.col_rnn.append(SingleRNN(rnn_type, feat_dim, hidden_dim))
+ self.row_norm.append(nn.GroupNorm(1, feat_dim, eps=1e-8))
+ self.col_norm.append(nn.GroupNorm(1, feat_dim, eps=1e-8))
+ self.conv = nn.Sequential(
+ nn.Conv2d(feat_dim, d_model, 1),
+ nn.PReLU(),
+ )
+ self.chunk_size = chunk_size
+ self.chunk_stride = chunk_stride
+
+ def pad_chunk(self, x: torch.Tensor) -> Tuple[torch.Tensor, int]:
+ # input shape: (B, N, T)
+ seq_len = x.shape[-1]
+
+ rest = self.chunk_size - (self.chunk_stride + seq_len % self.chunk_size) % self.chunk_size
+ out = F.pad(x, [self.chunk_stride, rest + self.chunk_stride])
+
+ return out, rest
+
+ def chunking(self, x: torch.Tensor) -> Tuple[torch.Tensor, int]:
+ out, rest = self.pad_chunk(x)
+ batch_size, feat_dim, seq_len = out.shape
+
+ segments1 = out[:, :, : -self.chunk_stride].contiguous().view(batch_size, feat_dim, -1, self.chunk_size)
+ segments2 = out[:, :, self.chunk_stride :].contiguous().view(batch_size, feat_dim, -1, self.chunk_size)
+ out = torch.cat([segments1, segments2], dim=3)
+ out = out.view(batch_size, feat_dim, -1, self.chunk_size).transpose(2, 3).contiguous()
+
+ return out, rest
+
+ def merging(self, x: torch.Tensor, rest: int) -> torch.Tensor:
+ batch_size, dim, _, _ = x.shape
+ out = x.transpose(2, 3).contiguous().view(batch_size, dim, -1, self.chunk_size * 2)
+ out1 = out[:, :, :, : self.chunk_size].contiguous().view(batch_size, dim, -1)[:, :, self.chunk_stride :]
+ out2 = out[:, :, :, self.chunk_size :].contiguous().view(batch_size, dim, -1)[:, :, : -self.chunk_stride]
+ out = out1 + out2
+ if rest > 0:
+ out = out[:, :, :-rest]
+ out = out.contiguous()
+ return out
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x, rest = self.chunking(x)
+ batch_size, _, dim1, dim2 = x.shape
+ out = x
+ for row_rnn, row_norm, col_rnn, col_norm in zip(self.row_rnn, self.row_norm, self.col_rnn, self.col_norm):
+ row_in = out.permute(0, 3, 2, 1).contiguous().view(batch_size * dim2, dim1, -1).contiguous()
+ row_out = row_rnn(row_in)
+ row_out = row_out.view(batch_size, dim2, dim1, -1).permute(0, 3, 2, 1).contiguous()
+ row_out = row_norm(row_out)
+ out = out + row_out
+
+ col_in = out.permute(0, 2, 3, 1).contiguous().view(batch_size * dim1, dim2, -1).contiguous()
+ col_out = col_rnn(col_in)
+ col_out = col_out.view(batch_size, dim1, dim2, -1).permute(0, 3, 1, 2).contiguous()
+ col_out = col_norm(col_out)
+ out = out + col_out
+ out = self.conv(out)
+ out = self.merging(out, rest)
+ out = out.transpose(1, 2).contiguous()
+ return out
+
+
+class AutoPool(nn.Module):
+ def __init__(self, pool_dim: int = 1) -> None:
+ super(AutoPool, self).__init__()
+ self.pool_dim: int = pool_dim
+ self.softmax: nn.modules.Module = nn.Softmax(dim=pool_dim)
+ self.register_parameter("alpha", nn.Parameter(torch.ones(1)))
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ weight = self.softmax(torch.mul(x, self.alpha))
+ out = torch.sum(torch.mul(x, weight), dim=self.pool_dim)
+ return out
+
+
+class SquimObjective(nn.Module):
+ """Speech Quality and Intelligibility Measures (SQUIM) model that predicts **objective** metric scores
+ for speech enhancement (e.g., STOI, PESQ, and SI-SDR).
+
+ Args:
+ encoder (torch.nn.Module): Encoder module to transform 1D waveform to 2D feature representation.
+ dprnn (torch.nn.Module): DPRNN module to model sequential feature.
+ branches (torch.nn.ModuleList): Transformer branches in which each branch estimate one objective metirc score.
+ """
+
+ def __init__(
+ self,
+ encoder: nn.Module,
+ dprnn: nn.Module,
+ branches: nn.ModuleList,
+ ):
+ super(SquimObjective, self).__init__()
+ self.encoder = encoder
+ self.dprnn = dprnn
+ self.branches = branches
+
+ def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+ """
+ Args:
+ x (torch.Tensor): Input waveforms. Tensor with dimensions `(batch, time)`.
+
+ Returns:
+ List(torch.Tensor): List of score Tenosrs. Each Tensor is with dimension `(batch,)`.
+ """
+ if x.ndim != 2:
+ raise ValueError(f"The input must be a 2D Tensor. Found dimension {x.ndim}.")
+ x = x / (torch.mean(x**2, dim=1, keepdim=True) ** 0.5 * 20)
+ out = self.encoder(x)
+ out = self.dprnn(out)
+ scores = []
+ for branch in self.branches:
+ scores.append(branch(out).squeeze(dim=1))
+ return scores
+
+
+def _create_branch(d_model: int, nhead: int, metric: str) -> nn.modules.Module:
+ """Create branch module after DPRNN model for predicting metric score.
+
+ Args:
+ d_model (int): The number of expected features in the input.
+ nhead (int): Number of heads in the multi-head attention model.
+ metric (str): The metric name to predict.
+
+ Returns:
+ (nn.Module): Returned module to predict corresponding metric score.
+ """
+ layer1 = nn.TransformerEncoderLayer(d_model, nhead, d_model * 4, dropout=0.0, batch_first=True)
+ layer2 = AutoPool()
+ if metric == "stoi":
+ layer3 = nn.Sequential(
+ nn.Linear(d_model, d_model),
+ nn.PReLU(),
+ nn.Linear(d_model, 1),
+ RangeSigmoid(),
+ )
+ elif metric == "pesq":
+ layer3 = nn.Sequential(
+ nn.Linear(d_model, d_model),
+ nn.PReLU(),
+ nn.Linear(d_model, 1),
+ RangeSigmoid(val_range=PESQRange),
+ )
+ else:
+ layer3: nn.modules.Module = nn.Sequential(nn.Linear(d_model, d_model), nn.PReLU(), nn.Linear(d_model, 1))
+ return nn.Sequential(layer1, layer2, layer3)
+
+
+def squim_objective_model(
+ feat_dim: int,
+ win_len: int,
+ d_model: int,
+ nhead: int,
+ hidden_dim: int,
+ num_blocks: int,
+ rnn_type: str,
+ chunk_size: int,
+ chunk_stride: Optional[int] = None,
+) -> SquimObjective:
+ """Build a custome :class:`torchaudio.prototype.models.SquimObjective` model.
+
+ Args:
+ feat_dim (int, optional): The feature dimension after Encoder module.
+ win_len (int): Kernel size in the Encoder module.
+ d_model (int): The number of expected features in the input.
+ nhead (int): Number of heads in the multi-head attention model.
+ hidden_dim (int): Hidden dimension in the RNN layer of DPRNN.
+ num_blocks (int): Number of DPRNN layers.
+ rnn_type (str): Type of RNN in DPRNN. Valid options are ["RNN", "LSTM", "GRU"].
+ chunk_size (int): Chunk size of input for DPRNN.
+ chunk_stride (int or None, optional): Stride of chunk input for DPRNN.
+ """
+ if chunk_stride is None:
+ chunk_stride = chunk_size // 2
+ encoder = Encoder(feat_dim, win_len)
+ dprnn = DPRNN(feat_dim, hidden_dim, num_blocks, rnn_type, d_model, chunk_size, chunk_stride)
+ branches = nn.ModuleList(
+ [
+ _create_branch(d_model, nhead, "stoi"),
+ _create_branch(d_model, nhead, "pesq"),
+ _create_branch(d_model, nhead, "sisdr"),
+ ]
+ )
+ return SquimObjective(encoder, dprnn, branches)
+
+
+def squim_objective_base() -> SquimObjective:
+ """Build :class:`torchaudio.prototype.models.SquimObjective` model with default arguments."""
+ return squim_objective_model(
+ feat_dim=256,
+ win_len=64,
+ d_model=256,
+ nhead=4,
+ hidden_dim=256,
+ num_blocks=2,
+ rnn_type="LSTM",
+ chunk_size=71,
+ )
diff --git a/lib/python3.12/site-packages/torchaudio/models/squim/subjective.py b/lib/python3.12/site-packages/torchaudio/models/squim/subjective.py
new file mode 100644
index 0000000000000000000000000000000000000000..4be681c91c5f67a2b888b49ec8269b74762360ab
--- /dev/null
+++ b/lib/python3.12/site-packages/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,
+ )
diff --git a/lib/python3.12/site-packages/torchaudio/models/tacotron2.py b/lib/python3.12/site-packages/torchaudio/models/tacotron2.py
new file mode 100644
index 0000000000000000000000000000000000000000..978fb97c88db9c64a9b216a340e63075e53e2295
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/tacotron2.py
@@ -0,0 +1,1046 @@
+# *****************************************************************************
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+# * Redistributions of source code must retain the above copyright
+# notice, this list of conditions and the following disclaimer.
+# * Redistributions in binary form must reproduce the above copyright
+# notice, this list of conditions and the following disclaimer in the
+# documentation and/or other materials provided with the distribution.
+# * Neither the name of the NVIDIA CORPORATION nor the
+# names of its contributors may be used to endorse or promote products
+# derived from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# *****************************************************************************
+
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+from torch import nn, Tensor
+from torch.nn import functional as F
+
+
+__all__ = [
+ "Tacotron2",
+]
+
+
+def _get_linear_layer(in_dim: int, out_dim: int, bias: bool = True, w_init_gain: str = "linear") -> torch.nn.Linear:
+ r"""Linear layer with xavier uniform initialization.
+
+ Args:
+ in_dim (int): Size of each input sample.
+ out_dim (int): Size of each output sample.
+ bias (bool, optional): If set to ``False``, the layer will not learn an additive bias. (Default: ``True``)
+ w_init_gain (str, optional): Parameter passed to ``torch.nn.init.calculate_gain``
+ for setting the gain parameter of ``xavier_uniform_``. (Default: ``linear``)
+
+ Returns:
+ (torch.nn.Linear): The corresponding linear layer.
+ """
+ linear = torch.nn.Linear(in_dim, out_dim, bias=bias)
+ torch.nn.init.xavier_uniform_(linear.weight, gain=torch.nn.init.calculate_gain(w_init_gain))
+ return linear
+
+
+def _get_conv1d_layer(
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int = 1,
+ stride: int = 1,
+ padding: Optional[Union[str, int, Tuple[int]]] = None,
+ dilation: int = 1,
+ bias: bool = True,
+ w_init_gain: str = "linear",
+) -> torch.nn.Conv1d:
+ r"""1D convolution with xavier uniform initialization.
+
+ Args:
+ in_channels (int): Number of channels in the input image.
+ out_channels (int): Number of channels produced by the convolution.
+ kernel_size (int, optional): Number of channels in the input image. (Default: ``1``)
+ stride (int, optional): Number of channels in the input image. (Default: ``1``)
+ padding (str, int or tuple, optional): Padding added to both sides of the input.
+ (Default: dilation * (kernel_size - 1) / 2)
+ dilation (int, optional): Number of channels in the input image. (Default: ``1``)
+ w_init_gain (str, optional): Parameter passed to ``torch.nn.init.calculate_gain``
+ for setting the gain parameter of ``xavier_uniform_``. (Default: ``linear``)
+
+ Returns:
+ (torch.nn.Conv1d): The corresponding Conv1D layer.
+ """
+ if padding is None:
+ if kernel_size % 2 != 1:
+ raise ValueError("kernel_size must be odd")
+ padding = int(dilation * (kernel_size - 1) / 2)
+
+ conv1d = torch.nn.Conv1d(
+ in_channels,
+ out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ bias=bias,
+ )
+
+ torch.nn.init.xavier_uniform_(conv1d.weight, gain=torch.nn.init.calculate_gain(w_init_gain))
+
+ return conv1d
+
+
+def _get_mask_from_lengths(lengths: Tensor) -> Tensor:
+ r"""Returns a binary mask based on ``lengths``. The ``i``-th row and ``j``-th column of the mask
+ is ``1`` if ``j`` is smaller than ``i``-th element of ``lengths.
+
+ Args:
+ lengths (Tensor): The length of each element in the batch, with shape (n_batch, ).
+
+ Returns:
+ mask (Tensor): The binary mask, with shape (n_batch, max of ``lengths``).
+ """
+ max_len = torch.max(lengths).item()
+ ids = torch.arange(0, max_len, device=lengths.device, dtype=lengths.dtype)
+ mask = (ids < lengths.unsqueeze(1)).byte()
+ mask = torch.le(mask, 0)
+ return mask
+
+
+class _LocationLayer(nn.Module):
+ r"""Location layer used in the Attention model.
+
+ Args:
+ attention_n_filter (int): Number of filters for attention model.
+ attention_kernel_size (int): Kernel size for attention model.
+ attention_hidden_dim (int): Dimension of attention hidden representation.
+ """
+
+ def __init__(
+ self,
+ attention_n_filter: int,
+ attention_kernel_size: int,
+ attention_hidden_dim: int,
+ ):
+ super().__init__()
+ padding = int((attention_kernel_size - 1) / 2)
+ self.location_conv = _get_conv1d_layer(
+ 2,
+ attention_n_filter,
+ kernel_size=attention_kernel_size,
+ padding=padding,
+ bias=False,
+ stride=1,
+ dilation=1,
+ )
+ self.location_dense = _get_linear_layer(
+ attention_n_filter, attention_hidden_dim, bias=False, w_init_gain="tanh"
+ )
+
+ def forward(self, attention_weights_cat: Tensor) -> Tensor:
+ r"""Location layer used in the Attention model.
+
+ Args:
+ attention_weights_cat (Tensor): Cumulative and previous attention weights
+ with shape (n_batch, 2, max of ``text_lengths``).
+
+ Returns:
+ processed_attention (Tensor): Cumulative and previous attention weights
+ with shape (n_batch, ``attention_hidden_dim``).
+ """
+ # (n_batch, attention_n_filter, text_lengths.max())
+ processed_attention = self.location_conv(attention_weights_cat)
+ processed_attention = processed_attention.transpose(1, 2)
+ # (n_batch, text_lengths.max(), attention_hidden_dim)
+ processed_attention = self.location_dense(processed_attention)
+ return processed_attention
+
+
+class _Attention(nn.Module):
+ r"""Locally sensitive attention model.
+
+ Args:
+ attention_rnn_dim (int): Number of hidden units for RNN.
+ encoder_embedding_dim (int): Number of embedding dimensions in the Encoder.
+ attention_hidden_dim (int): Dimension of attention hidden representation.
+ attention_location_n_filter (int): Number of filters for Attention model.
+ attention_location_kernel_size (int): Kernel size for Attention model.
+ """
+
+ def __init__(
+ self,
+ attention_rnn_dim: int,
+ encoder_embedding_dim: int,
+ attention_hidden_dim: int,
+ attention_location_n_filter: int,
+ attention_location_kernel_size: int,
+ ) -> None:
+ super().__init__()
+ self.query_layer = _get_linear_layer(attention_rnn_dim, attention_hidden_dim, bias=False, w_init_gain="tanh")
+ self.memory_layer = _get_linear_layer(
+ encoder_embedding_dim, attention_hidden_dim, bias=False, w_init_gain="tanh"
+ )
+ self.v = _get_linear_layer(attention_hidden_dim, 1, bias=False)
+ self.location_layer = _LocationLayer(
+ attention_location_n_filter,
+ attention_location_kernel_size,
+ attention_hidden_dim,
+ )
+ self.score_mask_value = -float("inf")
+
+ def _get_alignment_energies(self, query: Tensor, processed_memory: Tensor, attention_weights_cat: Tensor) -> Tensor:
+ r"""Get the alignment vector.
+
+ Args:
+ query (Tensor): Decoder output with shape (n_batch, n_mels * n_frames_per_step).
+ processed_memory (Tensor): Processed Encoder outputs
+ with shape (n_batch, max of ``text_lengths``, attention_hidden_dim).
+ attention_weights_cat (Tensor): Cumulative and previous attention weights
+ with shape (n_batch, 2, max of ``text_lengths``).
+
+ Returns:
+ alignment (Tensor): attention weights, it is a tensor with shape (batch, max of ``text_lengths``).
+ """
+
+ processed_query = self.query_layer(query.unsqueeze(1))
+ processed_attention_weights = self.location_layer(attention_weights_cat)
+ energies = self.v(torch.tanh(processed_query + processed_attention_weights + processed_memory))
+
+ alignment = energies.squeeze(2)
+ return alignment
+
+ def forward(
+ self,
+ attention_hidden_state: Tensor,
+ memory: Tensor,
+ processed_memory: Tensor,
+ attention_weights_cat: Tensor,
+ mask: Tensor,
+ ) -> Tuple[Tensor, Tensor]:
+ r"""Pass the input through the Attention model.
+
+ Args:
+ attention_hidden_state (Tensor): Attention rnn last output with shape (n_batch, ``attention_rnn_dim``).
+ memory (Tensor): Encoder outputs with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``).
+ processed_memory (Tensor): Processed Encoder outputs
+ with shape (n_batch, max of ``text_lengths``, ``attention_hidden_dim``).
+ attention_weights_cat (Tensor): Previous and cumulative attention weights
+ with shape (n_batch, current_num_frames * 2, max of ``text_lengths``).
+ mask (Tensor): Binary mask for padded data with shape (n_batch, current_num_frames).
+
+ Returns:
+ attention_context (Tensor): Context vector with shape (n_batch, ``encoder_embedding_dim``).
+ attention_weights (Tensor): Attention weights with shape (n_batch, max of ``text_lengths``).
+ """
+ alignment = self._get_alignment_energies(attention_hidden_state, processed_memory, attention_weights_cat)
+
+ alignment = alignment.masked_fill(mask, self.score_mask_value)
+
+ attention_weights = F.softmax(alignment, dim=1)
+ attention_context = torch.bmm(attention_weights.unsqueeze(1), memory)
+ attention_context = attention_context.squeeze(1)
+
+ return attention_context, attention_weights
+
+
+class _Prenet(nn.Module):
+ r"""Prenet Module. It is consists of ``len(output_size)`` linear layers.
+
+ Args:
+ in_dim (int): The size of each input sample.
+ output_sizes (list): The output dimension of each linear layers.
+ """
+
+ def __init__(self, in_dim: int, out_sizes: List[int]) -> None:
+ super().__init__()
+ in_sizes = [in_dim] + out_sizes[:-1]
+ self.layers = nn.ModuleList(
+ [_get_linear_layer(in_size, out_size, bias=False) for (in_size, out_size) in zip(in_sizes, out_sizes)]
+ )
+
+ def forward(self, x: Tensor) -> Tensor:
+ r"""Pass the input through Prenet.
+
+ Args:
+ x (Tensor): The input sequence to Prenet with shape (n_batch, in_dim).
+
+ Return:
+ x (Tensor): Tensor with shape (n_batch, sizes[-1])
+ """
+
+ for linear in self.layers:
+ x = F.dropout(F.relu(linear(x)), p=0.5, training=True)
+ return x
+
+
+class _Postnet(nn.Module):
+ r"""Postnet Module.
+
+ Args:
+ n_mels (int): Number of mel bins.
+ postnet_embedding_dim (int): Postnet embedding dimension.
+ postnet_kernel_size (int): Postnet kernel size.
+ postnet_n_convolution (int): Number of postnet convolutions.
+ """
+
+ def __init__(
+ self,
+ n_mels: int,
+ postnet_embedding_dim: int,
+ postnet_kernel_size: int,
+ postnet_n_convolution: int,
+ ):
+ super().__init__()
+ self.convolutions = nn.ModuleList()
+
+ for i in range(postnet_n_convolution):
+ in_channels = n_mels if i == 0 else postnet_embedding_dim
+ out_channels = n_mels if i == (postnet_n_convolution - 1) else postnet_embedding_dim
+ init_gain = "linear" if i == (postnet_n_convolution - 1) else "tanh"
+ num_features = n_mels if i == (postnet_n_convolution - 1) else postnet_embedding_dim
+ self.convolutions.append(
+ nn.Sequential(
+ _get_conv1d_layer(
+ in_channels,
+ out_channels,
+ kernel_size=postnet_kernel_size,
+ stride=1,
+ padding=int((postnet_kernel_size - 1) / 2),
+ dilation=1,
+ w_init_gain=init_gain,
+ ),
+ nn.BatchNorm1d(num_features),
+ )
+ )
+
+ self.n_convs = len(self.convolutions)
+
+ def forward(self, x: Tensor) -> Tensor:
+ r"""Pass the input through Postnet.
+
+ Args:
+ x (Tensor): The input sequence with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``).
+
+ Return:
+ x (Tensor): Tensor with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``).
+ """
+
+ for i, conv in enumerate(self.convolutions):
+ if i < self.n_convs - 1:
+ x = F.dropout(torch.tanh(conv(x)), 0.5, training=self.training)
+ else:
+ x = F.dropout(conv(x), 0.5, training=self.training)
+
+ return x
+
+
+class _Encoder(nn.Module):
+ r"""Encoder Module.
+
+ Args:
+ encoder_embedding_dim (int): Number of embedding dimensions in the encoder.
+ encoder_n_convolution (int): Number of convolution layers in the encoder.
+ encoder_kernel_size (int): The kernel size in the encoder.
+
+ Examples
+ >>> encoder = _Encoder(3, 512, 5)
+ >>> input = torch.rand(10, 20, 30)
+ >>> output = encoder(input) # shape: (10, 30, 512)
+ """
+
+ def __init__(
+ self,
+ encoder_embedding_dim: int,
+ encoder_n_convolution: int,
+ encoder_kernel_size: int,
+ ) -> None:
+ super().__init__()
+
+ self.convolutions = nn.ModuleList()
+ for _ in range(encoder_n_convolution):
+ conv_layer = nn.Sequential(
+ _get_conv1d_layer(
+ encoder_embedding_dim,
+ encoder_embedding_dim,
+ kernel_size=encoder_kernel_size,
+ stride=1,
+ padding=int((encoder_kernel_size - 1) / 2),
+ dilation=1,
+ w_init_gain="relu",
+ ),
+ nn.BatchNorm1d(encoder_embedding_dim),
+ )
+ self.convolutions.append(conv_layer)
+
+ self.lstm = nn.LSTM(
+ encoder_embedding_dim,
+ int(encoder_embedding_dim / 2),
+ 1,
+ batch_first=True,
+ bidirectional=True,
+ )
+ self.lstm.flatten_parameters()
+
+ def forward(self, x: Tensor, input_lengths: Tensor) -> Tensor:
+ r"""Pass the input through the Encoder.
+
+ Args:
+ x (Tensor): The input sequences with shape (n_batch, encoder_embedding_dim, n_seq).
+ input_lengths (Tensor): The length of each input sequence with shape (n_batch, ).
+
+ Return:
+ x (Tensor): A tensor with shape (n_batch, n_seq, encoder_embedding_dim).
+ """
+
+ for conv in self.convolutions:
+ x = F.dropout(F.relu(conv(x)), 0.5, self.training)
+
+ x = x.transpose(1, 2)
+
+ input_lengths = input_lengths.cpu()
+ x = nn.utils.rnn.pack_padded_sequence(x, input_lengths, batch_first=True)
+
+ outputs, _ = self.lstm(x)
+ outputs, _ = nn.utils.rnn.pad_packed_sequence(outputs, batch_first=True)
+
+ return outputs
+
+
+class _Decoder(nn.Module):
+ r"""Decoder with Attention model.
+
+ Args:
+ n_mels (int): number of mel bins
+ n_frames_per_step (int): number of frames processed per step, only 1 is supported
+ encoder_embedding_dim (int): the number of embedding dimensions in the encoder.
+ decoder_rnn_dim (int): number of units in decoder LSTM
+ decoder_max_step (int): maximum number of output mel spectrograms
+ decoder_dropout (float): dropout probability for decoder LSTM
+ decoder_early_stopping (bool): stop decoding when all samples are finished
+ attention_rnn_dim (int): number of units in attention LSTM
+ attention_hidden_dim (int): dimension of attention hidden representation
+ attention_location_n_filter (int): number of filters for attention model
+ attention_location_kernel_size (int): kernel size for attention model
+ attention_dropout (float): dropout probability for attention LSTM
+ prenet_dim (int): number of ReLU units in prenet layers
+ gate_threshold (float): probability threshold for stop token
+ """
+
+ def __init__(
+ self,
+ n_mels: int,
+ n_frames_per_step: int,
+ encoder_embedding_dim: int,
+ decoder_rnn_dim: int,
+ decoder_max_step: int,
+ decoder_dropout: float,
+ decoder_early_stopping: bool,
+ attention_rnn_dim: int,
+ attention_hidden_dim: int,
+ attention_location_n_filter: int,
+ attention_location_kernel_size: int,
+ attention_dropout: float,
+ prenet_dim: int,
+ gate_threshold: float,
+ ) -> None:
+
+ super().__init__()
+ self.n_mels = n_mels
+ self.n_frames_per_step = n_frames_per_step
+ self.encoder_embedding_dim = encoder_embedding_dim
+ self.attention_rnn_dim = attention_rnn_dim
+ self.decoder_rnn_dim = decoder_rnn_dim
+ self.prenet_dim = prenet_dim
+ self.decoder_max_step = decoder_max_step
+ self.gate_threshold = gate_threshold
+ self.attention_dropout = attention_dropout
+ self.decoder_dropout = decoder_dropout
+ self.decoder_early_stopping = decoder_early_stopping
+
+ self.prenet = _Prenet(n_mels * n_frames_per_step, [prenet_dim, prenet_dim])
+
+ self.attention_rnn = nn.LSTMCell(prenet_dim + encoder_embedding_dim, attention_rnn_dim)
+
+ self.attention_layer = _Attention(
+ attention_rnn_dim,
+ encoder_embedding_dim,
+ attention_hidden_dim,
+ attention_location_n_filter,
+ attention_location_kernel_size,
+ )
+
+ self.decoder_rnn = nn.LSTMCell(attention_rnn_dim + encoder_embedding_dim, decoder_rnn_dim, True)
+
+ self.linear_projection = _get_linear_layer(decoder_rnn_dim + encoder_embedding_dim, n_mels * n_frames_per_step)
+
+ self.gate_layer = _get_linear_layer(
+ decoder_rnn_dim + encoder_embedding_dim, 1, bias=True, w_init_gain="sigmoid"
+ )
+
+ def _get_initial_frame(self, memory: Tensor) -> Tensor:
+ r"""Gets all zeros frames to use as the first decoder input.
+
+ Args:
+ memory (Tensor): Encoder outputs with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``).
+
+ Returns:
+ decoder_input (Tensor): all zeros frames with shape
+ (n_batch, max of ``text_lengths``, ``n_mels * n_frames_per_step``).
+ """
+
+ n_batch = memory.size(0)
+ dtype = memory.dtype
+ device = memory.device
+ decoder_input = torch.zeros(n_batch, self.n_mels * self.n_frames_per_step, dtype=dtype, device=device)
+ return decoder_input
+
+ def _initialize_decoder_states(
+ self, memory: Tensor
+ ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+ r"""Initializes attention rnn states, decoder rnn states, attention
+ weights, attention cumulative weights, attention context, stores memory
+ and stores processed memory.
+
+ Args:
+ memory (Tensor): Encoder outputs with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``).
+
+ Returns:
+ attention_hidden (Tensor): Hidden state of the attention LSTM with shape (n_batch, ``attention_rnn_dim``).
+ attention_cell (Tensor): Hidden state of the attention LSTM with shape (n_batch, ``attention_rnn_dim``).
+ decoder_hidden (Tensor): Hidden state of the decoder LSTM with shape (n_batch, ``decoder_rnn_dim``).
+ decoder_cell (Tensor): Hidden state of the decoder LSTM with shape (n_batch, ``decoder_rnn_dim``).
+ attention_weights (Tensor): Attention weights with shape (n_batch, max of ``text_lengths``).
+ attention_weights_cum (Tensor): Cumulated attention weights with shape (n_batch, max of ``text_lengths``).
+ attention_context (Tensor): Context vector with shape (n_batch, ``encoder_embedding_dim``).
+ processed_memory (Tensor): Processed encoder outputs
+ with shape (n_batch, max of ``text_lengths``, ``attention_hidden_dim``).
+ """
+ n_batch = memory.size(0)
+ max_time = memory.size(1)
+ dtype = memory.dtype
+ device = memory.device
+
+ attention_hidden = torch.zeros(n_batch, self.attention_rnn_dim, dtype=dtype, device=device)
+ attention_cell = torch.zeros(n_batch, self.attention_rnn_dim, dtype=dtype, device=device)
+
+ decoder_hidden = torch.zeros(n_batch, self.decoder_rnn_dim, dtype=dtype, device=device)
+ decoder_cell = torch.zeros(n_batch, self.decoder_rnn_dim, dtype=dtype, device=device)
+
+ attention_weights = torch.zeros(n_batch, max_time, dtype=dtype, device=device)
+ attention_weights_cum = torch.zeros(n_batch, max_time, dtype=dtype, device=device)
+ attention_context = torch.zeros(n_batch, self.encoder_embedding_dim, dtype=dtype, device=device)
+
+ processed_memory = self.attention_layer.memory_layer(memory)
+
+ return (
+ attention_hidden,
+ attention_cell,
+ decoder_hidden,
+ decoder_cell,
+ attention_weights,
+ attention_weights_cum,
+ attention_context,
+ processed_memory,
+ )
+
+ def _parse_decoder_inputs(self, decoder_inputs: Tensor) -> Tensor:
+ r"""Prepares decoder inputs.
+
+ Args:
+ decoder_inputs (Tensor): Inputs used for teacher-forced training, i.e. mel-specs,
+ with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``)
+
+ Returns:
+ inputs (Tensor): Processed decoder inputs with shape (max of ``mel_specgram_lengths``, n_batch, ``n_mels``).
+ """
+ # (n_batch, n_mels, mel_specgram_lengths.max()) -> (n_batch, mel_specgram_lengths.max(), n_mels)
+ decoder_inputs = decoder_inputs.transpose(1, 2)
+ decoder_inputs = decoder_inputs.view(
+ decoder_inputs.size(0),
+ int(decoder_inputs.size(1) / self.n_frames_per_step),
+ -1,
+ )
+ # (n_batch, mel_specgram_lengths.max(), n_mels) -> (mel_specgram_lengths.max(), n_batch, n_mels)
+ decoder_inputs = decoder_inputs.transpose(0, 1)
+ return decoder_inputs
+
+ def _parse_decoder_outputs(
+ self, mel_specgram: Tensor, gate_outputs: Tensor, alignments: Tensor
+ ) -> Tuple[Tensor, Tensor, Tensor]:
+ r"""Prepares decoder outputs for output
+
+ Args:
+ mel_specgram (Tensor): mel spectrogram with shape (max of ``mel_specgram_lengths``, n_batch, ``n_mels``)
+ gate_outputs (Tensor): predicted stop token with shape (max of ``mel_specgram_lengths``, n_batch)
+ alignments (Tensor): sequence of attention weights from the decoder
+ with shape (max of ``mel_specgram_lengths``, n_batch, max of ``text_lengths``)
+
+ Returns:
+ mel_specgram (Tensor): mel spectrogram with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``)
+ gate_outputs (Tensor): predicted stop token with shape (n_batch, max of ``mel_specgram_lengths``)
+ alignments (Tensor): sequence of attention weights from the decoder
+ with shape (n_batch, max of ``mel_specgram_lengths``, max of ``text_lengths``)
+ """
+ # (mel_specgram_lengths.max(), n_batch, text_lengths.max())
+ # -> (n_batch, mel_specgram_lengths.max(), text_lengths.max())
+ alignments = alignments.transpose(0, 1).contiguous()
+ # (mel_specgram_lengths.max(), n_batch) -> (n_batch, mel_specgram_lengths.max())
+ gate_outputs = gate_outputs.transpose(0, 1).contiguous()
+ # (mel_specgram_lengths.max(), n_batch, n_mels) -> (n_batch, mel_specgram_lengths.max(), n_mels)
+ mel_specgram = mel_specgram.transpose(0, 1).contiguous()
+ # decouple frames per step
+ shape = (mel_specgram.shape[0], -1, self.n_mels)
+ mel_specgram = mel_specgram.view(*shape)
+ # (n_batch, mel_specgram_lengths.max(), n_mels) -> (n_batch, n_mels, T_out)
+ mel_specgram = mel_specgram.transpose(1, 2)
+
+ return mel_specgram, gate_outputs, alignments
+
+ def decode(
+ self,
+ decoder_input: Tensor,
+ attention_hidden: Tensor,
+ attention_cell: Tensor,
+ decoder_hidden: Tensor,
+ decoder_cell: Tensor,
+ attention_weights: Tensor,
+ attention_weights_cum: Tensor,
+ attention_context: Tensor,
+ memory: Tensor,
+ processed_memory: Tensor,
+ mask: Tensor,
+ ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+ r"""Decoder step using stored states, attention and memory
+
+ Args:
+ decoder_input (Tensor): Output of the Prenet with shape (n_batch, ``prenet_dim``).
+ attention_hidden (Tensor): Hidden state of the attention LSTM with shape (n_batch, ``attention_rnn_dim``).
+ attention_cell (Tensor): Hidden state of the attention LSTM with shape (n_batch, ``attention_rnn_dim``).
+ decoder_hidden (Tensor): Hidden state of the decoder LSTM with shape (n_batch, ``decoder_rnn_dim``).
+ decoder_cell (Tensor): Hidden state of the decoder LSTM with shape (n_batch, ``decoder_rnn_dim``).
+ attention_weights (Tensor): Attention weights with shape (n_batch, max of ``text_lengths``).
+ attention_weights_cum (Tensor): Cumulated attention weights with shape (n_batch, max of ``text_lengths``).
+ attention_context (Tensor): Context vector with shape (n_batch, ``encoder_embedding_dim``).
+ memory (Tensor): Encoder output with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``).
+ processed_memory (Tensor): Processed Encoder outputs
+ with shape (n_batch, max of ``text_lengths``, ``attention_hidden_dim``).
+ mask (Tensor): Binary mask for padded data with shape (n_batch, current_num_frames).
+
+ Returns:
+ decoder_output: Predicted mel spectrogram for the current frame with shape (n_batch, ``n_mels``).
+ gate_prediction (Tensor): Prediction of the stop token with shape (n_batch, ``1``).
+ attention_hidden (Tensor): Hidden state of the attention LSTM with shape (n_batch, ``attention_rnn_dim``).
+ attention_cell (Tensor): Hidden state of the attention LSTM with shape (n_batch, ``attention_rnn_dim``).
+ decoder_hidden (Tensor): Hidden state of the decoder LSTM with shape (n_batch, ``decoder_rnn_dim``).
+ decoder_cell (Tensor): Hidden state of the decoder LSTM with shape (n_batch, ``decoder_rnn_dim``).
+ attention_weights (Tensor): Attention weights with shape (n_batch, max of ``text_lengths``).
+ attention_weights_cum (Tensor): Cumulated attention weights with shape (n_batch, max of ``text_lengths``).
+ attention_context (Tensor): Context vector with shape (n_batch, ``encoder_embedding_dim``).
+ """
+ cell_input = torch.cat((decoder_input, attention_context), -1)
+
+ attention_hidden, attention_cell = self.attention_rnn(cell_input, (attention_hidden, attention_cell))
+ attention_hidden = F.dropout(attention_hidden, self.attention_dropout, self.training)
+
+ attention_weights_cat = torch.cat((attention_weights.unsqueeze(1), attention_weights_cum.unsqueeze(1)), dim=1)
+ attention_context, attention_weights = self.attention_layer(
+ attention_hidden, memory, processed_memory, attention_weights_cat, mask
+ )
+
+ attention_weights_cum += attention_weights
+ decoder_input = torch.cat((attention_hidden, attention_context), -1)
+
+ decoder_hidden, decoder_cell = self.decoder_rnn(decoder_input, (decoder_hidden, decoder_cell))
+ decoder_hidden = F.dropout(decoder_hidden, self.decoder_dropout, self.training)
+
+ decoder_hidden_attention_context = torch.cat((decoder_hidden, attention_context), dim=1)
+ decoder_output = self.linear_projection(decoder_hidden_attention_context)
+
+ gate_prediction = self.gate_layer(decoder_hidden_attention_context)
+
+ return (
+ decoder_output,
+ gate_prediction,
+ attention_hidden,
+ attention_cell,
+ decoder_hidden,
+ decoder_cell,
+ attention_weights,
+ attention_weights_cum,
+ attention_context,
+ )
+
+ def forward(
+ self, memory: Tensor, mel_specgram_truth: Tensor, memory_lengths: Tensor
+ ) -> Tuple[Tensor, Tensor, Tensor]:
+ r"""Decoder forward pass for training.
+
+ Args:
+ memory (Tensor): Encoder outputs
+ with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``).
+ mel_specgram_truth (Tensor): Decoder ground-truth mel-specs for teacher forcing
+ with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``).
+ memory_lengths (Tensor): Encoder output lengths for attention masking
+ (the same as ``text_lengths``) with shape (n_batch, ).
+
+ Returns:
+ mel_specgram (Tensor): Predicted mel spectrogram
+ with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``).
+ gate_outputs (Tensor): Predicted stop token for each timestep
+ with shape (n_batch, max of ``mel_specgram_lengths``).
+ alignments (Tensor): Sequence of attention weights from the decoder
+ with shape (n_batch, max of ``mel_specgram_lengths``, max of ``text_lengths``).
+ """
+
+ decoder_input = self._get_initial_frame(memory).unsqueeze(0)
+ decoder_inputs = self._parse_decoder_inputs(mel_specgram_truth)
+ decoder_inputs = torch.cat((decoder_input, decoder_inputs), dim=0)
+ decoder_inputs = self.prenet(decoder_inputs)
+
+ mask = _get_mask_from_lengths(memory_lengths)
+ (
+ attention_hidden,
+ attention_cell,
+ decoder_hidden,
+ decoder_cell,
+ attention_weights,
+ attention_weights_cum,
+ attention_context,
+ processed_memory,
+ ) = self._initialize_decoder_states(memory)
+
+ mel_outputs, gate_outputs, alignments = [], [], []
+ while len(mel_outputs) < decoder_inputs.size(0) - 1:
+ decoder_input = decoder_inputs[len(mel_outputs)]
+ (
+ mel_output,
+ gate_output,
+ attention_hidden,
+ attention_cell,
+ decoder_hidden,
+ decoder_cell,
+ attention_weights,
+ attention_weights_cum,
+ attention_context,
+ ) = self.decode(
+ decoder_input,
+ attention_hidden,
+ attention_cell,
+ decoder_hidden,
+ decoder_cell,
+ attention_weights,
+ attention_weights_cum,
+ attention_context,
+ memory,
+ processed_memory,
+ mask,
+ )
+
+ mel_outputs += [mel_output.squeeze(1)]
+ gate_outputs += [gate_output.squeeze(1)]
+ alignments += [attention_weights]
+
+ mel_specgram, gate_outputs, alignments = self._parse_decoder_outputs(
+ torch.stack(mel_outputs), torch.stack(gate_outputs), torch.stack(alignments)
+ )
+
+ return mel_specgram, gate_outputs, alignments
+
+ def _get_go_frame(self, memory: Tensor) -> Tensor:
+ """Gets all zeros frames to use as the first decoder input
+
+ args:
+ memory (Tensor): Encoder outputs
+ with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``).
+
+ returns:
+ decoder_input (Tensor): All zeros frames with shape(n_batch, ``n_mels`` * ``n_frame_per_step``).
+ """
+
+ n_batch = memory.size(0)
+ dtype = memory.dtype
+ device = memory.device
+ decoder_input = torch.zeros(n_batch, self.n_mels * self.n_frames_per_step, dtype=dtype, device=device)
+ return decoder_input
+
+ @torch.jit.export
+ def infer(self, memory: Tensor, memory_lengths: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+ """Decoder inference
+
+ Args:
+ memory (Tensor): Encoder outputs
+ with shape (n_batch, max of ``text_lengths``, ``encoder_embedding_dim``).
+ memory_lengths (Tensor): Encoder output lengths for attention masking
+ (the same as ``text_lengths``) with shape (n_batch, ).
+
+ Returns:
+ mel_specgram (Tensor): Predicted mel spectrogram
+ with shape (n_batch, ``n_mels``, max of ``mel_specgram_lengths``).
+ mel_specgram_lengths (Tensor): the length of the predicted mel spectrogram (n_batch, ))
+ gate_outputs (Tensor): Predicted stop token for each timestep
+ with shape (n_batch, max of ``mel_specgram_lengths``).
+ alignments (Tensor): Sequence of attention weights from the decoder
+ with shape (n_batch, max of ``mel_specgram_lengths``, max of ``text_lengths``).
+ """
+ batch_size, device = memory.size(0), memory.device
+
+ decoder_input = self._get_go_frame(memory)
+
+ mask = _get_mask_from_lengths(memory_lengths)
+ (
+ attention_hidden,
+ attention_cell,
+ decoder_hidden,
+ decoder_cell,
+ attention_weights,
+ attention_weights_cum,
+ attention_context,
+ processed_memory,
+ ) = self._initialize_decoder_states(memory)
+
+ mel_specgram_lengths = torch.zeros([batch_size], dtype=torch.int32, device=device)
+ finished = torch.zeros([batch_size], dtype=torch.bool, device=device)
+ mel_specgrams: List[Tensor] = []
+ gate_outputs: List[Tensor] = []
+ alignments: List[Tensor] = []
+ for _ in range(self.decoder_max_step):
+ decoder_input = self.prenet(decoder_input)
+ (
+ mel_specgram,
+ gate_output,
+ attention_hidden,
+ attention_cell,
+ decoder_hidden,
+ decoder_cell,
+ attention_weights,
+ attention_weights_cum,
+ attention_context,
+ ) = self.decode(
+ decoder_input,
+ attention_hidden,
+ attention_cell,
+ decoder_hidden,
+ decoder_cell,
+ attention_weights,
+ attention_weights_cum,
+ attention_context,
+ memory,
+ processed_memory,
+ mask,
+ )
+
+ mel_specgrams.append(mel_specgram.unsqueeze(0))
+ gate_outputs.append(gate_output.transpose(0, 1))
+ alignments.append(attention_weights)
+ mel_specgram_lengths[~finished] += 1
+
+ finished |= torch.sigmoid(gate_output.squeeze(1)) > self.gate_threshold
+ if self.decoder_early_stopping and torch.all(finished):
+ break
+
+ decoder_input = mel_specgram
+
+ if len(mel_specgrams) == self.decoder_max_step:
+ warnings.warn(
+ "Reached max decoder steps. The generated spectrogram might not cover " "the whole transcript."
+ )
+
+ mel_specgrams = torch.cat(mel_specgrams, dim=0)
+ gate_outputs = torch.cat(gate_outputs, dim=0)
+ alignments = torch.cat(alignments, dim=0)
+
+ mel_specgrams, gate_outputs, alignments = self._parse_decoder_outputs(mel_specgrams, gate_outputs, alignments)
+
+ return mel_specgrams, mel_specgram_lengths, gate_outputs, alignments
+
+
+class Tacotron2(nn.Module):
+ r"""Tacotron2 model from *Natural TTS Synthesis by Conditioning WaveNet on Mel Spectrogram Predictions*
+ :cite:`shen2018natural` based on the implementation from
+ `Nvidia Deep Learning Examples `_.
+
+ See Also:
+ * :class:`torchaudio.pipelines.Tacotron2TTSBundle`: TTS pipeline with pretrained model.
+
+ Args:
+ mask_padding (bool, optional): Use mask padding (Default: ``False``).
+ n_mels (int, optional): Number of mel bins (Default: ``80``).
+ n_symbol (int, optional): Number of symbols for the input text (Default: ``148``).
+ n_frames_per_step (int, optional): Number of frames processed per step, only 1 is supported (Default: ``1``).
+ symbol_embedding_dim (int, optional): Input embedding dimension (Default: ``512``).
+ encoder_n_convolution (int, optional): Number of encoder convolutions (Default: ``3``).
+ encoder_kernel_size (int, optional): Encoder kernel size (Default: ``5``).
+ encoder_embedding_dim (int, optional): Encoder embedding dimension (Default: ``512``).
+ decoder_rnn_dim (int, optional): Number of units in decoder LSTM (Default: ``1024``).
+ decoder_max_step (int, optional): Maximum number of output mel spectrograms (Default: ``2000``).
+ decoder_dropout (float, optional): Dropout probability for decoder LSTM (Default: ``0.1``).
+ decoder_early_stopping (bool, optional): Continue decoding after all samples are finished (Default: ``True``).
+ attention_rnn_dim (int, optional): Number of units in attention LSTM (Default: ``1024``).
+ attention_hidden_dim (int, optional): Dimension of attention hidden representation (Default: ``128``).
+ attention_location_n_filter (int, optional): Number of filters for attention model (Default: ``32``).
+ attention_location_kernel_size (int, optional): Kernel size for attention model (Default: ``31``).
+ attention_dropout (float, optional): Dropout probability for attention LSTM (Default: ``0.1``).
+ prenet_dim (int, optional): Number of ReLU units in prenet layers (Default: ``256``).
+ postnet_n_convolution (int, optional): Number of postnet convolutions (Default: ``5``).
+ postnet_kernel_size (int, optional): Postnet kernel size (Default: ``5``).
+ postnet_embedding_dim (int, optional): Postnet embedding dimension (Default: ``512``).
+ gate_threshold (float, optional): Probability threshold for stop token (Default: ``0.5``).
+ """
+
+ def __init__(
+ self,
+ mask_padding: bool = False,
+ n_mels: int = 80,
+ n_symbol: int = 148,
+ n_frames_per_step: int = 1,
+ symbol_embedding_dim: int = 512,
+ encoder_embedding_dim: int = 512,
+ encoder_n_convolution: int = 3,
+ encoder_kernel_size: int = 5,
+ decoder_rnn_dim: int = 1024,
+ decoder_max_step: int = 2000,
+ decoder_dropout: float = 0.1,
+ decoder_early_stopping: bool = True,
+ attention_rnn_dim: int = 1024,
+ attention_hidden_dim: int = 128,
+ attention_location_n_filter: int = 32,
+ attention_location_kernel_size: int = 31,
+ attention_dropout: float = 0.1,
+ prenet_dim: int = 256,
+ postnet_n_convolution: int = 5,
+ postnet_kernel_size: int = 5,
+ postnet_embedding_dim: int = 512,
+ gate_threshold: float = 0.5,
+ ) -> None:
+ super().__init__()
+
+ self.mask_padding = mask_padding
+ self.n_mels = n_mels
+ self.n_frames_per_step = n_frames_per_step
+ self.embedding = nn.Embedding(n_symbol, symbol_embedding_dim)
+ torch.nn.init.xavier_uniform_(self.embedding.weight)
+ self.encoder = _Encoder(encoder_embedding_dim, encoder_n_convolution, encoder_kernel_size)
+ self.decoder = _Decoder(
+ n_mels,
+ n_frames_per_step,
+ encoder_embedding_dim,
+ decoder_rnn_dim,
+ decoder_max_step,
+ decoder_dropout,
+ decoder_early_stopping,
+ attention_rnn_dim,
+ attention_hidden_dim,
+ attention_location_n_filter,
+ attention_location_kernel_size,
+ attention_dropout,
+ prenet_dim,
+ gate_threshold,
+ )
+ self.postnet = _Postnet(n_mels, postnet_embedding_dim, postnet_kernel_size, postnet_n_convolution)
+
+ def forward(
+ self,
+ tokens: Tensor,
+ token_lengths: Tensor,
+ mel_specgram: Tensor,
+ mel_specgram_lengths: Tensor,
+ ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+ r"""Pass the input through the Tacotron2 model. This is in teacher
+ forcing mode, which is generally used for training.
+
+ The input ``tokens`` should be padded with zeros to length max of ``token_lengths``.
+ The input ``mel_specgram`` should be padded with zeros to length max of ``mel_specgram_lengths``.
+
+ Args:
+ tokens (Tensor): The input tokens to Tacotron2 with shape `(n_batch, max of token_lengths)`.
+ token_lengths (Tensor): The valid length of each sample in ``tokens`` with shape `(n_batch, )`.
+ mel_specgram (Tensor): The target mel spectrogram
+ with shape `(n_batch, n_mels, max of mel_specgram_lengths)`.
+ mel_specgram_lengths (Tensor): The length of each mel spectrogram with shape `(n_batch, )`.
+
+ Returns:
+ [Tensor, Tensor, Tensor, Tensor]:
+ Tensor
+ Mel spectrogram before Postnet with shape `(n_batch, n_mels, max of mel_specgram_lengths)`.
+ Tensor
+ Mel spectrogram after Postnet with shape `(n_batch, n_mels, max of mel_specgram_lengths)`.
+ Tensor
+ The output for stop token at each time step with shape `(n_batch, max of mel_specgram_lengths)`.
+ Tensor
+ Sequence of attention weights from the decoder with
+ shape `(n_batch, max of mel_specgram_lengths, max of token_lengths)`.
+ """
+
+ embedded_inputs = self.embedding(tokens).transpose(1, 2)
+
+ encoder_outputs = self.encoder(embedded_inputs, token_lengths)
+ mel_specgram, gate_outputs, alignments = self.decoder(
+ encoder_outputs, mel_specgram, memory_lengths=token_lengths
+ )
+
+ mel_specgram_postnet = self.postnet(mel_specgram)
+ mel_specgram_postnet = mel_specgram + mel_specgram_postnet
+
+ if self.mask_padding:
+ mask = _get_mask_from_lengths(mel_specgram_lengths)
+ mask = mask.expand(self.n_mels, mask.size(0), mask.size(1))
+ mask = mask.permute(1, 0, 2)
+
+ mel_specgram.masked_fill_(mask, 0.0)
+ mel_specgram_postnet.masked_fill_(mask, 0.0)
+ gate_outputs.masked_fill_(mask[:, 0, :], 1e3)
+
+ return mel_specgram, mel_specgram_postnet, gate_outputs, alignments
+
+ @torch.jit.export
+ def infer(self, tokens: Tensor, lengths: Optional[Tensor] = None) -> Tuple[Tensor, Tensor, Tensor]:
+ r"""Using Tacotron2 for inference. The input is a batch of encoded
+ sentences (``tokens``) and its corresponding lengths (``lengths``). The
+ output is the generated mel spectrograms, its corresponding lengths, and
+ the attention weights from the decoder.
+
+ The input `tokens` should be padded with zeros to length max of ``lengths``.
+
+ Args:
+ tokens (Tensor): The input tokens to Tacotron2 with shape `(n_batch, max of lengths)`.
+ lengths (Tensor or None, optional):
+ The valid length of each sample in ``tokens`` with shape `(n_batch, )`.
+ If ``None``, it is assumed that the all the tokens are valid. Default: ``None``
+
+ Returns:
+ (Tensor, Tensor, Tensor):
+ Tensor
+ The predicted mel spectrogram with shape `(n_batch, n_mels, max of mel_specgram_lengths)`.
+ Tensor
+ The length of the predicted mel spectrogram with shape `(n_batch, )`.
+ Tensor
+ Sequence of attention weights from the decoder with shape
+ `(n_batch, max of mel_specgram_lengths, max of lengths)`.
+ """
+ n_batch, max_length = tokens.shape
+ if lengths is None:
+ lengths = torch.tensor([max_length]).expand(n_batch).to(tokens.device, tokens.dtype)
+
+ assert lengths is not None # For TorchScript compiler
+ embedded_inputs = self.embedding(tokens).transpose(1, 2)
+ encoder_outputs = self.encoder(embedded_inputs, lengths)
+ mel_specgram, mel_specgram_lengths, _, alignments = self.decoder.infer(encoder_outputs, lengths)
+
+ mel_outputs_postnet = self.postnet(mel_specgram)
+ mel_outputs_postnet = mel_specgram + mel_outputs_postnet
+
+ alignments = alignments.unfold(1, n_batch, n_batch).transpose(0, 2)
+
+ return mel_outputs_postnet, mel_specgram_lengths, alignments
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2letter.py b/lib/python3.12/site-packages/torchaudio/models/wav2letter.py
new file mode 100644
index 0000000000000000000000000000000000000000..d776131686d1f65982a565088e72e45e7b7c107f
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/wav2letter.py
@@ -0,0 +1,72 @@
+from torch import nn, Tensor
+
+__all__ = [
+ "Wav2Letter",
+]
+
+
+class Wav2Letter(nn.Module):
+ r"""Wav2Letter model architecture from *Wav2Letter: an End-to-End ConvNet-based Speech
+ Recognition System* :cite:`collobert2016wav2letter`.
+
+ See Also:
+ * `Training example `__
+
+ Args:
+ num_classes (int, optional): Number of classes to be classified. (Default: ``40``)
+ input_type (str, optional): Wav2Letter can use as input: ``waveform``, ``power_spectrum``
+ or ``mfcc`` (Default: ``waveform``).
+ num_features (int, optional): Number of input features that the network will receive (Default: ``1``).
+ """
+
+ def __init__(self, num_classes: int = 40, input_type: str = "waveform", num_features: int = 1) -> None:
+ super().__init__()
+
+ acoustic_num_features = 250 if input_type == "waveform" else num_features
+ acoustic_model = nn.Sequential(
+ nn.Conv1d(in_channels=acoustic_num_features, out_channels=250, kernel_size=48, stride=2, padding=23),
+ nn.ReLU(inplace=True),
+ nn.Conv1d(in_channels=250, out_channels=250, kernel_size=7, stride=1, padding=3),
+ nn.ReLU(inplace=True),
+ nn.Conv1d(in_channels=250, out_channels=250, kernel_size=7, stride=1, padding=3),
+ nn.ReLU(inplace=True),
+ nn.Conv1d(in_channels=250, out_channels=250, kernel_size=7, stride=1, padding=3),
+ nn.ReLU(inplace=True),
+ nn.Conv1d(in_channels=250, out_channels=250, kernel_size=7, stride=1, padding=3),
+ nn.ReLU(inplace=True),
+ nn.Conv1d(in_channels=250, out_channels=250, kernel_size=7, stride=1, padding=3),
+ nn.ReLU(inplace=True),
+ nn.Conv1d(in_channels=250, out_channels=250, kernel_size=7, stride=1, padding=3),
+ nn.ReLU(inplace=True),
+ nn.Conv1d(in_channels=250, out_channels=250, kernel_size=7, stride=1, padding=3),
+ nn.ReLU(inplace=True),
+ nn.Conv1d(in_channels=250, out_channels=2000, kernel_size=32, stride=1, padding=16),
+ nn.ReLU(inplace=True),
+ nn.Conv1d(in_channels=2000, out_channels=2000, kernel_size=1, stride=1, padding=0),
+ nn.ReLU(inplace=True),
+ nn.Conv1d(in_channels=2000, out_channels=num_classes, kernel_size=1, stride=1, padding=0),
+ nn.ReLU(inplace=True),
+ )
+
+ if input_type == "waveform":
+ waveform_model = nn.Sequential(
+ nn.Conv1d(in_channels=num_features, out_channels=250, kernel_size=250, stride=160, padding=45),
+ nn.ReLU(inplace=True),
+ )
+ self.acoustic_model = nn.Sequential(waveform_model, acoustic_model)
+
+ if input_type in ["power_spectrum", "mfcc"]:
+ self.acoustic_model = acoustic_model
+
+ def forward(self, x: Tensor) -> Tensor:
+ r"""
+ Args:
+ x (torch.Tensor): Tensor of dimension (batch_size, num_features, input_length).
+
+ Returns:
+ Tensor: Predictor tensor of dimension (batch_size, number_of_classes, input_length).
+ """
+
+ x = self.acoustic_model(x)
+ x = nn.functional.log_softmax(x, dim=1)
+ return x
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__init__.py b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..bb83403f5719b68c790d2f9f934f8c80acea3557
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__init__.py
@@ -0,0 +1,45 @@
+from . import utils
+from .model import (
+ hubert_base,
+ hubert_large,
+ hubert_pretrain_base,
+ hubert_pretrain_large,
+ hubert_pretrain_model,
+ hubert_pretrain_xlarge,
+ hubert_xlarge,
+ HuBERTPretrainModel,
+ wav2vec2_base,
+ wav2vec2_large,
+ wav2vec2_large_lv60k,
+ wav2vec2_model,
+ wav2vec2_xlsr_1b,
+ wav2vec2_xlsr_2b,
+ wav2vec2_xlsr_300m,
+ Wav2Vec2Model,
+ wavlm_base,
+ wavlm_large,
+ wavlm_model,
+)
+
+__all__ = [
+ "Wav2Vec2Model",
+ "HuBERTPretrainModel",
+ "wavlm_model",
+ "wavlm_base",
+ "wavlm_large",
+ "wav2vec2_model",
+ "wav2vec2_base",
+ "wav2vec2_large",
+ "wav2vec2_large_lv60k",
+ "hubert_base",
+ "hubert_large",
+ "hubert_xlarge",
+ "hubert_pretrain_model",
+ "hubert_pretrain_base",
+ "hubert_pretrain_large",
+ "hubert_pretrain_xlarge",
+ "utils",
+ "wav2vec2_xlsr_300m",
+ "wav2vec2_xlsr_1b",
+ "wav2vec2_xlsr_2b",
+]
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..845cbfb01386561451fd15d5784b895b8618638e
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__pycache__/components.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__pycache__/components.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..d6d932abe0e0086d43241e5091ec92a5b64b32eb
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__pycache__/components.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__pycache__/model.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__pycache__/model.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..11ec57a365cb512e7499afbf740298593c09009b
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__pycache__/model.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__pycache__/wavlm_attention.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__pycache__/wavlm_attention.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..80f35275e999f41c355e809d4e326dbe6341ccd7
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/__pycache__/wavlm_attention.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/components.py b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/components.py
new file mode 100644
index 0000000000000000000000000000000000000000..480a6ae50921efebf5930dc21caaa3a1a44945dd
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/components.py
@@ -0,0 +1,1167 @@
+import logging
+from typing import List, Optional, Tuple
+
+import torch
+from torch import nn, Tensor
+from torch.nn import Module, Parameter
+
+from .wavlm_attention import WavLMSelfAttention
+
+_LG = logging.getLogger(__name__)
+
+
+def _init_transformer_params(module):
+ """
+ Initialize the weights of Transformer module in Wav2Vec2/HuBERT.
+
+ If the module is ``nn.Linear``, normalize the weight with mean 0 and standard deviation 0.02.
+ If ``bias`` is set to ``True`` in the module, set ``bias`` to 0.
+
+ If the module is ``nn.Embedding``, normalize the weight with mean 0 and standard deviation 0.02.
+ If ``padding_idx`` is not None, set the weight of padding to 0.
+
+ Note:
+ Ths method corresponds to
+ `init_bert_params
+ `__
+ in the original ``fairseq`` implementation.
+ """
+
+ def normal_(data):
+ data.copy_(data.cpu().normal_(mean=0.0, std=0.02).to(data.device))
+
+ if isinstance(module, nn.Linear):
+ normal_(module.weight.data)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ if isinstance(module, nn.Embedding):
+ normal_(module.weight.data)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+
+
+class LayerNorm(nn.LayerNorm):
+ """Layer norm with transpose"""
+
+ def forward(self, input: Tensor) -> Tensor:
+ x = input.transpose(-2, -1)
+ x = nn.functional.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+ x = x.transpose(-2, -1)
+ return x
+
+
+class ConvLayerBlock(Module):
+ """Convolution unit of FeatureExtractor"""
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int,
+ bias: bool,
+ layer_norm: Optional[Module],
+ ):
+ super().__init__()
+ self.kernel_size = kernel_size
+ self.stride = stride
+ self.layer_norm = layer_norm
+ self.conv = nn.Conv1d(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ bias=bias,
+ )
+
+ def forward(
+ self,
+ x: Tensor,
+ length: Optional[Tensor],
+ ) -> Tuple[Tensor, Optional[Tensor]]:
+ """
+ Args:
+ x (Tensor): Shape: ``[batch, in_channels, in_frame]``.
+ length (Tensor or None, optional): Shape ``[batch, ]``.
+ Returns:
+ Tensor: Shape ``[batch, out_channels, out_frames]``.
+ Optional[Tensor]: Shape ``[batch, ]``.
+ """
+ x = self.conv(x)
+ if self.layer_norm is not None:
+ x = self.layer_norm(x)
+ x = nn.functional.gelu(x)
+
+ if length is not None:
+ length = torch.div(length - self.kernel_size, self.stride, rounding_mode="floor") + 1
+ # When input length is 0, the resulting length can be negative. So fix it here.
+ length = torch.max(torch.zeros_like(length), length)
+ return x, length
+
+
+class FeatureExtractor(Module):
+ """Extract features from audio
+
+ Args:
+ conv_layers (nn.ModuleList):
+ convolution layers
+ """
+
+ def __init__(
+ self,
+ conv_layers: nn.ModuleList,
+ ):
+ super().__init__()
+ self.conv_layers = conv_layers
+
+ def forward(
+ self,
+ x: Tensor,
+ length: Optional[Tensor],
+ ) -> Tuple[Tensor, Optional[Tensor]]:
+ """
+ Args:
+ x (Tensor):
+ Input Tensor representing a batch of audio,
+ shape: ``[batch, time]``.
+ length (Tensor or None, optional):
+ Valid length of each input sample. shape: ``[batch, ]``.
+
+ Returns:
+ Tensor:
+ The resulting feature, shape: ``[batch, frame, feature]``
+ Optional[Tensor]:
+ Valid length of each output sample. shape: ``[batch, ]``.
+ """
+ if x.ndim != 2:
+ raise ValueError(f"Expected the input Tensor to be 2D (batch, time). Found: {list(x.shape)}")
+
+ x = x.unsqueeze(1) # (batch, channel==1, frame)
+ for layer in self.conv_layers:
+ x, length = layer(x, length) # (batch, feature, frame)
+ x = x.transpose(1, 2) # (batch, frame, feature)
+ return x, length
+
+
+class FeatureProjection(Module):
+ """Layer that connects FeatureExtractor and Encoder
+
+ Projects features to encoder dimension.
+
+ Args:
+ in_features (int): Input feature dim.
+ out_features (int): Output feature dim.
+ dropout (float): Dropout probability.
+ """
+
+ def __init__(
+ self,
+ in_features: int,
+ out_features: int,
+ dropout: float,
+ ):
+ super().__init__()
+ self.layer_norm = nn.LayerNorm(in_features)
+ self.projection = nn.Linear(
+ in_features,
+ out_features,
+ )
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(self, x):
+ """
+ Args:
+ x (Tensor):
+ Feature Tensor. shape: ``[batch, frame, in_feature]``
+ Returns:
+ Tensor: Projected features. ``[batch, frame, out_feature]``.
+ """
+ x = self.layer_norm(x)
+ x = self.projection(x)
+ x = self.dropout(x)
+ return x
+
+
+class ConvolutionalPositionalEmbedding(Module):
+ """Positional embedding which is placed at the beginning of Transformer.
+
+ Args:
+ embed_dim (int): Feature dimension of the input Tensor.
+ kernel_size (int): The number of frames to be use.
+ groups (int): The number of groups in feature dimensions.
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ kernel_size: int,
+ groups: int,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.kernel_size = kernel_size
+ self.conv = nn.Conv1d(
+ in_channels=embed_dim,
+ out_channels=embed_dim,
+ kernel_size=kernel_size,
+ padding=kernel_size // 2,
+ groups=groups,
+ )
+
+ 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):
+ 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):
+ """
+ Args:
+ x (Tensor): shape ``[batch, frame, feature]``.
+
+ Returns:
+ Tensor: The resulting feature. Shape ``[batch, frame, feature]``.
+ """
+ x = x.transpose(-2, -1)
+ x = self.conv(x)
+ if self.num_remove > 0:
+ x = x[..., : -self.num_remove]
+ x = torch.nn.functional.gelu(x)
+ x = x.transpose(-2, -1)
+ return x
+
+
+class SelfAttention(Module):
+ """Multihead Self Attention module
+
+ Args:
+ embed_dim (int): Total dimension of the model.
+ num_heads (int): The number of heads.
+ dropout (float, optional):
+ Dropout probability on attn_output_weights. Default: ``0.0``
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ dropout: float = 0.0,
+ ):
+ super().__init__()
+ head_dim = embed_dim // num_heads
+ if head_dim * num_heads != embed_dim:
+ raise ValueError(f"`embed_dim ({embed_dim})` is not divisible by `num_heads ({num_heads})`")
+
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.head_dim = head_dim
+
+ self.scaling = self.head_dim**-0.5
+
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+
+ def forward(
+ self,
+ x: Tensor,
+ attention_mask: Optional[Tensor] = None,
+ position_bias: Optional[Tensor] = None,
+ key_padding_mask: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Optional[Tensor]]:
+ """
+ Args:
+ x (Tensor): shape: ``[batch_size, sequence_length, embed_dim]``.
+ attention_mask (Tensor or ``None``, optional):
+ shape: ``[batch_size, 1, sequence_length, sequence_length]``
+ position_bias: Not used. Only for the compatibility with :py:class:`WavLMSelfAttention`.
+ key_padding_mask (Tensor or ``None``): Not used. Only for the compatibility with
+ :py:class:`WavLMSelfAttention`.
+ Returns:
+ (Tensor, ``None``): The resulting attention output and ``None`` (necessary for compatibility
+ with :py:class:`WavLMSelAttention`).
+ Attention output shape: ``[batch, sequence_length, embed_dim]``.
+ """
+ if x.ndim != 3 or x.shape[2] != self.embed_dim:
+ raise ValueError(
+ f"The expected input shape is (batch, sequence, embed_dim=={self.embed_dim}). " f"Found {x.shape}."
+ )
+ batch_size, length, embed_dim = x.size()
+ if attention_mask is not None:
+ shape_ = (batch_size, 1, length, length)
+ if attention_mask.size() != shape_:
+ raise ValueError(f"The expected attention mask shape is {shape_}. " f"Found {attention_mask.size()}.")
+
+ shape = (batch_size, length, self.num_heads, self.head_dim)
+ q = self.q_proj(x).view(*shape).transpose(2, 1) # B, nH, L, Hd
+ k = self.k_proj(x).view(*shape).transpose(2, 1) # B, nH, L, Hd
+ v = self.v_proj(x).view(*shape).transpose(2, 1) # B, nH, L, Hd
+ dropout = self.dropout if self.training else 0.0
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
+ q, k, v, attn_mask=attention_mask, dropout_p=dropout, is_causal=False
+ )
+ attn_output = attn_output.transpose(1, 2).reshape(batch_size, -1, self.num_heads * self.head_dim)
+ output = self.out_proj(attn_output)
+ return output, None # Necessary for compatibility with WavLMSelAttention
+
+
+class FeedForward(Module):
+ """Layer that follows attention layer in encoder layer."""
+
+ def __init__(
+ self,
+ io_features: int,
+ intermediate_features: int,
+ intermediate_dropout: float,
+ output_dropout: float,
+ ):
+ super().__init__()
+ self.intermediate_dense = nn.Linear(io_features, intermediate_features)
+ self.intermediate_dropout = nn.Dropout(intermediate_dropout)
+ self.output_dense = nn.Linear(intermediate_features, io_features)
+ self.output_dropout = nn.Dropout(output_dropout)
+
+ def forward(self, x):
+ """
+ Args:
+ x (Tensor): shape: `(batch, sequence_length, io_features)`
+ Returns:
+ x (Tensor): shape: `(batch, sequence_length, io_features)`
+ """
+ x = self.intermediate_dense(x)
+ x = torch.nn.functional.gelu(x)
+ x = self.intermediate_dropout(x)
+
+ x = self.output_dense(x)
+ x = self.output_dropout(x)
+ return x
+
+
+class EncoderLayer(Module):
+ """A layer unit in encoder. Combines multihead self attention and feed forward."""
+
+ def __init__(
+ self,
+ attention: Module,
+ dropout: float,
+ layer_norm_first: bool,
+ feed_forward: Module,
+ ):
+ super().__init__()
+ self.attention = attention
+ self.dropout = nn.Dropout(dropout)
+ self.layer_norm = nn.LayerNorm(attention.embed_dim)
+ self.layer_norm_first = layer_norm_first
+ self.feed_forward = feed_forward
+ self.final_layer_norm = nn.LayerNorm(attention.embed_dim)
+
+ def forward(
+ self,
+ x: Tensor,
+ attention_mask: Optional[Tensor] = None,
+ position_bias: Optional[Tensor] = None,
+ key_padding_mask: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Optional[Tensor]]:
+ """
+ Args:
+ x (Tensor): Input of shape ``(batch, sequence_length, embed_dim)``.
+ attention_mask (Tensor or ``None``, optional): attention mask
+ of shape ``(batch, 1, sequence_length, sequence_length)``. (Default: ``None``)
+ position_bias (Tensor or ``None``, optional): position bias of shape
+ ``(batch_size * num_heads, src_len, src_len)``.
+ Only necessary for WavLM model, ``None`` otherwise. (Default: ``None``)
+ key_padding_mask (Tensor or ``None``, optional): key padding mask of shape ``(batch_size, src_len)``.
+ Only used for WavLM model, ignored otherwise. (Default: ``None``)
+ Returns:
+ (x, position_bias): Shapes are the same as in the input. Position bias is only relevant for WaLM model,
+ ``None`` otherwise.
+ """
+ residual = x
+
+ if self.layer_norm_first:
+ x = self.layer_norm(x)
+
+ x, position_bias = self.attention(
+ x, attention_mask=attention_mask, position_bias=position_bias, key_padding_mask=key_padding_mask
+ )
+
+ x = self.dropout(x)
+ x = residual + x
+
+ if self.layer_norm_first:
+ x = x + self.feed_forward(self.final_layer_norm(x))
+ else:
+ x = self.layer_norm(x)
+ x = self.final_layer_norm(x + self.feed_forward(x))
+ return x, position_bias
+
+
+class Transformer(Module):
+ def __init__(
+ self,
+ pos_conv_embed: Module,
+ dropout: float,
+ layers: Module,
+ layer_norm_first: bool,
+ layer_drop: float,
+ ):
+ super().__init__()
+ self.pos_conv_embed = pos_conv_embed
+ self.layer_norm = nn.LayerNorm(pos_conv_embed.embed_dim)
+ self.layer_norm_first = layer_norm_first
+ self.layer_drop = layer_drop
+ self.dropout = nn.Dropout(dropout)
+ self.layers = layers
+
+ def _preprocess(self, x: Tensor):
+ x = x + self.pos_conv_embed(x)
+
+ if self.layer_norm_first:
+ x = self.layer_norm(x)
+
+ x = self.dropout(x)
+ return x
+
+ def forward(
+ self,
+ x: Tensor,
+ attention_mask: Optional[Tensor] = None,
+ position_bias: Optional[Tensor] = None,
+ ) -> Tensor:
+ x = self._preprocess(x)
+ for layer in self.layers:
+ if not (self.training and torch.rand(1).item() <= self.layer_drop):
+ x, position_bias = layer(x, attention_mask, position_bias=position_bias)
+
+ if not self.layer_norm_first:
+ x = self.layer_norm(x)
+ return x
+
+ def get_intermediate_outputs(
+ self,
+ x: Tensor,
+ attention_mask: Optional[Tensor] = None,
+ num_layers: Optional[int] = None,
+ ) -> List[Tensor]:
+ if num_layers is not None:
+ if not 0 < num_layers <= len(self.layers):
+ 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, 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
+ return ret
+
+
+class Encoder(Module):
+ def __init__(
+ self,
+ feature_projection: Module,
+ transformer: Module,
+ ):
+ super().__init__()
+ self.feature_projection = feature_projection
+ self.transformer = transformer
+
+ def _preprocess(
+ self,
+ features: Tensor,
+ lengths: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Optional[Tensor]]:
+ x = self.feature_projection(features)
+
+ mask: Optional[Tensor] = None
+ if lengths is not None:
+ batch_size, max_len, _ = x.shape
+ # create mask for padded elements and zero-out them
+ mask = torch.arange(max_len, device=lengths.device).expand(batch_size, max_len) >= lengths[:, None]
+ x[mask] = 0.0
+ # extend the mask to attention shape and set weight
+ mask = -10000.0 * mask[:, None, None, :].to(dtype=features.dtype)
+ mask = mask.expand(batch_size, 1, max_len, max_len)
+ return x, mask
+
+ def forward(
+ self,
+ features: Tensor,
+ lengths: Optional[Tensor] = None,
+ ) -> Tensor:
+ x, mask = self._preprocess(features, lengths)
+ x = self.transformer(x, attention_mask=mask)
+ return x
+
+ def extract_features(
+ self,
+ features: Tensor,
+ lengths: Optional[Tensor] = None,
+ num_layers: Optional[int] = None,
+ ) -> List[Tensor]:
+ x, masks = self._preprocess(features, lengths)
+ return self.transformer.get_intermediate_outputs(x, attention_mask=masks, num_layers=num_layers)
+
+
+################################################################################
+def _get_feature_extractor(
+ norm_mode: str,
+ shapes: List[Tuple[int, int, int]],
+ bias: bool,
+) -> FeatureExtractor:
+ """
+ Args:
+ norm_mode (str):
+ Either "group_norm" or "layer_norm".
+ If "group_norm", then a single normalization is applied
+ in the first convolution block. Otherwise, all the convolution
+ blocks will have layer normalization.
+ This option corresponds to "extractor_mode" from fairseq.
+ Expected values are "group_norm" for Base arch, and
+ "layer_norm" for Large arch.
+ shapes (list of tuple of int):
+ Configuration of convolution layers. List of convolution configuration,
+ i.e. ``[(output_channel, kernel_size, stride), ...]``
+ This option corresponds to "conv_feature_layers" from fairseq.
+ Expected values are
+ ``[(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512, 2, 2)] * 2``
+ for all the architectures.
+ bias (bool):
+ Whether to include bias term to each convolution operation.
+ This option corresponds to "conv_bias" from fairseq.
+ Expected values are False for Base arch, and True for Large arch.
+
+ See Also:
+ * Original implementation
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L666-L733
+ * "extractor_mode"
+ - Def and base:
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L38-L45
+ - Large:
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L52
+ * "conv_feature_layers"
+ - Def, base and large:
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L94-L100
+ * "conv_bias"
+ - Def and base:
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L101-L103
+ - Large:
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L61
+ """
+ if norm_mode not in ["group_norm", "layer_norm"]:
+ raise ValueError("Invalid norm mode")
+ blocks = []
+ in_channels = 1
+ for i, (out_channels, kernel_size, stride) in enumerate(shapes):
+ normalization = None
+ if norm_mode == "group_norm" and i == 0:
+ normalization = nn.GroupNorm(
+ num_groups=out_channels,
+ num_channels=out_channels,
+ affine=True,
+ )
+ elif norm_mode == "layer_norm":
+ normalization = LayerNorm(
+ normalized_shape=out_channels,
+ elementwise_affine=True,
+ )
+ blocks.append(
+ ConvLayerBlock(
+ in_channels=in_channels,
+ out_channels=out_channels,
+ kernel_size=kernel_size,
+ stride=stride,
+ bias=bias,
+ layer_norm=normalization,
+ )
+ )
+ in_channels = out_channels
+ return FeatureExtractor(nn.ModuleList(blocks))
+
+
+def _get_encoder(
+ in_features: int,
+ embed_dim: int,
+ dropout_input: float,
+ pos_conv_kernel: int,
+ pos_conv_groups: int,
+ num_layers: int,
+ num_heads: int,
+ attention_dropout: float,
+ ff_interm_features: int,
+ ff_interm_dropout: float,
+ dropout: float,
+ layer_norm_first: bool,
+ layer_drop: float,
+) -> Encoder:
+ """
+ Args:
+ in_features (int): The number of input features.
+ embed_dim (int):
+ The dimension of embedding.
+ This option corresponds to "encoder_embed_dim" from fairseq.
+ Expected values are 768 for Base arch, and 1024 for Large arch.
+ dropout_input (float):
+ The dropout probability applied after the input feature is projected
+ to ``embed_dim``.
+ This option corresponds to "dropout_input" from fairseq.
+ Expected values are 0.1 for both Base and Large arch.
+ pos_conv_kernel (int):
+ The kernel size of convolutional positional embeddings.
+ This option corresponds to "conv_pos" from fairseq.
+ Expected values are 128 for both Base and Large arch.
+ pos_conv_groups (int):
+ The number of groups of convolutional positional embeddings.
+ This option corresponds to "conv_pos_groups" from fairseq.
+ Expected values are 16 for both Base and Large arch.
+ num_layers (int):
+ The number of self attention layers in transformer block.
+ This option corresponds to "encoder_layers" from fairseq.
+ Expected values are 12 for Base and 24 for Large arch.
+ num_heads (int):
+ The number of heads in self attention layers.
+ This option corresponds to "encoder_attention_heads" from fairseq.
+ Expected values are 12 for Base and 16 for Large arch.
+ attention_dropout (float):
+ The dropout probability applied after softmax in self-attention layer.
+ This option corresponds to "attention_dropout" from fairseq.
+ Expected values are 0.1 for Base and 0.0 for Large arch.
+ ff_interm_features (int):
+ The dimension of hidden features in feed forward layer.
+ This option corresponds to "encoder_ffn_embed_dim" from fairseq.
+ Expected values are 3072 for Base and 4096 for Large arch.
+ ff_interm_dropout (float):
+ The dropout probability applied in feedforward layer.
+ This option correspinds to "activation_dropout" from fairseq.
+ Expected values are 0.1 for both Base and Large arch.
+ dropout (float):
+ The dropout probability applied at the end of feed forward layer.
+ This option corresponds to "dropout" from fairseq.
+ Expected values are 0.1 for Base and 0.0 for Large arch.
+ layer_norm_first (bool):
+ Control the order of layer norm in transformer layer and each encoder layer.
+ If True, in transformer layer, layer norm is applied before features are fed
+ to encoder layers. In encoder layer, two layer norms are applied before and after
+ self attention.
+ If False, in transformer layer, layer norm is applied after features are fed
+ to encoder layers. In encoder layer, two layer norms are applied after self
+ attention, before and after feed forward.
+ This option corresponds to "layer_norm_first" from fairseq.
+ Expected values are False for Base and True for Large arch.
+ layer_drop (float):
+ Probability to drop each encoder layer during training.
+ This option corresponds to "layerdrop" from fairseq.
+ Expected values are 0.1 for both Base and Large arch.
+
+ See Also:
+ * "encoder_embed_dim"
+ - Def and base
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L49-L51
+ - Large
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L64
+ * "dropout_input"
+ - Def, base and large
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L75-L78
+ * "conv_pos"
+ - Def, base and large
+ NOTE: The description is wrong.
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L204-L207
+ - Usage
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L756
+ * "conv_pos_groups"
+ - Def, base and large
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L208-L211
+ * "encoder_layers"
+ - Def and base
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L46-L48
+ - Large
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L63
+ * "encoder_attention_heads"
+ - Def and base
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L55-L57
+ - Large
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L66
+ * "attention_dropout"
+ - Def and base
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L66-L68
+ - Large
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L60
+ * "encoder_ffn_embed_dim"
+ - Def and base
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L52-L54
+ - Large
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L65
+ * "activation_dropout"
+ - Def
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L69-L71
+ - Base
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/finetuning/base_960h.yaml#L55
+ - Large
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/finetuning/vox_960h.yaml#L55
+ * "dropout"
+ - Def and base
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L63-L65
+ - Large
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L59
+ * "layer_norm_first"
+ - Def and base
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L91-L93
+ - Large
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml#L53
+ * "layerdrop"
+ - Def
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L72-L74
+ - Base
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/finetuning/base_960h.yaml#L54
+ - Large
+ https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/examples/wav2vec/config/finetuning/vox_960h.yaml#L54
+ """
+ feature_projection = FeatureProjection(in_features, embed_dim, dropout_input)
+ pos_conv = ConvolutionalPositionalEmbedding(embed_dim, pos_conv_kernel, pos_conv_groups)
+
+ # Original impl
+ # https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L768-L782
+ encoder_layers = nn.ModuleList()
+ for _ in range(num_layers):
+ attention = SelfAttention(
+ embed_dim=embed_dim,
+ num_heads=num_heads,
+ dropout=attention_dropout,
+ )
+ feed_forward = FeedForward(
+ io_features=embed_dim,
+ intermediate_features=ff_interm_features,
+ intermediate_dropout=ff_interm_dropout,
+ output_dropout=dropout,
+ )
+ encoder_layers.append(
+ EncoderLayer(
+ attention=attention,
+ dropout=dropout,
+ layer_norm_first=layer_norm_first,
+ feed_forward=feed_forward,
+ )
+ )
+ transformer = Transformer(
+ pos_conv_embed=pos_conv,
+ dropout=dropout,
+ layers=encoder_layers,
+ layer_norm_first=not layer_norm_first,
+ layer_drop=layer_drop,
+ )
+ return Encoder(feature_projection, transformer)
+
+
+def _get_wavlm_encoder(
+ in_features: int,
+ embed_dim: int,
+ dropout_input: float,
+ pos_conv_kernel: int,
+ pos_conv_groups: int,
+ num_layers: int,
+ num_heads: int,
+ num_buckets: int,
+ max_distance: int,
+ attention_dropout: float,
+ ff_interm_features: int,
+ ff_interm_dropout: float,
+ dropout: float,
+ layer_norm_first: bool,
+ layer_drop: float,
+) -> Encoder:
+ """
+ Construct encoder for WavLM model :cite:`chen2022wavlm`. The structure of the encoder and most of the argments are
+ the same as in :py:func:`_get_encoder` so refer there for documentation. The only difference from Wav2Vec2 encoder
+ is usage of `WavLMSelfAttention` instead of `SelfAttention` and two additional parameters: `num_buckets` and
+ `max_distance`.
+ Args:
+ in_features (int): See :py:func:`_get_encoder`.
+ embed_dim (int): See :py:func:`_get_encoder`.
+ dropout_input (float): See :py:func:`_get_encoder`.
+ pos_conv_kernel (int): See :py:func:`_get_encoder`.
+ pos_conv_groups (int): See :py:func:`_get_encoder`.
+ num_layers (int): See :py:func:`_get_encoder`.
+ num_heads (int): See :py:func:`_get_encoder`.
+ num_buckets (int): Number of buckets for relative position embedding.
+ max_distance (int): Maximum distance for relative position embedding.
+ attention_dropout (float): See :py:func:`_get_encoder`.
+ ff_interm_features (int): See :py:func:`_get_encoder`.
+ ff_interm_dropout (float): See :py:func:`_get_encoder`.
+ dropout (float): See :py:func:`_get_encoder`.
+ layer_norm_first (bool): See :py:func:`_get_encoder`.
+ layer_drop (float): See :py:func:`_get_encoder`.
+
+ """
+ feature_projection = FeatureProjection(in_features, embed_dim, dropout_input)
+ pos_conv = ConvolutionalPositionalEmbedding(embed_dim, pos_conv_kernel, pos_conv_groups)
+
+ # Original impl
+ # https://github.com/pytorch/fairseq/blob/425c36eafff535fe7337f8bdd5ace22ebacc78cb/fairseq/models/wav2vec/wav2vec2.py#L768-L782
+ encoder_layers = nn.ModuleList()
+ for i in range(num_layers):
+ attention = WavLMSelfAttention(
+ embed_dim=embed_dim,
+ num_heads=num_heads,
+ num_buckets=num_buckets,
+ max_distance=max_distance,
+ dropout=attention_dropout,
+ has_relative_attention_bias=(i == 0), # Position embedding is only necessary in the first layer.
+ )
+ feed_forward = FeedForward(
+ io_features=embed_dim,
+ intermediate_features=ff_interm_features,
+ intermediate_dropout=ff_interm_dropout,
+ output_dropout=dropout,
+ )
+ encoder_layers.append(
+ EncoderLayer(
+ attention=attention,
+ dropout=dropout,
+ layer_norm_first=layer_norm_first,
+ feed_forward=feed_forward,
+ )
+ )
+ transformer = Transformer(
+ pos_conv_embed=pos_conv,
+ dropout=dropout,
+ layers=encoder_layers,
+ layer_norm_first=not layer_norm_first,
+ layer_drop=layer_drop,
+ )
+ return Encoder(feature_projection, transformer)
+
+
+def _compute_mask_indices(
+ shape: Tuple[int, int],
+ padding_mask: Optional[Tensor],
+ mask_prob: float,
+ mask_length: int,
+ mask_type: str = "static",
+ mask_other: float = 0.0,
+ min_masks: int = 0,
+ no_overlap: bool = False,
+ min_space: int = 0,
+) -> Tensor:
+ """Computes random mask spans for a given shape.
+ Args:
+ shape (int, int): The shape for which to compute masks.
+ The first element is batch size and second is the number of frames.
+ padding_mask (Tensor or None): The padding mask of the same dimension as shape,
+ which will prevent masking padded elements.
+ mask_prob (float): Probability for each token to be chosen as start of the span to be masked.
+ This will be multiplied by number of timesteps divided by length of mask span to mask
+ approximately this percentage of all elements. However due to overlaps, the actual number
+ will be smaller (unless no_overlap is True).
+ mask_type (str): How to compute mask lengths. Options: [``static``, ``uniform``, ``normal``, ``poisson``].
+ ``static``: Fixed size
+ ``uniform``: Sample from uniform distribution [mask_other, mask_length*2]
+ ``normal``: Sample from normal distribution with mean ``mask_length`` and stdev ``mask_other``.
+ ``poisson``: Sample from possion distribution with lambda = ``mask_length``.
+ min_masks (int): Minimum number of masked spans.
+ no_overlap (bool): If false, will switch to an alternative recursive algorithm
+ that prevents spans from overlapping.
+ min_space (int): How many frames to keep unmasked between spans (Only used if no_overlap is True).
+
+ Returns:
+ (Tensor): The mask indices of dimension `[batch, frame]`.
+ """
+
+ batch_size, frame = shape
+ mask = torch.full((batch_size, frame), False)
+ # add a random number for probabilistic rounding
+ all_num_mask = int(mask_prob * frame / float(mask_length) + torch.rand(1))
+
+ all_num_mask = max(min_masks, all_num_mask)
+
+ mask_idcs = []
+ for i in range(batch_size):
+ if padding_mask is not None:
+ sz = frame - padding_mask[i].long().sum().item()
+ # add a random number for probabilistic rounding
+ num_mask = int(mask_prob * sz / float(mask_length) + torch.rand(1))
+ num_mask = max(min_masks, num_mask)
+ else:
+ sz = frame
+ num_mask = all_num_mask
+
+ if mask_type == "static":
+ lengths = torch.full((num_mask,), mask_length)
+ elif mask_type == "uniform":
+ lengths = torch.randint(int(mask_other), mask_length * 2 + 1, size=(num_mask,))
+ elif mask_type == "normal":
+ lengths = torch.normal(mask_length, mask_other, size=(num_mask,))
+ lengths = torch.maximum(torch.ones(1), torch.round(lengths)).int()
+ elif mask_type == "poisson":
+ lengths = torch.poisson(mask_length, size=(num_mask,))
+ lengths = torch.round(lengths).int()
+ else:
+ raise Exception(f"unknown mask selection: {mask_type}")
+
+ if sum(lengths) == 0:
+ lengths[0] = min(mask_length, sz - 1)
+
+ if no_overlap:
+ mask_idc = []
+
+ def arrange(s, e, length, keep_length):
+ span_start = torch.randint(s, e - length, size=(1,))
+ mask_idc.extend(span_start + i for i in range(length))
+
+ new_parts = []
+ if span_start - s - min_space >= keep_length:
+ new_parts.append((s, span_start - min_space + 1))
+ if e - span_start - keep_length - min_space > keep_length:
+ new_parts.append((span_start + length + min_space, e))
+ return new_parts
+
+ parts = [(0, sz)]
+ min_length = min(lengths)
+ for length in sorted(lengths, reverse=True):
+ lens = torch.tensor([e - s for s, e in parts], dtype=torch.int)
+ lens[lens < length + min_space] = 0
+ l_sum = lens.sum()
+ if l_sum == 0:
+ break
+ probs = lens / l_sum
+ c = torch.distributions.categorical.Categorical(probs).sample()
+ s, e = parts.pop(c)
+ parts.extend(arrange(s, e, length, min_length))
+ mask_idc = torch.tensor(mask_idc)
+ else:
+ min_len = min(lengths)
+ if sz - min_len <= num_mask:
+ min_len = sz - num_mask - 1
+
+ mask_idc = torch.randperm(sz - min_len)[:num_mask]
+ mask_idc = torch.tensor(
+ [mask_idc[j] + offset for j in range(len(mask_idc)) for offset in range(lengths[j])]
+ )
+
+ mask_idcs.append(torch.unique(mask_idc[mask_idc < sz]))
+
+ min_len = min([len(m) for m in mask_idcs])
+ for i, mask_idc in enumerate(mask_idcs):
+ if len(mask_idc) > min_len:
+ mask_idc = mask_idc[torch.randperm(len(mask_idc))[:min_len].long()]
+ mask[i, mask_idc] = True
+
+ return mask
+
+
+def _get_padding_mask(input: Tensor, lengths: Tensor) -> Tensor:
+ """Generate the padding mask given the padded input and the lengths Tensors.
+ Args:
+ input (Tensor): The padded Tensor of dimension `[batch, max_len, frequency]`.
+ lengths (Tensor): The lengths Tensor of dimension `[batch,]`.
+
+ Returns:
+ (Tensor): The padding mask.
+ """
+ batch_size, max_len, _ = input.shape
+ mask = torch.arange(max_len, device=lengths.device).expand(batch_size, max_len) >= lengths[:, None]
+ return mask
+
+
+class MaskGenerator(Module):
+ """Generate the masks for masked prediction.
+ Args:
+ encoder_embed_dim (int): The dimension of the transformer embedding output.
+ mask_prob (float): Probability for each token to be chosen as start of the span to be masked.
+ This will be multiplied by number of timesteps divided by length of mask span to mask
+ approximately this percentage of all elements. However due to overlaps, the actual number
+ will be smaller (unless no_overlap is True).
+ mask_selection (str): How to choose the mask length.
+ Options: [``static``, ``uniform``, ``normal``, ``poisson``].
+ mask_other (float): Secondary mask argument (used for more complex distributions).
+ mask_length (int): The lengths of the mask.
+ no_mask_overlap (bool): Whether to allow masks to overlap.
+ mask_min_space (int): Minimum space between spans (if no overlap is enabled).
+ mask_channel_prob (float): The probability of replacing a feature with 0.
+ mask_channel_selection (str): How to choose the mask length for channel masking.
+ Options: [``static``, ``uniform``, ``normal``, ``poisson``].
+ mask_channel_other (float): Secondary mask argument for channel masking(used for more complex distributions).
+ mask_channel_length (int): Minimum space between spans (if no overlap is enabled) for channel masking.
+ no_mask_channel_overlap (bool): Whether to allow channel masks to overlap.
+ mask_channel_min_space (int): Minimum space between spans for channel masking(if no overlap is enabled).
+ """
+
+ def __init__(
+ self,
+ encoder_embed_dim: int,
+ mask_prob: float,
+ mask_selection: str,
+ mask_other: float,
+ mask_length: int,
+ no_mask_overlap: bool,
+ mask_min_space: int,
+ mask_channel_prob: float,
+ mask_channel_selection: str,
+ mask_channel_other: float,
+ mask_channel_length: int,
+ no_mask_channel_overlap: bool,
+ mask_channel_min_space: int,
+ ):
+ super().__init__()
+ self.mask_prob = mask_prob
+ self.mask_selection = mask_selection
+ self.mask_other = mask_other
+ self.mask_length = mask_length
+ self.no_mask_overlap = no_mask_overlap
+ self.mask_min_space = mask_min_space
+ self.mask_channel_prob = mask_channel_prob
+ self.mask_channel_selection = mask_channel_selection
+ self.mask_channel_other = mask_channel_other
+ self.mask_channel_length = mask_channel_length
+ self.no_mask_channel_overlap = no_mask_channel_overlap
+ self.mask_channel_min_space = mask_channel_min_space
+ self.mask_embedding = Parameter(torch.FloatTensor(encoder_embed_dim))
+ torch.nn.init.uniform_(self.mask_embedding)
+
+ def forward(self, x: Tensor, padding_mask: Optional[Tensor]) -> Tensor:
+ """
+ Args:
+ x (Tensor): The encoded representations after feature extraction module.
+ padding_mask (Tensor or None): The padding mask of the same dimension as shape,
+ which will prevent masking padded elements.
+
+ Returns:
+ Tensor: The feature representations after masking.
+ Tensor: The generated mask indices.
+ """
+ B, T, C = x.shape
+ if self.mask_prob > 0:
+ mask_indices = _compute_mask_indices(
+ (B, T),
+ padding_mask,
+ self.mask_prob,
+ self.mask_length,
+ self.mask_selection,
+ self.mask_other,
+ min_masks=2,
+ no_overlap=self.no_mask_overlap,
+ min_space=self.mask_min_space,
+ )
+ mask_indices = mask_indices.to(x.device)
+ # change dtype of mask_embedding to x for mixed-precision training.
+ # see https://github.com/pytorch/audio/issues/2847 for details.
+ x[mask_indices] = self.mask_embedding.to(x.dtype)
+ else:
+ mask_indices = None
+
+ if self.mask_channel_prob > 0:
+ mask_channel_indices = _compute_mask_indices(
+ (B, C),
+ None,
+ self.mask_channel_prob,
+ self.mask_channel_length,
+ self.mask_channel_selection,
+ self.mask_channel_other,
+ no_overlap=self.no_mask_channel_overlap,
+ min_space=self.mask_channel_min_space,
+ )
+ mask_channel_indices = mask_channel_indices.to(x.device).unsqueeze(1).expand(-1, T, -1)
+ x[mask_channel_indices] = 0
+
+ return x, mask_indices
+
+
+def _compute_logits(
+ proj_x: Tensor,
+ target: Tensor,
+ label_embeddings: Parameter,
+) -> Tensor:
+ """Compute the logits of the embeddings.
+ Args:
+ proj_x (Tensor): The projected masked representations of dimension `[batch, frame, final_dim]`.
+ target (Tensor): The target Tensor of dimension `[batch, frame, final_dim]`.
+ label_embeddings (Parameter): The trainable embeddings of target of dimension `[num_class, final_dim]`.
+
+ Returns:
+ (Tensor): The logits of the inputs.
+ """
+ logit_temp = 0.1
+ pos = torch.index_select(label_embeddings, 0, target.long())
+ negs = label_embeddings.unsqueeze(1).expand(-1, proj_x.size(0), -1)
+ neg_is_pos = (pos == negs).all(-1)
+ pos = pos.unsqueeze(0)
+ targets = torch.cat([pos, negs], dim=0)
+
+ logits = torch.cosine_similarity(proj_x.float(), targets.float(), dim=-1).type_as(proj_x)
+ logits /= logit_temp
+ if neg_is_pos.any():
+ logits[1:][neg_is_pos] = float("-inf")
+ logits = logits.transpose(0, 1) # (num_x, num_cls+1)
+ return logits
+
+
+class LogitGenerator(Module):
+ """Generate the logits of masked and unmasked inputs.
+ Args:
+ encoder_embed_dim (int): The dimension of the transformer embedding output.
+ num_classes (int): The number of classes in the labels.
+ final_dim (int): Project final representations and targets to `final_dim`.
+ skip_masked (bool): If True, skip computing losses over masked frames.
+ skip_nomask (bool): If True, skip computing losses over unmasked frames.
+ """
+
+ def __init__(
+ self,
+ encoder_embed_dim: int,
+ num_classes: int,
+ final_dim: int,
+ skip_masked: bool,
+ skip_nomask: bool,
+ ):
+ super().__init__()
+ self.label_embeddings = Parameter(torch.FloatTensor(num_classes, final_dim))
+ torch.nn.init.uniform_(self.label_embeddings)
+ self.final_proj = torch.nn.Linear(encoder_embed_dim, final_dim)
+ self.skip_masked = skip_masked
+ self.skip_nomask = skip_nomask
+
+ def forward(self, x: Tensor, label: Tensor, mask_m: Tensor, mask_u: Tensor) -> Tuple[Tensor, Tensor]:
+ """
+ Args:
+ x (Tensor): The feature representation of the last transformer layer.
+ label (Tensor): The label Tensor of dimension `[batch, frame]`.
+ mask_m (Tensor): The masked indices of dimension `[batch, frame]`.
+ mask_u (Tensor): The unmasked indices of dimension `[batch, frame]`.
+
+ Returns:
+ Tensor: The logits of masked frames. Tensor of dimension `[masked_frame, final_dim]`.
+ Tensor: The logits of unmasked frames. Tensor of dimension `[unmasked_frame, final_dim]`.
+ """
+ proj_x = self.final_proj(x)
+ if self.skip_masked:
+ logit_m = None
+ else:
+ proj_x_m = proj_x[mask_m]
+ label_m = label[mask_m]
+ logit_m = _compute_logits(proj_x_m, label_m, self.label_embeddings)
+
+ if self.skip_nomask:
+ logit_u = None
+ else:
+ proj_x_u = proj_x[mask_u]
+ label_u = label[mask_u]
+ logit_u = _compute_logits(proj_x_u, label_u, self.label_embeddings)
+ return logit_m, logit_u
+
+
+class GradMultiply(torch.autograd.Function):
+ @staticmethod
+ def forward(ctx, x, scale):
+ ctx.scale = scale
+ res = x.new(x)
+ return res
+
+ @staticmethod
+ def backward(ctx, grad):
+ return grad * ctx.scale, None
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/model.py b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..254122f0eee21906ec50f3d4238a5b3024e74a0a
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/model.py
@@ -0,0 +1,1579 @@
+import math
+from typing import List, Optional, Tuple
+
+import torch
+from torch import Tensor
+from torch.nn import Module
+
+from . import components
+
+
+class Wav2Vec2Model(Module):
+ """Acoustic model used in *wav2vec 2.0* :cite:`baevski2020wav2vec`.
+
+ Note:
+ To build the model, please use one of the factory functions.
+
+ See Also:
+ * :class:`torchaudio.pipelines.Wav2Vec2Bundle`: Pretrained models (without fine-tuning)
+ * :class:`torchaudio.pipelines.Wav2Vec2ASRBundle`: ASR pipelines with pretrained models.
+
+ Args:
+ feature_extractor (torch.nn.Module):
+ Feature extractor that extracts feature vectors from raw audio Tensor.
+
+ encoder (torch.nn.Module):
+ Encoder that converts the audio features into the sequence of probability
+ distribution (in negative log-likelihood) over labels.
+
+ aux (torch.nn.Module or None, optional):
+ Auxiliary module. If provided, the output from encoder is passed to this module.
+ """ # noqa: E501
+
+ def __init__(
+ self,
+ feature_extractor: Module,
+ encoder: Module,
+ aux: Optional[Module] = None,
+ ):
+ super().__init__()
+ self.feature_extractor = feature_extractor
+ self.encoder = encoder
+ self.aux = aux
+
+ @torch.jit.export
+ def extract_features(
+ self,
+ waveforms: Tensor,
+ lengths: Optional[Tensor] = None,
+ num_layers: Optional[int] = None,
+ ) -> Tuple[List[Tensor], Optional[Tensor]]:
+ """Extract feature vectors from raw waveforms
+
+ This returns the list of outputs from the intermediate layers of
+ transformer block in encoder.
+
+ Args:
+ waveforms (Tensor): Audio tensor of shape `(batch, frames)`.
+ lengths (Tensor or None, optional):
+ Indicates the valid length of each audio in the batch.
+ Shape: `(batch, )`.
+ When the ``waveforms`` contains audios with different durations,
+ by providing ``lengths`` argument, the model will compute
+ the corresponding valid output lengths and apply proper mask in
+ transformer attention layer.
+ If ``None``, it is assumed that the entire audio waveform
+ length is valid.
+ num_layers (int or None, optional):
+ If given, limit the number of intermediate layers to go through.
+ Providing `1` will stop the computation after going through one
+ intermediate layers. If not given, the outputs from all the
+ intermediate layers are returned.
+
+ Returns:
+ (List[Tensor], Optional[Tensor]):
+ List of Tensors
+ Features from requested layers.
+ Each Tensor is of shape: `(batch, time frame, feature dimension)`
+ Tensor or None
+ If ``lengths`` argument was provided, a Tensor of shape `(batch, )`
+ is returned.
+ It indicates the valid length in time axis of each feature Tensor.
+ """
+ x, lengths = self.feature_extractor(waveforms, lengths)
+ x = self.encoder.extract_features(x, lengths, num_layers)
+ return x, lengths
+
+ def forward(
+ self,
+ waveforms: Tensor,
+ lengths: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Optional[Tensor]]:
+ """Compute the sequence of probability distribution over labels.
+
+ Args:
+ waveforms (Tensor): Audio tensor of shape `(batch, frames)`.
+ lengths (Tensor or None, optional):
+ Indicates the valid length of each audio in the batch.
+ Shape: `(batch, )`.
+ When the ``waveforms`` contains audios with different durations,
+ by providing ``lengths`` argument, the model will compute
+ the corresponding valid output lengths and apply proper mask in
+ transformer attention layer.
+ If ``None``, it is assumed that all the audio in ``waveforms``
+ have valid length. Default: ``None``.
+
+ Returns:
+ (Tensor, Optional[Tensor]):
+ Tensor
+ The sequences of probability distribution (in logit) over labels.
+ Shape: `(batch, frames, num labels)`.
+ Tensor or None
+ If ``lengths`` argument was provided, a Tensor of shape `(batch, )`
+ is returned.
+ It indicates the valid length in time axis of the output Tensor.
+ """
+ x, lengths = self.feature_extractor(waveforms, lengths)
+ x = self.encoder(x, lengths)
+ if self.aux is not None:
+ x = self.aux(x)
+ return x, lengths
+
+
+class HuBERTPretrainModel(Module):
+ """HuBERTPretrainModel()
+
+ HuBERT model used for pretraining in *HuBERT* :cite:`hsu2021hubert`.
+
+ Note:
+ To build the model, please use one of the factory functions.
+
+ See Also:
+ `HuBERT Pre-training and Fine-tuning Recipes
+ `__
+
+ Args:
+ wav2vec2 (Wav2Vec2Model):
+ Wav2Vec2 encoder that generates the transformer outputs.
+
+ mask_generator (torch.nn.Module):
+ Mask generator that generates the mask for masked prediction during the training.
+
+ logit_generator (torch.nn.Module):
+ Logit generator that predicts the logits of the masked and unmasked inputs.
+
+ feature_grad_mult (float or None):
+ The factor to scale the convolutional feature extraction layer gradients by.
+ If ``None``, the gradients of feature extraction layers are not affected.
+ The scale factor will not affect the forward pass.
+ """
+
+ def __init__(
+ self,
+ wav2vec2: Wav2Vec2Model,
+ mask_generator: Module,
+ logit_generator: Module,
+ feature_grad_mult: Optional[float],
+ ):
+ super().__init__()
+ self.wav2vec2 = wav2vec2
+ self.mask_generator = mask_generator
+ self.logit_generator = logit_generator
+ if feature_grad_mult is not None and not 0.0 < feature_grad_mult < 1.0:
+ raise ValueError(
+ f"The value of `feature_grad_mult` must be ``None``or between (0, 1). Found {feature_grad_mult}"
+ )
+ self.feature_grad_mult = feature_grad_mult
+
+ def forward(
+ self,
+ waveforms: Tensor,
+ labels: Tensor,
+ audio_lengths: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Optional[Tensor]]:
+ """Compute the sequence of probability distribution over labels.
+
+ Args:
+ waveforms (Tensor): Audio tensor of dimension `[batch, frames]`.
+ labels (Tensor): Label for pre-training. A Tensor of dimension `[batch, frames]`.
+ audio_lengths (Tensor or None, optional):
+ Indicates the valid length of each audio in the batch.
+ Shape: `[batch, ]`.
+ When the ``waveforms`` contains audios with different durations,
+ by providing ``lengths`` argument, the model will compute
+ the corresponding valid output lengths and apply proper mask in
+ transformer attention layer.
+ If ``None``, it is assumed that all the audio in ``waveforms``
+ have valid length. Default: ``None``.
+
+ Returns:
+ (Tensor, Tensor, Tensor):
+ Tensor
+ The masked sequences of probability distribution (in logit).
+ Shape: `(masked_frames, num labels)`.
+ Tensor
+ The unmasked sequence of probability distribution (in logit).
+ Shape: `(unmasked_frames, num labels)`.
+ Tensor
+ The feature mean value for additional penalty loss.
+ Shape: `(1,)`.
+ """
+ x, lengths = self.wav2vec2.feature_extractor(waveforms, audio_lengths)
+ if self.feature_grad_mult is not None and self.feature_grad_mult < 1.0:
+ x = components.GradMultiply.apply(x, self.feature_grad_mult)
+ features_pen = x.float().pow(2).mean()
+ if lengths is not None:
+ padding_mask = components._get_padding_mask(x, lengths)
+ else:
+ padding_mask = None
+ x, attention_mask = self.wav2vec2.encoder._preprocess(x, lengths)
+ x, mask = self.mask_generator(x, padding_mask)
+ x = self.wav2vec2.encoder.transformer(x, attention_mask=attention_mask)
+ if x.shape[1] != labels.shape[1]:
+ raise ValueError("The length of label must match that of HuBERT model output")
+ if padding_mask is not None:
+ mask_m = torch.logical_and(~padding_mask, mask)
+ mask_u = torch.logical_and(~padding_mask, ~mask_m)
+ else:
+ mask_m = mask
+ mask_u = ~mask_m
+
+ logit_m, logit_u = self.logit_generator(x, labels, mask_m, mask_u)
+
+ return logit_m, logit_u, features_pen
+
+
+def wav2vec2_model(
+ extractor_mode: str,
+ extractor_conv_layer_config: Optional[List[Tuple[int, int, int]]],
+ extractor_conv_bias: bool,
+ encoder_embed_dim: int,
+ encoder_projection_dropout: float,
+ encoder_pos_conv_kernel: int,
+ encoder_pos_conv_groups: int,
+ encoder_num_layers: int,
+ encoder_num_heads: int,
+ encoder_attention_dropout: float,
+ encoder_ff_interm_features: int,
+ encoder_ff_interm_dropout: float,
+ encoder_dropout: float,
+ encoder_layer_norm_first: bool,
+ encoder_layer_drop: float,
+ aux_num_out: Optional[int],
+) -> Wav2Vec2Model:
+ """Builds custom :class:`~torchaudio.models.Wav2Vec2Model`.
+
+ Note:
+ The "feature extractor" below corresponds to
+ `ConvFeatureExtractionModel `__
+ in the original ``fairseq`` implementation.
+ This is referred as "(convolutional) feature encoder" in the *wav2vec 2.0*
+ :cite:`baevski2020wav2vec` paper.
+
+ The "encoder" below corresponds to `TransformerEncoder `__,
+ and this is referred as "Transformer" in the paper.
+
+ Args:
+ extractor_mode (str): Operation mode of feature extractor.
+ Valid values are ``"group_norm"`` or ``"layer_norm"``.
+ If ``"group_norm"``, then a single normalization is applied
+ in the first convolution block. Otherwise, all the convolution
+ blocks will have layer normalization.
+
+ This option corresponds to ``extractor_mode`` from ``fairseq``.
+ extractor_conv_layer_config (list of integer tuples or None):
+ Configuration of convolution layers in feature extractor.
+ List of convolution configuration,
+ i.e. ``[(output_channel, kernel_size, stride), ...]``
+
+ If ``None`` is provided, then the following default value is used.
+
+ .. code-block:: python
+
+ [
+ (512, 10, 5),
+ (512, 3, 2),
+ (512, 3, 2),
+ (512, 3, 2),
+ (512, 3, 2),
+ (512, 2, 2),
+ (512, 2, 2),
+ ]
+
+ This option corresponds to ``conv_feature_layers`` from ``fairseq``.
+
+ extractor_conv_bias (bool):
+ Whether to include bias term to each convolution operation.
+
+ This option corresponds to ``conv_bias`` from ``fairseq``.
+
+ encoder_embed_dim (int):
+ The dimension of embedding in encoder.
+
+ This option corresponds to ``encoder_embed_dim`` from ``fairseq``.
+
+ encoder_projection_dropout (float):
+ The dropout probability applied after the input feature is projected
+ to ``encoder_embed_dim``.
+
+ This option corresponds to ``dropout_input`` from ``fairseq``.
+
+ encoder_pos_conv_kernel (int):
+ The kernel size of convolutional positional embeddings.
+
+ This option corresponds to ``conv_pos`` from ``fairseq``.
+
+ encoder_pos_conv_groups (int):
+ The number of groups of convolutional positional embeddings.
+
+ This option corresponds to ``conv_pos_groups`` from ``fairseq``.
+
+ encoder_num_layers (int):
+ The number of self attention layers in transformer block.
+
+ This option corresponds to ``encoder_layers`` from ``fairseq``.
+
+ encoder_num_heads (int):
+ The number of heads in self attention layers.
+
+ This option corresponds to ``encoder_attention_heads`` from ``fairseq``.
+
+ encoder_attention_dropout (float):
+ The dropout probability applied after softmax in self-attention layer.
+
+ This option corresponds to ``attention_dropout`` from ``fairseq``.
+
+ encoder_ff_interm_features (int):
+ The dimension of hidden features in feed forward layer.
+
+ This option corresponds to ``encoder_ffn_embed_dim`` from ``fairseq``.
+
+ encoder_ff_interm_dropout (float):
+ The dropout probability applied in feedforward layer.
+
+ This option correspinds to ``activation_dropout`` from ``fairseq``.
+
+ encoder_dropout (float):
+ The dropout probability applied at the end of feed forward layer.
+
+ This option corresponds to ``dropout`` from ``fairseq``.
+
+ encoder_layer_norm_first (bool):
+ Control the order of layer norm in transformer layer and each encoder layer.
+ If True, in transformer layer, layer norm is applied before features are fed
+ to encoder layers. In encoder layer, two layer norms are applied before and after
+ self attention.
+ If False, in transformer layer, layer norm is applied after features are fed
+ to encoder layers. In encoder layer, two layer norms are applied after self
+ attention, before and after feed forward.
+
+ This option corresponds to ``layer_norm_first`` from ``fairseq``.
+
+ encoder_layer_drop (float):
+ Probability to drop each encoder layer during training.
+
+ This option corresponds to ``layerdrop`` from ``fairseq``.
+
+ aux_num_out (int or None):
+ When provided, attach an extra linear layer on top of encoder, which can be
+ used for fine-tuning.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """ # noqa: E501
+ if extractor_conv_layer_config is None:
+ extractor_conv_layer_config = [(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512, 2, 2)] * 2
+
+ feature_extractor = components._get_feature_extractor(
+ extractor_mode, extractor_conv_layer_config, extractor_conv_bias
+ )
+ encoder = components._get_encoder(
+ in_features=extractor_conv_layer_config[-1][0],
+ embed_dim=encoder_embed_dim,
+ dropout_input=encoder_projection_dropout,
+ pos_conv_kernel=encoder_pos_conv_kernel,
+ pos_conv_groups=encoder_pos_conv_groups,
+ num_layers=encoder_num_layers,
+ num_heads=encoder_num_heads,
+ attention_dropout=encoder_attention_dropout,
+ ff_interm_features=encoder_ff_interm_features,
+ ff_interm_dropout=encoder_ff_interm_dropout,
+ dropout=encoder_dropout,
+ layer_norm_first=encoder_layer_norm_first,
+ layer_drop=encoder_layer_drop,
+ )
+ aux = None
+ if aux_num_out is not None:
+ aux = torch.nn.Linear(in_features=encoder_embed_dim, out_features=aux_num_out)
+ return Wav2Vec2Model(feature_extractor, encoder, aux)
+
+
+def wav2vec2_base(
+ encoder_projection_dropout: float = 0.1,
+ encoder_attention_dropout: float = 0.1,
+ encoder_ff_interm_dropout: float = 0.1,
+ encoder_dropout: float = 0.1,
+ encoder_layer_drop: float = 0.1,
+ aux_num_out: Optional[int] = None,
+) -> Wav2Vec2Model:
+ """Builds "base" :class:`~torchaudio.models.Wav2Vec2Model` from *wav2vec 2.0* :cite:`baevski2020wav2vec`
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_layer_drop (float):
+ See :py:func:`wav2vec2_model`.
+ aux_num_out (int or None, optional):
+ See :py:func:`wav2vec2_model`.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """ # noqa: E501
+ return wav2vec2_model(
+ extractor_mode="group_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=False,
+ encoder_embed_dim=768,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=12,
+ encoder_num_heads=12,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=3072,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=False,
+ encoder_layer_drop=encoder_layer_drop,
+ aux_num_out=aux_num_out,
+ )
+
+
+def wav2vec2_large(
+ encoder_projection_dropout: float = 0.1,
+ encoder_attention_dropout: float = 0.1,
+ encoder_ff_interm_dropout: float = 0.1,
+ encoder_dropout: float = 0.1,
+ encoder_layer_drop: float = 0.1,
+ aux_num_out: Optional[int] = None,
+) -> Wav2Vec2Model:
+ """Builds "large" :class:`~torchaudio.models.Wav2Vec2Model` from *wav2vec 2.0* :cite:`baevski2020wav2vec`
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_layer_drop (float):
+ See :py:func:`wav2vec2_model`.
+ aux_num_out (int or None, optional):
+ See :py:func:`wav2vec2_model`.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """ # noqa: E501
+ return wav2vec2_model(
+ extractor_mode="group_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=False,
+ encoder_embed_dim=1024,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=24,
+ encoder_num_heads=16,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=4096,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=False,
+ encoder_layer_drop=encoder_layer_drop,
+ aux_num_out=aux_num_out,
+ )
+
+
+def wav2vec2_large_lv60k(
+ encoder_projection_dropout: float = 0.1,
+ encoder_attention_dropout: float = 0.0,
+ encoder_ff_interm_dropout: float = 0.1,
+ encoder_dropout: float = 0.0,
+ encoder_layer_drop: float = 0.1,
+ aux_num_out: Optional[int] = None,
+) -> Wav2Vec2Model:
+ """Builds "large lv-60k" :class:`~torchaudio.models.Wav2Vec2Model` from *wav2vec 2.0* :cite:`baevski2020wav2vec`
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_layer_drop (float):
+ See :py:func:`wav2vec2_model`.
+ aux_num_out (int or None, optional):
+ See :py:func:`wav2vec2_model`.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """ # noqa: E501
+ return wav2vec2_model(
+ extractor_mode="layer_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=True,
+ encoder_embed_dim=1024,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=24,
+ encoder_num_heads=16,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=4096,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=True,
+ encoder_layer_drop=encoder_layer_drop,
+ aux_num_out=aux_num_out,
+ )
+
+
+def hubert_base(
+ encoder_projection_dropout: float = 0.1,
+ encoder_attention_dropout: float = 0.1,
+ encoder_ff_interm_dropout: float = 0.0,
+ encoder_dropout: float = 0.1,
+ encoder_layer_drop: float = 0.05,
+ aux_num_out: Optional[int] = None,
+) -> Wav2Vec2Model:
+ """Builds "base" :class:`HuBERT ` from *HuBERT* :cite:`hsu2021hubert`
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_layer_drop (float):
+ See :py:func:`wav2vec2_model`.
+ aux_num_out (int or None, optional):
+ See :py:func:`wav2vec2_model`.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """ # noqa: E501
+ return wav2vec2_model(
+ extractor_mode="group_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=False,
+ encoder_embed_dim=768,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=12,
+ encoder_num_heads=12,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=3072,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=False,
+ encoder_layer_drop=encoder_layer_drop,
+ aux_num_out=aux_num_out,
+ )
+
+
+def hubert_large(
+ encoder_projection_dropout: float = 0.0,
+ encoder_attention_dropout: float = 0.0,
+ encoder_ff_interm_dropout: float = 0.0,
+ encoder_dropout: float = 0.0,
+ encoder_layer_drop: float = 0.0,
+ aux_num_out: Optional[int] = None,
+) -> Wav2Vec2Model:
+ """Builds "large" :class:`HuBERT ` from *HuBERT* :cite:`hsu2021hubert`
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_layer_drop (float):
+ See :py:func:`wav2vec2_model`.
+ aux_num_out (int or None, optional):
+ See :py:func:`wav2vec2_model`.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """ # noqa: E501
+ return wav2vec2_model(
+ extractor_mode="layer_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=False,
+ encoder_embed_dim=1024,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=24,
+ encoder_num_heads=16,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=4096,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=True,
+ encoder_layer_drop=encoder_layer_drop,
+ aux_num_out=aux_num_out,
+ )
+
+
+def hubert_xlarge(
+ encoder_projection_dropout: float = 0.0,
+ encoder_attention_dropout: float = 0.0,
+ encoder_ff_interm_dropout: float = 0.0,
+ encoder_dropout: float = 0.0,
+ encoder_layer_drop: float = 0.0,
+ aux_num_out: Optional[int] = None,
+) -> Wav2Vec2Model:
+ """Builds "extra large" :class:`HuBERT ` from *HuBERT* :cite:`hsu2021hubert`
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_dropout (float):
+ See :py:func:`wav2vec2_model`.
+ encoder_layer_drop (float):
+ See :py:func:`wav2vec2_model`.
+ aux_num_out (int or None, optional):
+ See :py:func:`wav2vec2_model`.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """ # noqa: E501
+ return wav2vec2_model(
+ extractor_mode="layer_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=False,
+ encoder_embed_dim=1280,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=48,
+ encoder_num_heads=16,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=5120,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=True,
+ encoder_layer_drop=encoder_layer_drop,
+ aux_num_out=aux_num_out,
+ )
+
+
+def _init_hubert_pretrain_model(module):
+ if isinstance(module, components.ConvLayerBlock):
+ torch.nn.init.kaiming_normal_(module.conv.weight)
+ elif isinstance(module, components.ConvolutionalPositionalEmbedding):
+ # normalize the weight to normal distribution.
+ std = math.sqrt(4.0 / (module.embed_dim * module.kernel_size))
+ torch.nn.init.normal_(module.conv.weight, mean=0.0, std=std)
+ torch.nn.init.constant_(module.conv.bias, 0.0)
+ elif isinstance(module, components.SelfAttention):
+ # normalize the query, key, value, and out_proj parameters in self attention module.
+ torch.nn.init.xavier_uniform_(module.k_proj.weight, gain=1 / math.sqrt(2))
+ torch.nn.init.xavier_uniform_(module.v_proj.weight, gain=1 / math.sqrt(2))
+ torch.nn.init.xavier_uniform_(module.q_proj.weight, gain=1 / math.sqrt(2))
+ torch.nn.init.xavier_uniform_(module.out_proj.weight)
+ torch.nn.init.constant_(module.out_proj.bias, 0.0)
+ elif isinstance(module, components.Transformer):
+ module.apply(components._init_transformer_params)
+ else:
+ pass
+
+
+def hubert_pretrain_model(
+ extractor_mode: str,
+ extractor_conv_layer_config: Optional[List[Tuple[int, int, int]]],
+ extractor_conv_bias: bool,
+ encoder_embed_dim: int,
+ encoder_projection_dropout: float,
+ encoder_pos_conv_kernel: int,
+ encoder_pos_conv_groups: int,
+ encoder_num_layers: int,
+ encoder_num_heads: int,
+ encoder_attention_dropout: float,
+ encoder_ff_interm_features: int,
+ encoder_ff_interm_dropout: float,
+ encoder_dropout: float,
+ encoder_layer_norm_first: bool,
+ encoder_layer_drop: float,
+ mask_prob: float,
+ mask_selection: str,
+ mask_other: float,
+ mask_length: int,
+ no_mask_overlap: bool,
+ mask_min_space: int,
+ mask_channel_prob: float,
+ mask_channel_selection: str,
+ mask_channel_other: float,
+ mask_channel_length: int,
+ no_mask_channel_overlap: bool,
+ mask_channel_min_space: int,
+ skip_masked: bool,
+ skip_nomask: bool,
+ num_classes: int,
+ final_dim: int,
+ feature_grad_mult: Optional[float],
+) -> HuBERTPretrainModel:
+ """Builds custom :class:`HuBERTPretrainModel` for training from scratch
+
+ Note:
+ The "feature extractor" below corresponds to
+ `ConvFeatureExtractionModel `__
+ in the original ``fairseq`` implementation.
+ This is referred as "(convolutional) feature encoder" in the *wav2vec 2.0*
+ :cite:`baevski2020wav2vec` paper.
+
+ The "encoder" below corresponds to `TransformerEncoder `__,
+ and this is referred as "Transformer" in the paper.
+
+ Args:
+ extractor_mode (str): Operation mode of feature extractor.
+ Valid values are ``"group_norm"`` or ``"layer_norm"``.
+ If ``"group_norm"``, then a single normalization is applied
+ in the first convolution block. Otherwise, all the convolution
+ blocks will have layer normalization.
+
+ This option corresponds to ``extractor_mode`` from ``fairseq``.
+
+ extractor_conv_layer_config (list of integer tuples or None):
+ Configuration of convolution layers in feature extractor.
+ List of convolution configuration,
+ i.e. ``[(output_channel, kernel_size, stride), ...]``
+
+ If ``None`` is provided, then the following default value is used.
+
+ .. code-block:: python
+
+ [
+ (512, 10, 5),
+ (512, 3, 2),
+ (512, 3, 2),
+ (512, 3, 2),
+ (512, 3, 2),
+ (512, 2, 2),
+ (512, 2, 2),
+ ]
+
+ This option corresponds to ``conv_feature_layers`` from ``fairseq``.
+
+ extractor_conv_bias (bool):
+ Whether to include bias term to each convolution operation.
+
+ This option corresponds to ``conv_bias`` from ``fairseq``.
+
+ encoder_embed_dim (int):
+ The dimension of embedding in encoder.
+
+ This option corresponds to ``encoder_embed_dim`` from ``fairseq``.
+
+ encoder_projection_dropout (float):
+ The dropout probability applied after the input feature is projected
+ to ``encoder_embed_dim``.
+
+ This option corresponds to ``dropout_input`` from ``fairseq``.
+
+ encoder_pos_conv_kernel (int):
+ The kernel size of convolutional positional embeddings.
+
+ This option corresponds to ``conv_pos`` from ``fairseq``.
+
+ encoder_pos_conv_groups (int):
+ The number of groups of convolutional positional embeddings.
+
+ This option corresponds to ``conv_pos_groups`` from ``fairseq``.
+
+ encoder_num_layers (int):
+ The number of self attention layers in transformer block.
+
+ This option corresponds to ``encoder_layers`` from ``fairseq``.
+
+ encoder_num_heads (int):
+ The number of heads in self attention layers.
+
+ This option corresponds to ``encoder_attention_heads`` from ``fairseq``.
+
+ encoder_attention_dropout (float):
+ The dropout probability applied after softmax in self-attention layer.
+
+ This option corresponds to ``attention_dropout`` from ``fairseq``.
+
+ encoder_ff_interm_features (int):
+ The dimension of hidden features in feed forward layer.
+
+ This option corresponds to ``encoder_ffn_embed_dim`` from ``fairseq``.
+
+ encoder_ff_interm_dropout (float):
+ The dropout probability applied in feedforward layer.
+
+ This option correspinds to ``activation_dropout`` from ``fairseq``.
+
+ encoder_dropout (float):
+ The dropout probability applied at the end of feed forward layer.
+
+ This option corresponds to ``dropout`` from ``fairseq``.
+
+ encoder_layer_norm_first (bool):
+ Control the order of layer norm in transformer layer and each encoder layer.
+ If True, in transformer layer, layer norm is applied before features are fed
+ to encoder layers. In encoder layer, two layer norms are applied before and after
+ self attention.
+ If False, in transformer layer, layer norm is applied after features are fed
+ to encoder layers. In encoder layer, two layer norms are applied after self
+ attention, before and after feed forward.
+
+ This option corresponds to ``layer_norm_first`` from ``fairseq``.
+
+ encoder_layer_drop (float):
+ Probability to drop each encoder layer during training.
+
+ This option corresponds to ``layerdrop`` from ``fairseq``.
+
+ mask_prob (float):
+ Probability for each token to be chosen as start of the span to be masked. this will be multiplied by
+ number of timesteps divided by length of mask span to mask approximately this percentage of all elements.
+ However due to overlaps, the actual number will be smaller (unless no_overlap is True).
+
+ This option corresponds to ``mask_prob`` from ``fairseq``.
+
+ mask_selection (str):
+ How to choose the mask length. Options: [``static``, ``uniform``, ``normal``, ``poisson``].
+
+ This option corresponds to ``mask_selection`` from ``fairseq``.
+
+ mask_other (float):
+ Secondary mask argument (used for more complex distributions).
+
+ This option corresponds to ``mask_other`` from ``fairseq``.
+
+ mask_length (int):
+ The lengths of the mask.
+
+ This option corresponds to ``mask_length`` from ``fairseq``.
+
+ no_mask_overlap (bool):
+ Whether to allow masks to overlap.
+
+ This option corresponds to ``no_mask_overlap`` from ``fairseq``.
+
+ mask_min_space (int):
+ Minimum space between spans (if no overlap is enabled).
+
+ This option corresponds to ``mask_min_space`` from ``fairseq``.
+
+ mask_channel_prob: (float):
+ The probability of replacing a feature with 0.
+
+ This option corresponds to ``mask_channel_prob`` from ``fairseq``.
+
+ mask_channel_selection (str):
+ How to choose the mask length for channel masking. Options: [``static``, ``uniform``, ``normal``, ``poisson``].
+
+ This option corresponds to ``mask_channel_selection`` from ``fairseq``.
+
+ mask_channel_other (float):
+ Secondary mask argument for channel masking(used for more complex distributions).
+
+ This option corresponds to ``mask_channel_other`` from ``fairseq``.
+
+ mask_channel_length (int):
+ Minimum space between spans (if no overlap is enabled) for channel masking.
+
+ This option corresponds to ``mask_channel_length`` from ``fairseq``.
+
+ no_mask_channel_overlap (bool):
+ Whether to allow channel masks to overlap.
+
+ This option corresponds to ``no_mask_channel_overlap`` from ``fairseq``.
+
+ mask_channel_min_space (int):
+ Minimum space between spans for channel masking(if no overlap is enabled).
+
+ This option corresponds to ``mask_channel_min_space`` from ``fairseq``.
+
+ skip_masked (bool):
+ If True, skip computing losses over masked frames.
+
+ This option corresponds to ``skip_masked`` from ``fairseq``.
+
+ skip_nomask (bool):
+ If True, skip computing losses over unmasked frames.
+
+ This option corresponds to ``skip_nomask`` from ``fairseq``.
+
+ num_classes (int):
+ The number of classes in the labels.
+
+ final_dim (int):
+ Project final representations and targets to `final_dim`.
+
+ This option corresponds to ``final_dim`` from ``fairseq``.
+
+ feature_grad_mult (float or None):
+ The factor to scale the convolutional feature extraction layer gradients by.
+ The scale factor will not affect the forward pass.
+
+ This option corresponds to ``feature_grad_mult`` from ``fairseq``.
+
+ Returns:
+ HuBERTPretrainModel:
+ The resulting model.
+ """ # noqa: E501
+ if extractor_conv_layer_config is None:
+ extractor_conv_layer_config = [(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512, 2, 2)] * 2
+
+ feature_extractor = components._get_feature_extractor(
+ extractor_mode, extractor_conv_layer_config, extractor_conv_bias
+ )
+ encoder = components._get_encoder(
+ in_features=extractor_conv_layer_config[-1][0],
+ embed_dim=encoder_embed_dim,
+ dropout_input=encoder_projection_dropout,
+ pos_conv_kernel=encoder_pos_conv_kernel,
+ pos_conv_groups=encoder_pos_conv_groups,
+ num_layers=encoder_num_layers,
+ num_heads=encoder_num_heads,
+ attention_dropout=encoder_attention_dropout,
+ ff_interm_features=encoder_ff_interm_features,
+ ff_interm_dropout=encoder_ff_interm_dropout,
+ dropout=encoder_dropout,
+ layer_norm_first=encoder_layer_norm_first,
+ layer_drop=encoder_layer_drop,
+ )
+ wav2vec2 = Wav2Vec2Model(feature_extractor, encoder)
+ mask_generator = components.MaskGenerator(
+ encoder_embed_dim,
+ mask_prob,
+ mask_selection,
+ mask_other,
+ mask_length,
+ no_mask_overlap,
+ mask_min_space,
+ mask_channel_prob,
+ mask_channel_selection,
+ mask_channel_other,
+ mask_channel_length,
+ no_mask_channel_overlap,
+ mask_channel_min_space,
+ )
+ logit_generator = components.LogitGenerator(
+ encoder_embed_dim,
+ num_classes,
+ final_dim,
+ skip_masked,
+ skip_nomask,
+ )
+ model = HuBERTPretrainModel(
+ wav2vec2=wav2vec2,
+ mask_generator=mask_generator,
+ logit_generator=logit_generator,
+ feature_grad_mult=feature_grad_mult,
+ )
+ # initialize the model for pre-training
+ model.apply(_init_hubert_pretrain_model)
+ return model
+
+
+def hubert_pretrain_base(
+ encoder_projection_dropout: float = 0.1,
+ encoder_attention_dropout: float = 0.1,
+ encoder_ff_interm_dropout: float = 0.0,
+ encoder_dropout: float = 0.1,
+ encoder_layer_drop: float = 0.05,
+ mask_prob: float = 0.8,
+ mask_channel_prob: float = 0.0,
+ mask_channel_length: int = 10,
+ feature_grad_mult: Optional[float] = 0.1,
+ num_classes: int = 100,
+) -> HuBERTPretrainModel:
+ """Builds "base" :class:`HuBERTPretrainModel` from *HuBERT* :cite:`hsu2021hubert` for pretraining.
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`hubert_pretrain_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`hubert_pretrain_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`hubert_pretrain_model`.
+ encoder_dropout (float):
+ See :py:func:`hubert_pretrain_model`.
+ encoder_layer_drop (float):
+ See :py:func:`hubert_pretrain_model`.
+ mask_prob (float):
+ See :py:func:`hubert_pretrain_model`.
+ mask_channel_prob (float):
+ See :py:func:`hubert_pretrain_model`.
+ mask_channel_length (int):
+ See :py:func:`hubert_pretrain_model`.
+ feature_grad_mult (float or None):
+ See :py:func:`hubert_pretrain_model`.
+ num_classes (int, optional):
+ See :py:func:`hubert_pretrain_model`.
+
+ Returns:
+ HuBERTPretrainModel:
+ The resulting model.
+ """ # noqa: E501
+ return hubert_pretrain_model(
+ extractor_mode="group_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=False,
+ encoder_embed_dim=768,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=12,
+ encoder_num_heads=12,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=3072,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=False,
+ encoder_layer_drop=encoder_layer_drop,
+ mask_prob=mask_prob,
+ mask_selection="static",
+ mask_other=0.0,
+ mask_length=10,
+ no_mask_overlap=False,
+ mask_min_space=1,
+ mask_channel_prob=mask_channel_prob,
+ mask_channel_selection="static",
+ mask_channel_other=0.0,
+ mask_channel_length=mask_channel_length,
+ no_mask_channel_overlap=False,
+ mask_channel_min_space=1,
+ skip_masked=False,
+ skip_nomask=False,
+ num_classes=num_classes,
+ final_dim=256,
+ feature_grad_mult=feature_grad_mult,
+ )
+
+
+def hubert_pretrain_large(
+ encoder_projection_dropout: float = 0.0,
+ encoder_attention_dropout: float = 0.0,
+ encoder_ff_interm_dropout: float = 0.0,
+ encoder_dropout: float = 0.0,
+ encoder_layer_drop: float = 0.0,
+ mask_prob: float = 0.8,
+ mask_channel_prob: float = 0.0,
+ mask_channel_length: int = 10,
+ feature_grad_mult: Optional[float] = None,
+) -> HuBERTPretrainModel:
+ """Builds "large" :class:`HuBERTPretrainModel` from *HuBERT* :cite:`hsu2021hubert` for pretraining.
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`hubert_pretrain_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`hubert_pretrain_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`hubert_pretrain_model`.
+ encoder_dropout (float):
+ See :py:func:`hubert_pretrain_model`.
+ encoder_layer_drop (float):
+ See :py:func:`hubert_pretrain_model`.
+ mask_prob (float):
+ See :py:func:`hubert_pretrain_model`.
+ mask_channel_prob (float):
+ See :py:func:`hubert_pretrain_model`.
+ mask_channel_length (int):
+ See :py:func:`hubert_pretrain_model`.
+ feature_grad_mult (float or None):
+ See :py:func:`hubert_pretrain_model`.
+
+ Returns:
+ HuBERTPretrainModel:
+ The resulting model.
+ """ # noqa: E501
+ return hubert_pretrain_model(
+ extractor_mode="layer_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=False,
+ encoder_embed_dim=1024,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=24,
+ encoder_num_heads=16,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=4096,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=True,
+ encoder_layer_drop=encoder_layer_drop,
+ mask_prob=mask_prob,
+ mask_selection="static",
+ mask_other=0.0,
+ mask_length=10,
+ no_mask_overlap=False,
+ mask_min_space=1,
+ mask_channel_prob=mask_channel_prob,
+ mask_channel_selection="static",
+ mask_channel_other=0.0,
+ mask_channel_length=mask_channel_length,
+ no_mask_channel_overlap=False,
+ mask_channel_min_space=1,
+ skip_masked=False,
+ skip_nomask=False,
+ num_classes=500,
+ final_dim=768,
+ feature_grad_mult=feature_grad_mult,
+ )
+
+
+def hubert_pretrain_xlarge(
+ encoder_projection_dropout: float = 0.0,
+ encoder_attention_dropout: float = 0.0,
+ encoder_ff_interm_dropout: float = 0.0,
+ encoder_dropout: float = 0.0,
+ encoder_layer_drop: float = 0.0,
+ mask_prob: float = 0.8,
+ mask_channel_prob: float = 0.0,
+ mask_channel_length: int = 10,
+ feature_grad_mult: Optional[float] = None,
+) -> HuBERTPretrainModel:
+ """Builds "extra large" :class:`HuBERTPretrainModel` from *HuBERT* :cite:`hsu2021hubert` for pretraining.
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`hubert_pretrain_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`hubert_pretrain_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`hubert_pretrain_model`.
+ encoder_dropout (float):
+ See :py:func:`hubert_pretrain_model`.
+ encoder_layer_drop (float):
+ See :py:func:`hubert_pretrain_model`.
+ mask_prob (float):
+ See :py:func:`hubert_pretrain_model`.
+ mask_channel_prob (float):
+ See :py:func:`hubert_pretrain_model`.
+ mask_channel_length (int):
+ See :py:func:`hubert_pretrain_model`.
+ feature_grad_mult (float or None):
+ See :py:func:`hubert_pretrain_model`.
+
+ Returns:
+ HuBERTPretrainModel:
+ The resulting model.
+ """ # noqa: E501
+ return hubert_pretrain_model(
+ extractor_mode="layer_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=False,
+ encoder_embed_dim=1280,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=48,
+ encoder_num_heads=16,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=5120,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=True,
+ encoder_layer_drop=encoder_layer_drop,
+ mask_prob=mask_prob,
+ mask_selection="static",
+ mask_other=0.0,
+ mask_length=10,
+ no_mask_overlap=False,
+ mask_min_space=1,
+ mask_channel_prob=mask_channel_prob,
+ mask_channel_selection="static",
+ mask_channel_other=0.0,
+ mask_channel_length=mask_channel_length,
+ no_mask_channel_overlap=False,
+ mask_channel_min_space=1,
+ skip_masked=False,
+ skip_nomask=False,
+ num_classes=500,
+ final_dim=1024,
+ feature_grad_mult=feature_grad_mult,
+ )
+
+
+def wavlm_model(
+ extractor_mode: str,
+ extractor_conv_layer_config: Optional[List[Tuple[int, int, int]]],
+ extractor_conv_bias: bool,
+ encoder_embed_dim: int,
+ encoder_projection_dropout: float,
+ encoder_pos_conv_kernel: int,
+ encoder_pos_conv_groups: int,
+ encoder_num_layers: int,
+ encoder_num_heads: int,
+ encoder_num_buckets: int,
+ encoder_max_distance: int,
+ encoder_attention_dropout: float,
+ encoder_ff_interm_features: int,
+ encoder_ff_interm_dropout: float,
+ encoder_dropout: float,
+ encoder_layer_norm_first: bool,
+ encoder_layer_drop: float,
+ aux_num_out: Optional[int],
+) -> Wav2Vec2Model:
+ """Builds custom WaveLM model :cite:`chen2022wavlm`. The architecture is compatible
+ with Wav2Vec2 model :cite:`baevski2020wav2vec`, and so the output object is
+ :class:`~torchaudio.models.Wav2Vec2Model`. Most of the arguments have the same meaning
+ as in :py:func:`~torchaudio.models.wav2vec2_model` so please refer there for documentation.
+
+ Args:
+ extractor_mode (str): Operation mode of feature extractor.
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ extractor_conv_layer_config (list of integer tuples or None):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ extractor_conv_bias (bool):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ encoder_embed_dim (int):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ encoder_projection_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ encoder_pos_conv_kernel (int):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ encoder_pos_conv_groups (int):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ encoder_num_layers (int):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ encoder_num_heads (int):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ encoder_num_buckets (int):
+ Number of buckets for relative position embedding.
+ encoder_max_distance (int):
+ Maximum distance for relative position embedding.
+
+ encoder_attention_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ encoder_ff_interm_features (int):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ encoder_ff_interm_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ encoder_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ encoder_layer_norm_first (bool):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ encoder_layer_drop (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ aux_num_out (int or None):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """
+ if extractor_conv_layer_config is None:
+ extractor_conv_layer_config = [(512, 10, 5)] + [(512, 3, 2)] * 4 + [(512, 2, 2)] * 2
+
+ feature_extractor = components._get_feature_extractor(
+ extractor_mode, extractor_conv_layer_config, extractor_conv_bias
+ )
+ encoder = components._get_wavlm_encoder(
+ in_features=extractor_conv_layer_config[-1][0],
+ embed_dim=encoder_embed_dim,
+ dropout_input=encoder_projection_dropout,
+ pos_conv_kernel=encoder_pos_conv_kernel,
+ pos_conv_groups=encoder_pos_conv_groups,
+ num_layers=encoder_num_layers,
+ num_heads=encoder_num_heads,
+ num_buckets=encoder_num_buckets,
+ max_distance=encoder_max_distance,
+ attention_dropout=encoder_attention_dropout,
+ ff_interm_features=encoder_ff_interm_features,
+ ff_interm_dropout=encoder_ff_interm_dropout,
+ dropout=encoder_dropout,
+ layer_norm_first=encoder_layer_norm_first,
+ layer_drop=encoder_layer_drop,
+ )
+ aux = None
+ if aux_num_out is not None:
+ aux = torch.nn.Linear(in_features=encoder_embed_dim, out_features=aux_num_out)
+ return Wav2Vec2Model(feature_extractor, encoder, aux)
+
+
+def wavlm_base(
+ encoder_projection_dropout: float = 0.1,
+ encoder_attention_dropout: float = 0.1,
+ encoder_ff_interm_dropout: float = 0.1,
+ encoder_dropout: float = 0.1,
+ encoder_layer_drop: float = 0.1,
+ aux_num_out: Optional[int] = None,
+) -> Wav2Vec2Model:
+ """Builds "base" WaveLM model :cite:`chen2022wavlm`. The architecture is compatible
+ with Wav2Vec2 model :cite:`baevski2020wav2vec`, and so the output class is
+ :class:`~torchaudio.models.Wav2Vec2Model`.
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_layer_drop (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ aux_num_out (int, optional):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """
+ return wavlm_model(
+ extractor_mode="group_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=False,
+ encoder_embed_dim=768,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=12,
+ encoder_num_heads=12,
+ encoder_num_buckets=320,
+ encoder_max_distance=800,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=3072,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=False,
+ encoder_layer_drop=encoder_layer_drop,
+ aux_num_out=aux_num_out,
+ )
+
+
+def wavlm_large(
+ encoder_projection_dropout: float = 0.1,
+ encoder_attention_dropout: float = 0.1,
+ encoder_ff_interm_dropout: float = 0.0,
+ encoder_dropout: float = 0.1,
+ encoder_layer_drop: float = 0.1,
+ aux_num_out: Optional[int] = None,
+) -> Wav2Vec2Model:
+ """Builds "large" WaveLM model :cite:`chen2022wavlm`. The architecture is compatible
+ with Wav2Vec2 model :cite:`baevski2020wav2vec`, and so the output class is
+ :class:`~torchaudio.models.Wav2Vec2Model`.
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_layer_drop (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ aux_num_out (int, optional):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """
+ return wavlm_model(
+ extractor_mode="layer_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=False,
+ encoder_embed_dim=1024,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=24,
+ encoder_num_heads=16,
+ encoder_num_buckets=320,
+ encoder_max_distance=800,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=4096,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=True,
+ encoder_layer_drop=encoder_layer_drop,
+ aux_num_out=aux_num_out,
+ )
+
+
+def wav2vec2_xlsr_300m(
+ encoder_projection_dropout: float = 0.0,
+ encoder_attention_dropout: float = 0.0,
+ encoder_ff_interm_dropout: float = 0.0,
+ encoder_dropout: float = 0.0,
+ encoder_layer_drop: float = 0.0,
+ aux_num_out: Optional[int] = None,
+) -> Wav2Vec2Model:
+ """Builds XLS-R model :cite:`babu2021xls` with 300 millions of parameters. The architecture is compatible
+ with Wav2Vec2 model :cite:`baevski2020wav2vec`, and so the output class is
+ :class:`~torchaudio.models.Wav2Vec2Model`.
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_layer_drop (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ aux_num_out (int, optional):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """
+ return wav2vec2_model(
+ extractor_mode="layer_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=True,
+ encoder_embed_dim=1024,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=24,
+ encoder_num_heads=16,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=4096,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=True,
+ encoder_layer_drop=encoder_layer_drop,
+ aux_num_out=aux_num_out,
+ )
+
+
+def wav2vec2_xlsr_1b(
+ encoder_projection_dropout: float = 0.1,
+ encoder_attention_dropout: float = 0.0,
+ encoder_ff_interm_dropout: float = 0.0,
+ encoder_dropout: float = 0.0,
+ encoder_layer_drop: float = 0.0,
+ aux_num_out: Optional[int] = None,
+) -> Wav2Vec2Model:
+ """Builds XLS-R model :cite:`babu2021xls` with 1 billion of parameters. The architecture is compatible
+ with Wav2Vec2 model :cite:`baevski2020wav2vec`, and so the output class is
+ :class:`~torchaudio.models.Wav2Vec2Model`.
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_layer_drop (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ aux_num_out (int, optional):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """
+ return wav2vec2_model(
+ extractor_mode="layer_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=True,
+ encoder_embed_dim=1280,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=48,
+ encoder_num_heads=16,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=5120,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=True,
+ encoder_layer_drop=encoder_layer_drop,
+ aux_num_out=aux_num_out,
+ )
+
+
+def wav2vec2_xlsr_2b(
+ encoder_projection_dropout: float = 0.1,
+ encoder_attention_dropout: float = 0.0,
+ encoder_ff_interm_dropout: float = 0.0,
+ encoder_dropout: float = 0.0,
+ encoder_layer_drop: float = 0.0,
+ aux_num_out: Optional[int] = None,
+) -> Wav2Vec2Model:
+ """Builds XLS-R model :cite:`babu2021xls` with 2 billions of parameters. The architecture is compatible
+ with Wav2Vec2 model :cite:`baevski2020wav2vec`, and so the output class is
+ :class:`~torchaudio.models.Wav2Vec2Model`.
+
+ Args:
+ encoder_projection_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_attention_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_ff_interm_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_dropout (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ encoder_layer_drop (float):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+ aux_num_out (int, optional):
+ See :py:func:`~torchaudio.models.wav2vec2_model`.
+
+ Returns:
+ Wav2Vec2Model:
+ The resulting model.
+ """
+ return wav2vec2_model(
+ extractor_mode="layer_norm",
+ extractor_conv_layer_config=None,
+ extractor_conv_bias=True,
+ encoder_embed_dim=1920,
+ encoder_projection_dropout=encoder_projection_dropout,
+ encoder_pos_conv_kernel=128,
+ encoder_pos_conv_groups=16,
+ encoder_num_layers=48,
+ encoder_num_heads=16,
+ encoder_attention_dropout=encoder_attention_dropout,
+ encoder_ff_interm_features=7680,
+ encoder_ff_interm_dropout=encoder_ff_interm_dropout,
+ encoder_dropout=encoder_dropout,
+ encoder_layer_norm_first=True,
+ encoder_layer_drop=encoder_layer_drop,
+ aux_num_out=aux_num_out,
+ )
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/__init__.py b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0457b5dd707f7216adc3ea919ba8e257d86f4f71
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/__init__.py
@@ -0,0 +1,7 @@
+from .import_fairseq import import_fairseq_model
+from .import_huggingface import import_huggingface_model
+
+__all__ = [
+ "import_huggingface_model",
+ "import_fairseq_model",
+]
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..52395a0bddbd67e33c62f7fae55f88e7348f058d
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/__pycache__/import_fairseq.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/__pycache__/import_fairseq.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..01cd8ee7dfedc92949a715b23a7a372df7ccdf51
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/__pycache__/import_fairseq.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/__pycache__/import_huggingface.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/__pycache__/import_huggingface.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f7664b8b35fbe6b1407cf2c689a8aff513e8b397
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/__pycache__/import_huggingface.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/import_fairseq.py b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/import_fairseq.py
new file mode 100644
index 0000000000000000000000000000000000000000..39791e9b7d75ac3c2eb1fcf4f9c3517e7483048c
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/import_fairseq.py
@@ -0,0 +1,213 @@
+"""Import fariseq's wav2vec2.0 pretrained weights to torchaudios's format.
+
+For this module to work, you need `fairseq`.
+"""
+import re
+
+from torch.nn import Module
+
+from ..model import wav2vec2_model, Wav2Vec2Model
+
+
+def _parse_config(w2v_model):
+ encoder = w2v_model.encoder
+ conv_layers = w2v_model.feature_extractor.conv_layers
+
+ extractor_mode = "layer_norm"
+ if "GroupNorm" in conv_layers[0][2].__class__.__name__:
+ extractor_mode = "group_norm"
+ else:
+ extractor_mode = "layer_norm"
+
+ conv_layer_config = [(l[0].out_channels, l[0].kernel_size[0], l[0].stride[0]) for l in conv_layers]
+
+ if all(l[0].bias is None for l in conv_layers):
+ conv_bias = False
+ elif all(l[0].bias is not None for l in conv_layers):
+ conv_bias = True
+ else:
+ raise ValueError("Either all the convolutions layers have bias term or none of them should.")
+
+ config = {
+ "extractor_mode": extractor_mode,
+ "extractor_conv_layer_config": conv_layer_config,
+ "extractor_conv_bias": conv_bias,
+ "encoder_embed_dim": w2v_model.post_extract_proj.out_features,
+ "encoder_projection_dropout": w2v_model.dropout_input.p,
+ "encoder_pos_conv_kernel": encoder.pos_conv[0].kernel_size[0],
+ "encoder_pos_conv_groups": encoder.pos_conv[0].groups,
+ "encoder_num_layers": len(encoder.layers),
+ "encoder_num_heads": encoder.layers[0].self_attn.num_heads,
+ "encoder_attention_dropout": encoder.layers[0].self_attn.dropout_module.p,
+ "encoder_ff_interm_features": encoder.layers[0].fc1.out_features,
+ "encoder_ff_interm_dropout": encoder.layers[0].dropout2.p,
+ "encoder_dropout": encoder.layers[0].dropout3.p,
+ "encoder_layer_norm_first": encoder.layer_norm_first,
+ "encoder_layer_drop": encoder.layerdrop,
+ }
+ return config
+
+
+def _map_key(key):
+ key_ = key
+ if key.startswith("w2v_model."):
+ key = key.replace("w2v_model.", "")
+ if re.match(r"(mask_emb|quantizer|project_q|final_proj|mask_emb)", key):
+ return None
+ # Feature Extractor
+ # Group norm when "extractor_mode" is "default".
+ # (Only the first layer)
+ # "conv_layers.0.2.weight" -> "conv_layers.0.layer_norm.weight"
+ # "conv_layers.0.2.bias" -> "conv_layers.0.layer_norm.bias"
+ match = re.match(r"feature_extractor\.conv_layers\.0\.2\.(weight|bias)", key)
+ if match:
+ return f"feature_extractor.conv_layers.0.layer_norm.{match.group(1)}"
+ # Convolutions
+ # "conv_layers.X.0.weight" -> "conv_layers.X.conv.weight"
+ # "conv_layers.X.0.bias" -> "conv_layers.X.conv.bias"
+ match = re.match(r"feature_extractor\.conv_layers\.(\d+)\.0\.(weight|bias)", key)
+ if match:
+ return f"feature_extractor.conv_layers.{match.group(1)}.conv.{match.group(2)}"
+ # Layer norm when "extractor_mode" is "layer_norm".
+ # "conv_layers.X.2.1.weight" -> "conv_layers.X.layer_norm.weight"
+ # "conv_layers.X.2.1.bias" -> "conv_layers.X.layer_norm.bias"
+ match = re.match(r"feature_extractor\.conv_layers\.(\d+)\.2\.1\.(weight|bias)", key)
+ if match:
+ return f"feature_extractor.conv_layers.{match.group(1)}.layer_norm.{match.group(2)}"
+ match = re.match(r"post_extract_proj\.(weight|bias)", key)
+ # Encoder - Feature projection
+ if match:
+ return f"encoder.feature_projection.projection.{match.group(1)}"
+ match = re.match(r"layer_norm\.(weight|bias)", key)
+ if match:
+ return f"encoder.feature_projection.layer_norm.{match.group(1)}"
+ # Encoder - Transformer - Convolutional positional embedding
+ match = re.match(r"encoder\.pos_conv\.0\.(bias|weight_g|weight_v)", key)
+ if match:
+ return f"encoder.transformer.pos_conv_embed.conv.{match.group(1)}"
+ match = re.match(r"encoder\.layer_norm\.(weight|bias)", key)
+ if match:
+ return f"encoder.transformer.layer_norm.{match.group(1)}"
+ # Encoder - Transformer - Self attention layers
+ match = re.match(r"encoder\.layers\.(\d+)\.self_attn\.((k_|v_|q_|out_)proj\.(weight|bias))", key)
+ if match:
+ return f"encoder.transformer.layers.{match.group(1)}.attention.{match.group(2)}"
+ match = re.match(r"encoder\.layers\.(\d+)\.self_attn_layer_norm\.(weight|bias)", key)
+ if match:
+ return f"encoder.transformer.layers.{match.group(1)}.layer_norm.{match.group(2)}"
+ match = re.match(r"encoder\.layers\.(\d+)\.fc1\.(weight|bias)", key)
+ if match:
+ return f"encoder.transformer.layers.{match.group(1)}.feed_forward.intermediate_dense.{match.group(2)}"
+ match = re.match(r"encoder\.layers\.(\d+)\.fc2\.(weight|bias)", key)
+ if match:
+ return f"encoder.transformer.layers.{match.group(1)}.feed_forward.output_dense.{match.group(2)}"
+ match = re.match(r"encoder\.layers\.(\d+)\.final_layer_norm\.(weight|bias)", key)
+ if match:
+ return f"encoder.transformer.layers.{match.group(1)}.final_layer_norm.{match.group(2)}"
+ match = re.match(r"proj\.(weight|bias)", key)
+ # Auxiliary Module
+ # Only relevant when loading fine-tuned models
+ if match:
+ return f"aux.{match.group(1)}"
+ # HuBERT Extension
+ if key in ["label_embs_concat"]:
+ return key
+ raise ValueError(f"Unexpected key: {key_}")
+
+
+def _convert_state_dict(state_dict):
+ converted = {}
+ for k, v in state_dict.items():
+ k = _map_key(k)
+ if k is not None:
+ converted[k] = v
+ return converted
+
+
+def import_fairseq_model(original: Module) -> Wav2Vec2Model:
+ """Builds :class:`Wav2Vec2Model` from the corresponding model object of
+ `fairseq `_.
+
+ Args:
+ original (torch.nn.Module):
+ An instance of fairseq's Wav2Vec2.0 or HuBERT model.
+ One of ``fairseq.models.wav2vec.wav2vec2_asr.Wav2VecEncoder``,
+ ``fairseq.models.wav2vec.wav2vec2.Wav2Vec2Model`` or
+ ``fairseq.models.hubert.hubert_asr.HubertEncoder``.
+
+ Returns:
+ Wav2Vec2Model: Imported model.
+
+ Example - Loading pretrain-only model
+ >>> from torchaudio.models.wav2vec2.utils import import_fairseq_model
+ >>>
+ >>> # Load model using fairseq
+ >>> model_file = 'wav2vec_small.pt'
+ >>> model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task([model_file])
+ >>> original = model[0]
+ >>> imported = import_fairseq_model(original)
+ >>>
+ >>> # Perform feature extraction
+ >>> waveform, _ = torchaudio.load('audio.wav')
+ >>> features, _ = imported.extract_features(waveform)
+ >>>
+ >>> # Compare result with the original model from fairseq
+ >>> reference = original.feature_extractor(waveform).transpose(1, 2)
+ >>> torch.testing.assert_allclose(features, reference)
+
+ Example - Fine-tuned model
+ >>> from torchaudio.models.wav2vec2.utils import import_fairseq_model
+ >>>
+ >>> # Load model using fairseq
+ >>> model_file = 'wav2vec_small_960h.pt'
+ >>> model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task([model_file])
+ >>> original = model[0]
+ >>> imported = import_fairseq_model(original.w2v_encoder)
+ >>>
+ >>> # Perform encoding
+ >>> waveform, _ = torchaudio.load('audio.wav')
+ >>> emission, _ = imported(waveform)
+ >>>
+ >>> # Compare result with the original model from fairseq
+ >>> mask = torch.zeros_like(waveform)
+ >>> reference = original(waveform, mask)['encoder_out'].transpose(0, 1)
+ >>> torch.testing.assert_allclose(emission, reference)
+ """
+ class_ = original.__class__.__name__
+ if class_ == "Wav2Vec2Model":
+ return _import_wav2vec2_pretraining(original)
+ if class_ == "Wav2VecEncoder":
+ return _import_wav2vec2_finetuning(original)
+ if class_ == "HubertModel":
+ return _import_hubert_pretraining(original)
+ if class_ == "HubertEncoder":
+ return _import_hubert_finetuning(original)
+ raise ValueError(f"Expected an instance of `Wav2Vec2Model` or `Wav2VecEncoder`. Found: {class_}")
+
+
+def _import_wav2vec2_finetuning(original: Module) -> Wav2Vec2Model:
+ config = _parse_config(original.w2v_model)
+ model = wav2vec2_model(**config, aux_num_out=original.proj.out_features)
+ model.load_state_dict(_convert_state_dict(original.state_dict()))
+ return model
+
+
+def _import_wav2vec2_pretraining(original: Module) -> Wav2Vec2Model:
+ config = _parse_config(original)
+ model = wav2vec2_model(**config, aux_num_out=None)
+ model.load_state_dict(_convert_state_dict(original.state_dict()), strict=False)
+ return model
+
+
+def _import_hubert_finetuning(original: Module) -> Wav2Vec2Model:
+ config = _parse_config(original.w2v_model)
+ model = wav2vec2_model(**config, aux_num_out=original.proj.out_features)
+ model.load_state_dict(_convert_state_dict(original.state_dict()), strict=False)
+ return model
+
+
+def _import_hubert_pretraining(original: Module) -> Wav2Vec2Model:
+ config = _parse_config(original)
+ model = wav2vec2_model(**config, aux_num_out=None)
+ model.load_state_dict(_convert_state_dict(original.state_dict()), strict=False)
+ return model
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/import_huggingface.py b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/import_huggingface.py
new file mode 100644
index 0000000000000000000000000000000000000000..519d8c919f02be62b2f2e2aa0dd8db97222430d5
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/utils/import_huggingface.py
@@ -0,0 +1,134 @@
+"""Import Hugging Face transformers's wav2vec2.0 pretrained weights to torchaudios's format.
+"""
+import logging
+from typing import Any, Dict
+
+import torch
+from torch.nn import Module
+
+from ..model import wav2vec2_model, Wav2Vec2Model, wavlm_model
+
+_LG = logging.getLogger(__name__)
+
+
+def _get_config(cfg):
+ config = {
+ "extractor_mode": f"{cfg.feat_extract_norm}_norm",
+ "extractor_conv_layer_config": list(zip(cfg.conv_dim, cfg.conv_kernel, cfg.conv_stride)),
+ "extractor_conv_bias": cfg.conv_bias,
+ "encoder_embed_dim": cfg.hidden_size,
+ "encoder_projection_dropout": cfg.feat_proj_dropout,
+ "encoder_pos_conv_kernel": cfg.num_conv_pos_embeddings,
+ "encoder_pos_conv_groups": cfg.num_conv_pos_embedding_groups,
+ "encoder_num_layers": cfg.num_hidden_layers,
+ "encoder_num_heads": cfg.num_attention_heads,
+ "encoder_attention_dropout": cfg.attention_dropout,
+ "encoder_ff_interm_features": cfg.intermediate_size,
+ "encoder_ff_interm_dropout": cfg.activation_dropout,
+ "encoder_dropout": cfg.hidden_dropout,
+ "encoder_layer_norm_first": cfg.do_stable_layer_norm,
+ "encoder_layer_drop": cfg.layerdrop,
+ }
+ return config
+
+
+def _get_config_wavlm(cfg):
+ config = {
+ "extractor_mode": f"{cfg.feat_extract_norm}_norm",
+ "extractor_conv_layer_config": list(zip(cfg.conv_dim, cfg.conv_kernel, cfg.conv_stride)),
+ "extractor_conv_bias": cfg.conv_bias,
+ "encoder_embed_dim": cfg.hidden_size,
+ "encoder_projection_dropout": cfg.feat_proj_dropout,
+ "encoder_pos_conv_kernel": cfg.num_conv_pos_embeddings,
+ "encoder_pos_conv_groups": cfg.num_conv_pos_embedding_groups,
+ "encoder_num_layers": cfg.num_hidden_layers,
+ "encoder_num_heads": cfg.num_attention_heads,
+ "encoder_num_buckets": cfg.num_buckets,
+ "encoder_max_distance": cfg.max_bucket_distance,
+ "encoder_attention_dropout": cfg.attention_dropout,
+ "encoder_ff_interm_features": cfg.intermediate_size,
+ "encoder_ff_interm_dropout": cfg.activation_dropout,
+ "encoder_dropout": cfg.hidden_dropout,
+ "encoder_layer_norm_first": cfg.do_stable_layer_norm,
+ "encoder_layer_drop": cfg.layerdrop,
+ }
+ return config
+
+
+def _build(config, original):
+ is_for_ctc = original.__class__.__name__ in ["Wav2Vec2ForCTC", "WavLMForCTC"]
+ if is_for_ctc:
+ aux_num_out = original.config.vocab_size
+ wav2vec2 = original.wav2vec2
+ else:
+ _LG.warning(
+ "The model is not an instance of Wav2Vec2ForCTC or WavLMForCTC. " '"lm_head" module is not imported.'
+ )
+ aux_num_out = None
+ wav2vec2 = original
+ is_wavlm = original.__class__.__name__ in ["WavLMModel", "WavLMForCTC"]
+ if is_wavlm:
+ imported = wavlm_model(**config, aux_num_out=aux_num_out)
+ else:
+ imported = wav2vec2_model(**config, aux_num_out=aux_num_out)
+ imported.feature_extractor.load_state_dict(wav2vec2.feature_extractor.state_dict())
+ imported.encoder.feature_projection.load_state_dict(wav2vec2.feature_projection.state_dict())
+ encoder_state_dict = wav2vec2.encoder.state_dict()
+ if is_wavlm: # Rename paramaters of linear transformations for compatibility with the HF model
+ transform_wavlm_encoder_state(encoder_state_dict, config["encoder_num_layers"])
+ imported.encoder.transformer.load_state_dict(encoder_state_dict)
+ if is_for_ctc:
+ imported.aux.load_state_dict(original.lm_head.state_dict())
+ return imported
+
+
+def transform_wavlm_encoder_state(state: Dict[str, Any], encoder_num_layers: int):
+ """Converts WavLM encoder state from HuggingFace format. In particular, concatenates linear projection weights and
+ biases to align with the structure of ``torch.nn.MultiheadAttention``.
+ """
+ for i in range(encoder_num_layers):
+ q_proj_bias = state.pop(f"layers.{i}.attention.q_proj.bias")
+ k_proj_bias = state.pop(f"layers.{i}.attention.k_proj.bias")
+ v_proj_bias = state.pop(f"layers.{i}.attention.v_proj.bias")
+ q_proj_weight = state.pop(f"layers.{i}.attention.q_proj.weight")
+ k_proj_weight = state.pop(f"layers.{i}.attention.k_proj.weight")
+ v_proj_weight = state.pop(f"layers.{i}.attention.v_proj.weight")
+ state[f"layers.{i}.attention.attention.in_proj_bias"] = torch.cat((q_proj_bias, k_proj_bias, v_proj_bias))
+ state[f"layers.{i}.attention.attention.in_proj_weight"] = torch.cat(
+ (q_proj_weight, k_proj_weight, v_proj_weight)
+ )
+
+ state[f"layers.{i}.attention.attention.out_proj.weight"] = state.pop(f"layers.{i}.attention.out_proj.weight")
+ state[f"layers.{i}.attention.attention.out_proj.bias"] = state.pop(f"layers.{i}.attention.out_proj.bias")
+
+
+def import_huggingface_model(original: Module) -> Wav2Vec2Model:
+ """Builds :class:`Wav2Vec2Model` from the corresponding model object of
+ `Transformers `_.
+
+ Args:
+ original (torch.nn.Module): An instance of ``Wav2Vec2ForCTC`` from ``transformers``.
+
+ Returns:
+ Wav2Vec2Model: Imported model.
+
+ Example
+ >>> from torchaudio.models.wav2vec2.utils import import_huggingface_model
+ >>>
+ >>> original = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
+ >>> model = import_huggingface_model(original)
+ >>>
+ >>> waveforms, _ = torchaudio.load("audio.wav")
+ >>> logits, _ = model(waveforms)
+ """
+ _LG.info("Importing model.")
+ _LG.info("Loading model configuration.")
+ is_wavlm = original.__class__.__name__ in ["WavLMModel", "WavLMForCTC"]
+ if is_wavlm:
+ config = _get_config_wavlm(original.config)
+ else:
+ config = _get_config(original.config)
+ _LG.debug(" - config: %s", config)
+ _LG.info("Building model.")
+ imported = _build(config, original)
+ return imported
diff --git a/lib/python3.12/site-packages/torchaudio/models/wav2vec2/wavlm_attention.py b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/wavlm_attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..fafddfeb958cbcdfdc0a7781b49bc124fff78290
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/wav2vec2/wavlm_attention.py
@@ -0,0 +1,214 @@
+"""
+The MIT License (MIT)
+
+Copyright (c) Microsoft Corporation
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+
+import math
+from typing import Optional, Tuple
+
+import torch
+from torch import nn, Tensor
+
+
+class WavLMSelfAttention(nn.Module):
+ """Multi-headed self-attention for WavLM model :cite:`chen2022wavlm`.
+ Wraps around ``torch.nn.MultiheadAttention``, creating relaive position embeddings and passing them to multi-headed
+ attention as a mask.
+ Source: https://github.com/microsoft/unilm/blob/2d8302f09c99bca2b82e6e868d81d4281cceebc8/wavlm/modules.py#L303-L763
+
+ Args:
+ embed_dim (int): Total dimension of the model.
+ num_heads (int): The number of heads.
+ dropout (float, optional): Dropout probability on attn_output_weights. (Default: to ``0.0``)
+ bias (bool, optional): If ``True``, add bias to input / output projection layers. (Default: ``True``)
+ has_relative_attention_bias (bool, optional): If ``True``, apply relative position embedding.
+ Necessary in the first encoder layer, but not in the subsequent ones. (Default: ``False``)
+ num_buckets (int, optional): Number of buckets for relative position embedding. (Default: ``32``)
+ max_distance (int, optional): Naximum distance for relative position embedding. (Default: ``128``)
+ gru_rel_pos (bool, optional): If ``True``, apply gated relative position embedding. (Default: ``False``)
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ dropout: float = 0.0,
+ bias: bool = True,
+ has_relative_attention_bias: bool = False,
+ num_buckets: int = 32,
+ max_distance: int = 128,
+ gru_rel_pos: bool = True,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.has_relative_attention_bias = has_relative_attention_bias
+ self.num_buckets = num_buckets
+ self.max_distance = max_distance
+
+ if has_relative_attention_bias:
+ self.rel_attn_embed = nn.Embedding(num_buckets, num_heads)
+ else:
+ self.rel_attn_embed = None
+
+ self.head_dim = embed_dim // num_heads
+ assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+
+ self.dropout = dropout
+ self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, bias=bias, batch_first=True)
+
+ self.gru_rel_pos = gru_rel_pos
+ if self.gru_rel_pos:
+ self.gru_rel_pos_linear = nn.Linear(self.head_dim, 8)
+ self.gru_rel_pos_const = nn.Parameter(torch.ones(1, num_heads, 1, 1))
+ self.has_position_bias = True
+
+ def compute_bias(self, query_length: int, key_length: int) -> Tensor:
+ """Compute relative position embeddings for WavLM model.
+ Args:
+ query_length (int): Query position can take values between 0 and ``query_length - 1``.
+ key_length (int): Key position can take values between 0 and ``key_length - 1``.
+ Returns:
+ Tensor of shape `(num_heads, query_length, key_length)`, relative positions embeddings
+ """
+ context_position = torch.arange(query_length, dtype=torch.long)[:, None]
+ memory_position = torch.arange(key_length, dtype=torch.long)[None, :]
+ relative_position = memory_position - context_position # Shape (query_length, key_length)
+ relative_position_bucket = self._relative_positions_bucket(relative_position, bidirectional=True)
+ relative_position_bucket = relative_position_bucket.to(self.rel_attn_embed.weight.device)
+ values = self.rel_attn_embed(relative_position_bucket) # Shape (query_length, key_length, num_heads)
+ values = values.permute([2, 0, 1])
+ return values
+
+ def _relative_positions_bucket(self, relative_positions: Tensor, bidirectional: bool = True):
+ """Compute relative position buckets for WavLM model. Computation similar to formula (5) in WavLM
+ paper :cite:`chen2022wavlm`.
+ Args:
+ relative_positions (Tensor): Relative offsets between query and key positions,
+ of shape ``(query_length, key_length)``.
+ bidirectional (bool): If ``True``, values will be filled both above and below the diagonal in the resulting
+ matrix. If ``False``, the elements above the diagonal (i.e. with negative relative offsets) will be set
+ to zero. (Default ``True``)
+ Returns:
+ Tensor of shape ``(query_length, key_length)`` filled bucketed values of with relative positions.
+ """
+ num_buckets = self.num_buckets
+ max_distance = self.max_distance
+ # Shape (query_length, key_length)
+ relative_buckets = torch.zeros_like(relative_positions, dtype=torch.long)
+
+ if bidirectional:
+ num_buckets = num_buckets // 2
+ relative_buckets += (relative_positions > 0).to(torch.long) * num_buckets
+ relative_positions = torch.abs(relative_positions)
+ else:
+ relative_positions = -torch.min(relative_positions, torch.zeros_like(relative_positions))
+
+ max_exact = num_buckets // 2
+ is_small = relative_positions < max_exact
+
+ relative_postion_if_large = max_exact + (
+ torch.log(relative_positions.float() / max_exact)
+ / math.log(max_distance / max_exact)
+ * (num_buckets - max_exact)
+ ).to(torch.long)
+ relative_postion_if_large = torch.min(
+ relative_postion_if_large, torch.full_like(relative_postion_if_large, num_buckets - 1)
+ )
+
+ relative_buckets += torch.where(is_small, relative_positions, relative_postion_if_large)
+ return relative_buckets
+
+ def forward(
+ self,
+ query: Tensor,
+ key_padding_mask: Optional[Tensor] = None,
+ attention_mask: Optional[Tensor] = None,
+ position_bias: Optional[Tensor] = None,
+ ) -> Tuple[Tensor, Optional[Tensor]]:
+ """
+ Args:
+ query (Tensor): Input of shape ``(batch_size, src_len, embed_dim)``.
+ key_padding_mask (Tensor or None, optional): Mask to exclude keys that are pads, of shape
+ `(batch, src_len)`, where padding elements are indicated by 1s. (Default: ``None``)
+ attn_mask: Needs to be ``None``. The argument exists for compatibility with
+ ``EncoderLayer``. (Default: ``None``)
+ position_bias (Tensor or None, optional): Position bias of shape
+ ``(batch_size * num_heads, src_len, src_len)``. When used inside WavLM model encoder, will be
+ generated in the first layer and then passed from each encoder layer to the next one.
+ (Default: ``None``)
+ Returns:
+ attn_output (Tensor): Attention output of shape ``(batch_size, src_len, embed_dim)``.
+ position_bias (Tensor or None): Position bias of shape ``(batch_size * num_heads, src_len, src_len)``.
+ """
+ bsz, seq_len, embed_dim = query.size()
+ assert embed_dim == self.embed_dim
+ assert attention_mask is None
+
+ if self.rel_attn_embed is not None and position_bias is None:
+ position_bias = self.compute_bias(seq_len, seq_len)
+ position_bias = position_bias.unsqueeze(0).repeat(bsz, 1, 1, 1)
+
+ attn_mask_rel_pos: Optional[Tensor] = None
+ if position_bias is not None:
+ attn_mask_rel_pos = position_bias
+ if self.gru_rel_pos: # Apply gating on relative position bias
+ query_layer = query.view(bsz, seq_len, self.num_heads, -1)
+ query_layer = query_layer.permute(0, 2, 1, 3)
+
+ gate_a, gate_b = torch.sigmoid(
+ self.gru_rel_pos_linear(query_layer).view(bsz, self.num_heads, seq_len, 2, 4).sum(-1, keepdim=False)
+ ).chunk(2, dim=-1)
+ gate_a_1 = gate_a * (gate_b * self.gru_rel_pos_const - 1.0) + 2.0
+ attn_mask_rel_pos = gate_a_1.view(bsz, self.num_heads, -1, 1) * position_bias
+
+ attn_mask_rel_pos = attn_mask_rel_pos.view((bsz, self.num_heads, seq_len, seq_len))
+
+ if attn_mask_rel_pos is not None and key_padding_mask is not None:
+ key_padding_mask = key_padding_mask.view(bsz, 1, 1, seq_len).expand(-1, self.num_heads, -1, -1)
+ key_padding_mask = torch.nn.functional._canonical_mask(
+ mask=key_padding_mask,
+ mask_name="key_padding_mask",
+ other_type=torch.nn.functional._none_or_dtype(attn_mask_rel_pos),
+ other_name="",
+ target_type=query.dtype,
+ )
+ if attn_mask_rel_pos is not None and key_padding_mask is not None:
+ attn_mask_rel_pos = attn_mask_rel_pos + key_padding_mask
+ query_projected = torch.nn.functional.linear(query, self.attention.in_proj_weight, self.attention.in_proj_bias)
+ query, key, value = query_projected.chunk(3, -1)
+ shape = (bsz, seq_len, self.num_heads, self.head_dim)
+ query = query.view(shape).transpose(2, 1) # (batch, num_heads, seq_len, head_dim)
+ key = key.view(shape).transpose(2, 1) # (batch, num_heads, seq_len, head_dim)
+ value = value.view(shape).transpose(2, 1) # (batch, num_heads, seq_len, head_dim)
+ dropout = self.dropout if self.training else 0.0
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
+ query,
+ key,
+ value,
+ attn_mask=attn_mask_rel_pos,
+ dropout_p=dropout,
+ is_causal=False,
+ )
+ attn_output = attn_output.transpose(1, 2).reshape(bsz, -1, self.num_heads * self.head_dim)
+ attn_output = self.attention.out_proj(attn_output)
+ return attn_output, position_bias
diff --git a/lib/python3.12/site-packages/torchaudio/models/wavernn.py b/lib/python3.12/site-packages/torchaudio/models/wavernn.py
new file mode 100644
index 0000000000000000000000000000000000000000..8ae5a3e91675cd9ef7d4614f0daaec50f80dcdee
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/models/wavernn.py
@@ -0,0 +1,409 @@
+import math
+from typing import List, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+
+__all__ = [
+ "ResBlock",
+ "MelResNet",
+ "Stretch2d",
+ "UpsampleNetwork",
+ "WaveRNN",
+]
+
+
+class ResBlock(nn.Module):
+ r"""ResNet block based on *Efficient Neural Audio Synthesis* :cite:`kalchbrenner2018efficient`.
+
+ Args:
+ n_freq: the number of bins in a spectrogram. (Default: ``128``)
+
+ Examples
+ >>> resblock = ResBlock()
+ >>> input = torch.rand(10, 128, 512) # a random spectrogram
+ >>> output = resblock(input) # shape: (10, 128, 512)
+ """
+
+ def __init__(self, n_freq: int = 128) -> None:
+ super().__init__()
+
+ self.resblock_model = nn.Sequential(
+ nn.Conv1d(in_channels=n_freq, out_channels=n_freq, kernel_size=1, bias=False),
+ nn.BatchNorm1d(n_freq),
+ nn.ReLU(inplace=True),
+ nn.Conv1d(in_channels=n_freq, out_channels=n_freq, kernel_size=1, bias=False),
+ nn.BatchNorm1d(n_freq),
+ )
+
+ def forward(self, specgram: Tensor) -> Tensor:
+ r"""Pass the input through the ResBlock layer.
+ Args:
+ specgram (Tensor): the input sequence to the ResBlock layer (n_batch, n_freq, n_time).
+
+ Return:
+ Tensor shape: (n_batch, n_freq, n_time)
+ """
+
+ return self.resblock_model(specgram) + specgram
+
+
+class MelResNet(nn.Module):
+ r"""MelResNet layer uses a stack of ResBlocks on spectrogram.
+
+ Args:
+ n_res_block: the number of ResBlock in stack. (Default: ``10``)
+ n_freq: the number of bins in a spectrogram. (Default: ``128``)
+ n_hidden: the number of hidden dimensions of resblock. (Default: ``128``)
+ n_output: the number of output dimensions of melresnet. (Default: ``128``)
+ kernel_size: the number of kernel size in the first Conv1d layer. (Default: ``5``)
+
+ Examples
+ >>> melresnet = MelResNet()
+ >>> input = torch.rand(10, 128, 512) # a random spectrogram
+ >>> output = melresnet(input) # shape: (10, 128, 508)
+ """
+
+ def __init__(
+ self, n_res_block: int = 10, n_freq: int = 128, n_hidden: int = 128, n_output: int = 128, kernel_size: int = 5
+ ) -> None:
+ super().__init__()
+
+ ResBlocks = [ResBlock(n_hidden) for _ in range(n_res_block)]
+
+ self.melresnet_model = nn.Sequential(
+ nn.Conv1d(in_channels=n_freq, out_channels=n_hidden, kernel_size=kernel_size, bias=False),
+ nn.BatchNorm1d(n_hidden),
+ nn.ReLU(inplace=True),
+ *ResBlocks,
+ nn.Conv1d(in_channels=n_hidden, out_channels=n_output, kernel_size=1),
+ )
+
+ def forward(self, specgram: Tensor) -> Tensor:
+ r"""Pass the input through the MelResNet layer.
+ Args:
+ specgram (Tensor): the input sequence to the MelResNet layer (n_batch, n_freq, n_time).
+
+ Return:
+ Tensor shape: (n_batch, n_output, n_time - kernel_size + 1)
+ """
+
+ return self.melresnet_model(specgram)
+
+
+class Stretch2d(nn.Module):
+ r"""Upscale the frequency and time dimensions of a spectrogram.
+
+ Args:
+ time_scale: the scale factor in time dimension
+ freq_scale: the scale factor in frequency dimension
+
+ Examples
+ >>> stretch2d = Stretch2d(time_scale=10, freq_scale=5)
+
+ >>> input = torch.rand(10, 100, 512) # a random spectrogram
+ >>> output = stretch2d(input) # shape: (10, 500, 5120)
+ """
+
+ def __init__(self, time_scale: int, freq_scale: int) -> None:
+ super().__init__()
+
+ self.freq_scale = freq_scale
+ self.time_scale = time_scale
+
+ def forward(self, specgram: Tensor) -> Tensor:
+ r"""Pass the input through the Stretch2d layer.
+
+ Args:
+ specgram (Tensor): the input sequence to the Stretch2d layer (..., n_freq, n_time).
+
+ Return:
+ Tensor shape: (..., n_freq * freq_scale, n_time * time_scale)
+ """
+
+ return specgram.repeat_interleave(self.freq_scale, -2).repeat_interleave(self.time_scale, -1)
+
+
+class UpsampleNetwork(nn.Module):
+ r"""Upscale the dimensions of a spectrogram.
+
+ Args:
+ upsample_scales: the list of upsample scales.
+ n_res_block: the number of ResBlock in stack. (Default: ``10``)
+ n_freq: the number of bins in a spectrogram. (Default: ``128``)
+ n_hidden: the number of hidden dimensions of resblock. (Default: ``128``)
+ n_output: the number of output dimensions of melresnet. (Default: ``128``)
+ kernel_size: the number of kernel size in the first Conv1d layer. (Default: ``5``)
+
+ Examples
+ >>> upsamplenetwork = UpsampleNetwork(upsample_scales=[4, 4, 16])
+ >>> input = torch.rand(10, 128, 10) # a random spectrogram
+ >>> output = upsamplenetwork(input) # shape: (10, 128, 1536), (10, 128, 1536)
+ """
+
+ def __init__(
+ self,
+ upsample_scales: List[int],
+ n_res_block: int = 10,
+ n_freq: int = 128,
+ n_hidden: int = 128,
+ n_output: int = 128,
+ kernel_size: int = 5,
+ ) -> None:
+ super().__init__()
+
+ total_scale = 1
+ for upsample_scale in upsample_scales:
+ total_scale *= upsample_scale
+ self.total_scale: int = total_scale
+
+ self.indent = (kernel_size - 1) // 2 * total_scale
+ self.resnet = MelResNet(n_res_block, n_freq, n_hidden, n_output, kernel_size)
+ self.resnet_stretch = Stretch2d(total_scale, 1)
+
+ up_layers = []
+ for scale in upsample_scales:
+ stretch = Stretch2d(scale, 1)
+ conv = nn.Conv2d(
+ in_channels=1, out_channels=1, kernel_size=(1, scale * 2 + 1), padding=(0, scale), bias=False
+ )
+ torch.nn.init.constant_(conv.weight, 1.0 / (scale * 2 + 1))
+ up_layers.append(stretch)
+ up_layers.append(conv)
+ self.upsample_layers = nn.Sequential(*up_layers)
+
+ def forward(self, specgram: Tensor) -> Tuple[Tensor, Tensor]:
+ r"""Pass the input through the UpsampleNetwork layer.
+
+ Args:
+ specgram (Tensor): the input sequence to the UpsampleNetwork layer (n_batch, n_freq, n_time)
+
+ Return:
+ Tensor shape: (n_batch, n_freq, (n_time - kernel_size + 1) * total_scale),
+ (n_batch, n_output, (n_time - kernel_size + 1) * total_scale)
+ where total_scale is the product of all elements in upsample_scales.
+ """
+
+ resnet_output = self.resnet(specgram).unsqueeze(1)
+ resnet_output = self.resnet_stretch(resnet_output)
+ resnet_output = resnet_output.squeeze(1)
+
+ specgram = specgram.unsqueeze(1)
+ upsampling_output = self.upsample_layers(specgram)
+ upsampling_output = upsampling_output.squeeze(1)[:, :, self.indent : -self.indent]
+
+ return upsampling_output, resnet_output
+
+
+class WaveRNN(nn.Module):
+ r"""WaveRNN model from *Efficient Neural Audio Synthesis* :cite:`wavernn`
+ based on the implementation from `fatchord/WaveRNN `_.
+
+ The original implementation was introduced in *Efficient Neural Audio Synthesis*
+ :cite:`kalchbrenner2018efficient`. The input channels of waveform and spectrogram have to be 1.
+ The product of `upsample_scales` must equal `hop_length`.
+
+ See Also:
+ * `Training example `__
+ * :class:`torchaudio.pipelines.Tacotron2TTSBundle`: TTS pipeline with pretrained model.
+
+ Args:
+ upsample_scales: the list of upsample scales.
+ n_classes: the number of output classes.
+ hop_length: the number of samples between the starts of consecutive frames.
+ n_res_block: the number of ResBlock in stack. (Default: ``10``)
+ n_rnn: the dimension of RNN layer. (Default: ``512``)
+ n_fc: the dimension of fully connected layer. (Default: ``512``)
+ kernel_size: the number of kernel size in the first Conv1d layer. (Default: ``5``)
+ n_freq: the number of bins in a spectrogram. (Default: ``128``)
+ n_hidden: the number of hidden dimensions of resblock. (Default: ``128``)
+ n_output: the number of output dimensions of melresnet. (Default: ``128``)
+
+ Example
+ >>> wavernn = WaveRNN(upsample_scales=[5,5,8], n_classes=512, hop_length=200)
+ >>> waveform, sample_rate = torchaudio.load(file)
+ >>> # waveform shape: (n_batch, n_channel, (n_time - kernel_size + 1) * hop_length)
+ >>> specgram = MelSpectrogram(sample_rate)(waveform) # shape: (n_batch, n_channel, n_freq, n_time)
+ >>> output = wavernn(waveform, specgram)
+ >>> # output shape: (n_batch, n_channel, (n_time - kernel_size + 1) * hop_length, n_classes)
+ """
+
+ def __init__(
+ self,
+ upsample_scales: List[int],
+ n_classes: int,
+ hop_length: int,
+ n_res_block: int = 10,
+ n_rnn: int = 512,
+ n_fc: int = 512,
+ kernel_size: int = 5,
+ n_freq: int = 128,
+ n_hidden: int = 128,
+ n_output: int = 128,
+ ) -> None:
+ super().__init__()
+
+ self.kernel_size = kernel_size
+ self._pad = (kernel_size - 1 if kernel_size % 2 else kernel_size) // 2
+ self.n_rnn = n_rnn
+ self.n_aux = n_output // 4
+ self.hop_length = hop_length
+ self.n_classes = n_classes
+ self.n_bits: int = int(math.log2(self.n_classes))
+
+ total_scale = 1
+ for upsample_scale in upsample_scales:
+ total_scale *= upsample_scale
+ if total_scale != self.hop_length:
+ raise ValueError(f"Expected: total_scale == hop_length, but found {total_scale} != {hop_length}")
+
+ self.upsample = UpsampleNetwork(upsample_scales, n_res_block, n_freq, n_hidden, n_output, kernel_size)
+ self.fc = nn.Linear(n_freq + self.n_aux + 1, n_rnn)
+
+ self.rnn1 = nn.GRU(n_rnn, n_rnn, batch_first=True)
+ self.rnn2 = nn.GRU(n_rnn + self.n_aux, n_rnn, batch_first=True)
+
+ self.relu1 = nn.ReLU(inplace=True)
+ self.relu2 = nn.ReLU(inplace=True)
+
+ self.fc1 = nn.Linear(n_rnn + self.n_aux, n_fc)
+ self.fc2 = nn.Linear(n_fc + self.n_aux, n_fc)
+ self.fc3 = nn.Linear(n_fc, self.n_classes)
+
+ def forward(self, waveform: Tensor, specgram: Tensor) -> Tensor:
+ r"""Pass the input through the WaveRNN model.
+
+ Args:
+ waveform: the input waveform to the WaveRNN layer (n_batch, 1, (n_time - kernel_size + 1) * hop_length)
+ specgram: the input spectrogram to the WaveRNN layer (n_batch, 1, n_freq, n_time)
+
+ Return:
+ Tensor: shape (n_batch, 1, (n_time - kernel_size + 1) * hop_length, n_classes)
+ """
+
+ if waveform.size(1) != 1:
+ raise ValueError("Require the input channel of waveform is 1")
+ if specgram.size(1) != 1:
+ raise ValueError("Require the input channel of specgram is 1")
+ # remove channel dimension until the end
+ waveform, specgram = waveform.squeeze(1), specgram.squeeze(1)
+
+ batch_size = waveform.size(0)
+ h1 = torch.zeros(1, batch_size, self.n_rnn, dtype=waveform.dtype, device=waveform.device)
+ h2 = torch.zeros(1, batch_size, self.n_rnn, dtype=waveform.dtype, device=waveform.device)
+ # output of upsample:
+ # specgram: (n_batch, n_freq, (n_time - kernel_size + 1) * total_scale)
+ # aux: (n_batch, n_output, (n_time - kernel_size + 1) * total_scale)
+ specgram, aux = self.upsample(specgram)
+ specgram = specgram.transpose(1, 2)
+ aux = aux.transpose(1, 2)
+
+ aux_idx = [self.n_aux * i for i in range(5)]
+ a1 = aux[:, :, aux_idx[0] : aux_idx[1]]
+ a2 = aux[:, :, aux_idx[1] : aux_idx[2]]
+ a3 = aux[:, :, aux_idx[2] : aux_idx[3]]
+ a4 = aux[:, :, aux_idx[3] : aux_idx[4]]
+
+ x = torch.cat([waveform.unsqueeze(-1), specgram, a1], dim=-1)
+ x = self.fc(x)
+ res = x
+ x, _ = self.rnn1(x, h1)
+
+ x = x + res
+ res = x
+ x = torch.cat([x, a2], dim=-1)
+ x, _ = self.rnn2(x, h2)
+
+ x = x + res
+ x = torch.cat([x, a3], dim=-1)
+ x = self.fc1(x)
+ x = self.relu1(x)
+
+ x = torch.cat([x, a4], dim=-1)
+ x = self.fc2(x)
+ x = self.relu2(x)
+ x = self.fc3(x)
+
+ # bring back channel dimension
+ return x.unsqueeze(1)
+
+ @torch.jit.export
+ def infer(self, specgram: Tensor, lengths: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]:
+ r"""Inference method of WaveRNN.
+
+ This function currently only supports multinomial sampling, which assumes the
+ network is trained on cross entropy loss.
+
+ Args:
+ specgram (Tensor):
+ Batch of spectrograms. Shape: `(n_batch, n_freq, n_time)`.
+ lengths (Tensor or None, optional):
+ Indicates the valid length of each audio in the batch.
+ Shape: `(batch, )`.
+ When the ``specgram`` contains spectrograms with different durations,
+ by providing ``lengths`` argument, the model will compute
+ the corresponding valid output lengths.
+ If ``None``, it is assumed that all the audio in ``waveforms``
+ have valid length. Default: ``None``.
+
+ Returns:
+ (Tensor, Optional[Tensor]):
+ Tensor
+ The inferred waveform of size `(n_batch, 1, n_time)`.
+ 1 stands for a single channel.
+ Tensor or None
+ If ``lengths`` argument was provided, a Tensor of shape `(batch, )`
+ is returned.
+ It indicates the valid length in time axis of the output Tensor.
+ """
+
+ device = specgram.device
+ dtype = specgram.dtype
+
+ specgram = torch.nn.functional.pad(specgram, (self._pad, self._pad))
+ specgram, aux = self.upsample(specgram)
+ if lengths is not None:
+ lengths = lengths * self.upsample.total_scale
+
+ output: List[Tensor] = []
+ b_size, _, seq_len = specgram.size()
+
+ h1 = torch.zeros((1, b_size, self.n_rnn), device=device, dtype=dtype)
+ h2 = torch.zeros((1, b_size, self.n_rnn), device=device, dtype=dtype)
+ x = torch.zeros((b_size, 1), device=device, dtype=dtype)
+
+ aux_split = [aux[:, self.n_aux * i : self.n_aux * (i + 1), :] for i in range(4)]
+
+ for i in range(seq_len):
+
+ m_t = specgram[:, :, i]
+
+ a1_t, a2_t, a3_t, a4_t = [a[:, :, i] for a in aux_split]
+
+ x = torch.cat([x, m_t, a1_t], dim=1)
+ x = self.fc(x)
+ _, h1 = self.rnn1(x.unsqueeze(1), h1)
+
+ x = x + h1[0]
+ inp = torch.cat([x, a2_t], dim=1)
+ _, h2 = self.rnn2(inp.unsqueeze(1), h2)
+
+ x = x + h2[0]
+ x = torch.cat([x, a3_t], dim=1)
+ x = F.relu(self.fc1(x))
+
+ x = torch.cat([x, a4_t], dim=1)
+ x = F.relu(self.fc2(x))
+
+ logits = self.fc3(x)
+
+ posterior = F.softmax(logits, dim=1)
+
+ x = torch.multinomial(posterior, 1).float()
+ # Transform label [0, 2 ** n_bits - 1] to waveform [-1, 1]
+ x = 2 * x / (2**self.n_bits - 1.0) - 1.0
+
+ output.append(x)
+
+ return torch.stack(output).permute(1, 2, 0), lengths
diff --git a/lib/python3.12/site-packages/torchaudio/transforms/__init__.py b/lib/python3.12/site-packages/torchaudio/transforms/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..1fe77865a9740ab6c88465dbe4a5a41bc445a688
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/transforms/__init__.py
@@ -0,0 +1,75 @@
+from ._multi_channel import MVDR, PSD, RTFMVDR, SoudenMVDR
+from ._transforms import (
+ AddNoise,
+ AmplitudeToDB,
+ ComputeDeltas,
+ Convolve,
+ Deemphasis,
+ Fade,
+ FFTConvolve,
+ FrequencyMasking,
+ GriffinLim,
+ InverseMelScale,
+ InverseSpectrogram,
+ LFCC,
+ Loudness,
+ MelScale,
+ MelSpectrogram,
+ MFCC,
+ MuLawDecoding,
+ MuLawEncoding,
+ PitchShift,
+ Preemphasis,
+ Resample,
+ RNNTLoss,
+ SlidingWindowCmn,
+ SpecAugment,
+ SpectralCentroid,
+ Spectrogram,
+ Speed,
+ SpeedPerturbation,
+ TimeMasking,
+ TimeStretch,
+ Vad,
+ Vol,
+)
+
+
+__all__ = [
+ "AddNoise",
+ "AmplitudeToDB",
+ "ComputeDeltas",
+ "Convolve",
+ "Deemphasis",
+ "Fade",
+ "FFTConvolve",
+ "FrequencyMasking",
+ "GriffinLim",
+ "InverseMelScale",
+ "InverseSpectrogram",
+ "LFCC",
+ "Loudness",
+ "MFCC",
+ "MVDR",
+ "MelScale",
+ "MelSpectrogram",
+ "MuLawDecoding",
+ "MuLawEncoding",
+ "PSD",
+ "PitchShift",
+ "Preemphasis",
+ "RNNTLoss",
+ "RTFMVDR",
+ "Resample",
+ "SlidingWindowCmn",
+ "SoudenMVDR",
+ "SpecAugment",
+ "SpectralCentroid",
+ "Spectrogram",
+ "Speed",
+ "SpeedPerturbation",
+ "TimeMasking",
+ "TimeStretch",
+ "Vad",
+ "Vol",
+]
diff --git a/lib/python3.12/site-packages/torchaudio/transforms/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/transforms/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2801db2e134c12303a29052b07b8307ba751f4cb
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/transforms/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/transforms/__pycache__/_multi_channel.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/transforms/__pycache__/_multi_channel.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fafb75e66cccff4b3fc16852bc61919954619d8f
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/transforms/__pycache__/_multi_channel.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/transforms/_multi_channel.py b/lib/python3.12/site-packages/torchaudio/transforms/_multi_channel.py
new file mode 100644
index 0000000000000000000000000000000000000000..4ba3db7f454058de7c0fda1d57781ed346d7a65c
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/transforms/_multi_channel.py
@@ -0,0 +1,467 @@
+# -*- coding: utf-8 -*-
+
+import warnings
+from typing import Optional, Union
+
+import torch
+from torch import Tensor
+from torchaudio import functional as F
+
+
+__all__ = []
+
+
+def _get_mvdr_vector(
+ psd_s: torch.Tensor,
+ psd_n: torch.Tensor,
+ reference_vector: torch.Tensor,
+ solution: str = "ref_channel",
+ diagonal_loading: bool = True,
+ diag_eps: float = 1e-7,
+ eps: float = 1e-8,
+) -> torch.Tensor:
+ r"""Compute the MVDR beamforming weights with ``solution`` argument.
+
+ Args:
+ psd_s (torch.Tensor): The complex-valued power spectral density (PSD) matrix of target speech.
+ Tensor with dimensions `(..., freq, channel, channel)`.
+ psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise.
+ Tensor with dimensions `(..., freq, channel, channel)`.
+ reference_vector (torch.Tensor): one-hot reference channel matrix.
+ solution (str, optional): Solution to compute the MVDR beamforming weights.
+ Options: [``ref_channel``, ``stv_evd``, ``stv_power``]. (Default: ``ref_channel``)
+ diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to ``psd_n``.
+ (Default: ``True``)
+ diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading.
+ It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``)
+ eps (float, optional): Value to add to the denominator in the beamforming weight formula.
+ (Default: ``1e-8``)
+
+ Returns:
+ torch.Tensor: the mvdr beamforming weight matrix
+ """
+ if solution == "ref_channel":
+ beamform_vector = F.mvdr_weights_souden(psd_s, psd_n, reference_vector, diagonal_loading, diag_eps, eps)
+ else:
+ if solution == "stv_evd":
+ stv = F.rtf_evd(psd_s)
+ else:
+ stv = F.rtf_power(psd_s, psd_n, reference_vector, diagonal_loading=diagonal_loading, diag_eps=diag_eps)
+ beamform_vector = F.mvdr_weights_rtf(stv, psd_n, reference_vector, diagonal_loading, diag_eps, eps)
+
+ return beamform_vector
+
+
+class PSD(torch.nn.Module):
+ r"""Compute cross-channel power spectral density (PSD) matrix.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ multi_mask (bool, optional): If ``True``, only accepts multi-channel Time-Frequency masks. (Default: ``False``)
+ normalize (bool, optional): If ``True``, normalize the mask along the time dimension. (Default: ``True``)
+ eps (float, optional): Value to add to the denominator in mask normalization. (Default: ``1e-15``)
+ """
+
+ def __init__(self, multi_mask: bool = False, normalize: bool = True, eps: float = 1e-15):
+ super().__init__()
+ self.multi_mask = multi_mask
+ self.normalize = normalize
+ self.eps = eps
+
+ def forward(self, specgram: torch.Tensor, mask: Optional[torch.Tensor] = None):
+ """
+ Args:
+ specgram (torch.Tensor): Multi-channel complex-valued spectrum.
+ Tensor with dimensions `(..., channel, freq, time)`.
+ mask (torch.Tensor or None, optional): Time-Frequency mask for normalization.
+ Tensor with dimensions `(..., freq, time)` if multi_mask is ``False`` or
+ with dimensions `(..., channel, freq, time)` if multi_mask is ``True``.
+ (Default: ``None``)
+
+ Returns:
+ torch.Tensor: The complex-valued PSD matrix of the input spectrum.
+ Tensor with dimensions `(..., freq, channel, channel)`
+ """
+ if mask is not None:
+ if self.multi_mask:
+ # Averaging mask along channel dimension
+ mask = mask.mean(dim=-3) # (..., freq, time)
+ psd = F.psd(specgram, mask, self.normalize, self.eps)
+
+ return psd
+
+
+class MVDR(torch.nn.Module):
+ """Minimum Variance Distortionless Response (MVDR) module that performs MVDR beamforming with Time-Frequency masks.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Based on https://github.com/espnet/espnet/blob/master/espnet2/enh/layers/beamformer.py
+
+ We provide three solutions of MVDR beamforming. One is based on *reference channel selection*
+ :cite:`souden2009optimal` (``solution=ref_channel``).
+
+ .. math::
+ \\textbf{w}_{\\text{MVDR}}(f) =\
+ \\frac{{{\\bf{\\Phi}_{\\textbf{NN}}^{-1}}(f){\\bf{\\Phi}_{\\textbf{SS}}}}(f)}\
+ {\\text{Trace}({{{\\bf{\\Phi}_{\\textbf{NN}}^{-1}}(f) \\bf{\\Phi}_{\\textbf{SS}}}(f))}}\\bm{u}
+
+ where :math:`\\bf{\\Phi}_{\\textbf{SS}}` and :math:`\\bf{\\Phi}_{\\textbf{NN}}` are the covariance\
+ matrices of speech and noise, respectively. :math:`\\bf{u}` is an one-hot vector to determine the\
+ reference channel.
+
+ The other two solutions are based on the steering vector (``solution=stv_evd`` or ``solution=stv_power``).
+
+ .. math::
+ \\textbf{w}_{\\text{MVDR}}(f) =\
+ \\frac{{{\\bf{\\Phi}_{\\textbf{NN}}^{-1}}(f){\\bm{v}}(f)}}\
+ {{\\bm{v}^{\\mathsf{H}}}(f){\\bf{\\Phi}_{\\textbf{NN}}^{-1}}(f){\\bm{v}}(f)}
+
+ where :math:`\\bm{v}` is the acoustic transfer function or the steering vector.\
+ :math:`.^{\\mathsf{H}}` denotes the Hermitian Conjugate operation.
+
+ We apply either *eigenvalue decomposition*
+ :cite:`higuchi2016robust` or the *power method* :cite:`mises1929praktische` to get the
+ steering vector from the PSD matrix of speech.
+
+ After estimating the beamforming weight, the enhanced Short-time Fourier Transform (STFT) is obtained by
+
+ .. math::
+ \\hat{\\bf{S}} = {\\bf{w}^\\mathsf{H}}{\\bf{Y}}, {\\bf{w}} \\in \\mathbb{C}^{M \\times F}
+
+ where :math:`\\bf{Y}` and :math:`\\hat{\\bf{S}}` are the STFT of the multi-channel noisy speech and\
+ the single-channel enhanced speech, respectively.
+
+ For online streaming audio, we provide a *recursive method* :cite:`higuchi2017online` to update the
+ PSD matrices of speech and noise, respectively.
+
+ Args:
+ ref_channel (int, optional): Reference channel for beamforming. (Default: ``0``)
+ solution (str, optional): Solution to compute the MVDR beamforming weights.
+ Options: [``ref_channel``, ``stv_evd``, ``stv_power``]. (Default: ``ref_channel``)
+ multi_mask (bool, optional): If ``True``, only accepts multi-channel Time-Frequency masks. (Default: ``False``)
+ diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to the covariance matrix
+ of the noise. (Default: ``True``)
+ diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading.
+ It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``)
+ online (bool, optional): If ``True``, updates the MVDR beamforming weights based on
+ the previous covarience matrices. (Default: ``False``)
+
+ Note:
+ To improve the numerical stability, the input spectrogram will be converted to double precision
+ (``torch.complex128`` or ``torch.cdouble``) dtype for internal computation. The output spectrogram
+ is converted to the dtype of the input spectrogram to be compatible with other modules.
+
+ Note:
+ If you use ``stv_evd`` solution, the gradient of the same input may not be identical if the
+ eigenvalues of the PSD matrix are not distinct (i.e. some eigenvalues are close or identical).
+ """
+
+ def __init__(
+ self,
+ ref_channel: int = 0,
+ solution: str = "ref_channel",
+ multi_mask: bool = False,
+ diag_loading: bool = True,
+ diag_eps: float = 1e-7,
+ online: bool = False,
+ ):
+ super().__init__()
+ if solution not in [
+ "ref_channel",
+ "stv_evd",
+ "stv_power",
+ ]:
+ raise ValueError(
+ '`solution` must be one of ["ref_channel", "stv_evd", "stv_power"]. Given {}'.format(solution)
+ )
+ self.ref_channel = ref_channel
+ self.solution = solution
+ self.multi_mask = multi_mask
+ self.diag_loading = diag_loading
+ self.diag_eps = diag_eps
+ self.online = online
+ self.psd = PSD(multi_mask)
+
+ psd_s: torch.Tensor = torch.zeros(1)
+ psd_n: torch.Tensor = torch.zeros(1)
+ mask_sum_s: torch.Tensor = torch.zeros(1)
+ mask_sum_n: torch.Tensor = torch.zeros(1)
+ self.register_buffer("psd_s", psd_s)
+ self.register_buffer("psd_n", psd_n)
+ self.register_buffer("mask_sum_s", mask_sum_s)
+ self.register_buffer("mask_sum_n", mask_sum_n)
+
+ def _get_updated_mvdr_vector(
+ self,
+ psd_s: torch.Tensor,
+ psd_n: torch.Tensor,
+ mask_s: torch.Tensor,
+ mask_n: torch.Tensor,
+ reference_vector: torch.Tensor,
+ solution: str = "ref_channel",
+ diagonal_loading: bool = True,
+ diag_eps: float = 1e-7,
+ eps: float = 1e-8,
+ ) -> torch.Tensor:
+ r"""Recursively update the MVDR beamforming vector.
+
+ Args:
+ psd_s (torch.Tensor): The complex-valued power spectral density (PSD) matrix of target speech.
+ Tensor with dimensions `(..., freq, channel, channel)`.
+ psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise.
+ Tensor with dimensions `(..., freq, channel, channel)`.
+ mask_s (torch.Tensor): Time-Frequency mask of the target speech.
+ Tensor with dimensions `(..., freq, time)` if multi_mask is ``False``
+ or with dimensions `(..., channel, freq, time)` if multi_mask is ``True``.
+ mask_n (torch.Tensor or None, optional): Time-Frequency mask of the noise.
+ Tensor with dimensions `(..., freq, time)` if multi_mask is ``False``
+ or with dimensions `(..., channel, freq, time)` if multi_mask is ``True``.
+ reference_vector (torch.Tensor): One-hot reference channel matrix.
+ solution (str, optional): Solution to compute the MVDR beamforming weights.
+ Options: [``ref_channel``, ``stv_evd``, ``stv_power``]. (Default: ``ref_channel``)
+ diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to ``psd_n``.
+ (Default: ``True``)
+ diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading.
+ It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``)
+ eps (float, optional): Value to add to the denominator in the beamforming weight formula.
+ (Default: ``1e-8``)
+
+ Returns:
+ torch.Tensor: The MVDR beamforming weight matrix.
+ """
+ if self.multi_mask:
+ # Averaging mask along channel dimension
+ mask_s = mask_s.mean(dim=-3) # (..., freq, time)
+ mask_n = mask_n.mean(dim=-3) # (..., freq, time)
+ if self.psd_s.ndim == 1:
+ self.psd_s = psd_s
+ self.psd_n = psd_n
+ self.mask_sum_s = mask_s.sum(dim=-1)
+ self.mask_sum_n = mask_n.sum(dim=-1)
+ return _get_mvdr_vector(psd_s, psd_n, reference_vector, solution, diagonal_loading, diag_eps, eps)
+ else:
+ psd_s = self._get_updated_psd_speech(psd_s, mask_s)
+ psd_n = self._get_updated_psd_noise(psd_n, mask_n)
+ self.psd_s = psd_s
+ self.psd_n = psd_n
+ self.mask_sum_s = self.mask_sum_s + mask_s.sum(dim=-1)
+ self.mask_sum_n = self.mask_sum_n + mask_n.sum(dim=-1)
+ return _get_mvdr_vector(psd_s, psd_n, reference_vector, solution, diagonal_loading, diag_eps, eps)
+
+ def _get_updated_psd_speech(self, psd_s: torch.Tensor, mask_s: torch.Tensor) -> torch.Tensor:
+ r"""Update psd of speech recursively.
+
+ Args:
+ psd_s (torch.Tensor): The complex-valued power spectral density (PSD) matrix of target speech.
+ Tensor with dimensions `(..., freq, channel, channel)`.
+ mask_s (torch.Tensor): Time-Frequency mask of the target speech.
+ Tensor with dimensions `(..., freq, time)`.
+
+ Returns:
+ torch.Tensor: The updated PSD matrix of target speech.
+ """
+ numerator = self.mask_sum_s / (self.mask_sum_s + mask_s.sum(dim=-1))
+ denominator = 1 / (self.mask_sum_s + mask_s.sum(dim=-1))
+ psd_s = self.psd_s * numerator[..., None, None] + psd_s * denominator[..., None, None]
+ return psd_s
+
+ def _get_updated_psd_noise(self, psd_n: torch.Tensor, mask_n: torch.Tensor) -> torch.Tensor:
+ r"""Update psd of noise recursively.
+
+ Args:
+ psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise.
+ Tensor with dimensions `(..., freq, channel, channel)`.
+ mask_n (torch.Tensor or None, optional): Time-Frequency mask of the noise.
+ Tensor with dimensions `(..., freq, time)`.
+
+ Returns:
+ torch.Tensor: The updated PSD matrix of noise.
+ """
+ numerator = self.mask_sum_n / (self.mask_sum_n + mask_n.sum(dim=-1))
+ denominator = 1 / (self.mask_sum_n + mask_n.sum(dim=-1))
+ psd_n = self.psd_n * numerator[..., None, None] + psd_n * denominator[..., None, None]
+ return psd_n
+
+ def forward(
+ self, specgram: torch.Tensor, mask_s: torch.Tensor, mask_n: Optional[torch.Tensor] = None
+ ) -> torch.Tensor:
+ """Perform MVDR beamforming.
+
+ Args:
+ specgram (torch.Tensor): Multi-channel complex-valued spectrum.
+ Tensor with dimensions `(..., channel, freq, time)`
+ mask_s (torch.Tensor): Time-Frequency mask of target speech.
+ Tensor with dimensions `(..., freq, time)` if multi_mask is ``False``
+ or with dimensions `(..., channel, freq, time)` if multi_mask is ``True``.
+ mask_n (torch.Tensor or None, optional): Time-Frequency mask of noise.
+ Tensor with dimensions `(..., freq, time)` if multi_mask is ``False``
+ or with dimensions `(..., channel, freq, time)` if multi_mask is ``True``.
+ (Default: None)
+
+ Returns:
+ torch.Tensor: Single-channel complex-valued enhanced spectrum with dimensions `(..., freq, time)`.
+ """
+ dtype = specgram.dtype
+ if specgram.ndim < 3:
+ raise ValueError(f"Expected at least 3D tensor (..., channel, freq, time). Found: {specgram.shape}")
+ if not specgram.is_complex():
+ raise ValueError(
+ f"The type of ``specgram`` tensor must be ``torch.cfloat`` or ``torch.cdouble``.\
+ Found: {specgram.dtype}"
+ )
+ if specgram.dtype == torch.cfloat:
+ specgram = specgram.cdouble() # Convert specgram to ``torch.cdouble``.
+
+ if mask_n is None:
+ warnings.warn("``mask_n`` is not provided, use ``1 - mask_s`` as ``mask_n``.")
+ mask_n = 1 - mask_s
+
+ psd_s = self.psd(specgram, mask_s) # (..., freq, time, channel, channel)
+ psd_n = self.psd(specgram, mask_n) # (..., freq, time, channel, channel)
+
+ u = torch.zeros(specgram.size()[:-2], device=specgram.device, dtype=torch.cdouble) # (..., channel)
+ u[..., self.ref_channel].fill_(1)
+
+ if self.online:
+ w_mvdr = self._get_updated_mvdr_vector(
+ psd_s, psd_n, mask_s, mask_n, u, self.solution, self.diag_loading, self.diag_eps
+ )
+ else:
+ w_mvdr = _get_mvdr_vector(psd_s, psd_n, u, self.solution, self.diag_loading, self.diag_eps)
+
+ specgram_enhanced = F.apply_beamforming(w_mvdr, specgram)
+
+ return specgram_enhanced.to(dtype)
+
+
+class RTFMVDR(torch.nn.Module):
+ r"""Minimum Variance Distortionless Response (*MVDR* :cite:`capon1969high`) module
+ based on the relative transfer function (RTF) and power spectral density (PSD) matrix of noise.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Given the multi-channel complex-valued spectrum :math:`\textbf{Y}`, the relative transfer function (RTF) matrix
+ or the steering vector of target speech :math:`\bm{v}`, the PSD matrix of noise :math:`\bf{\Phi}_{\textbf{NN}}`, and
+ a one-hot vector that represents the reference channel :math:`\bf{u}`, the module computes the single-channel
+ complex-valued spectrum of the enhanced speech :math:`\hat{\textbf{S}}`. The formula is defined as:
+
+ .. math::
+ \hat{\textbf{S}}(f) = \textbf{w}_{\text{bf}}(f)^{\mathsf{H}} \textbf{Y}(f)
+
+ where :math:`\textbf{w}_{\text{bf}}(f)` is the MVDR beamforming weight for the :math:`f`-th frequency bin,
+ :math:`(.)^{\mathsf{H}}` denotes the Hermitian Conjugate operation.
+
+ The beamforming weight is computed by:
+
+ .. math::
+ \textbf{w}_{\text{MVDR}}(f) =
+ \frac{{{\bf{\Phi}_{\textbf{NN}}^{-1}}(f){\bm{v}}(f)}}
+ {{\bm{v}^{\mathsf{H}}}(f){\bf{\Phi}_{\textbf{NN}}^{-1}}(f){\bm{v}}(f)}
+ """
+
+ def forward(
+ self,
+ specgram: Tensor,
+ rtf: Tensor,
+ psd_n: Tensor,
+ reference_channel: Union[int, Tensor],
+ diagonal_loading: bool = True,
+ diag_eps: float = 1e-7,
+ eps: float = 1e-8,
+ ) -> Tensor:
+ """
+ Args:
+ specgram (torch.Tensor): Multi-channel complex-valued spectrum.
+ Tensor with dimensions `(..., channel, freq, time)`
+ rtf (torch.Tensor): The complex-valued RTF vector of target speech.
+ Tensor with dimensions `(..., freq, channel)`.
+ psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise.
+ Tensor with dimensions `(..., freq, channel, channel)`.
+ reference_channel (int or torch.Tensor): Specifies the reference channel.
+ If the dtype is ``int``, it represents the reference channel index.
+ If the dtype is ``torch.Tensor``, its shape is `(..., channel)`, where the ``channel`` dimension
+ is one-hot.
+ diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to ``psd_n``.
+ (Default: ``True``)
+ diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading.
+ It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``)
+ eps (float, optional): Value to add to the denominator in the beamforming weight formula.
+ (Default: ``1e-8``)
+
+ Returns:
+ torch.Tensor: Single-channel complex-valued enhanced spectrum with dimensions `(..., freq, time)`.
+ """
+ w_mvdr = F.mvdr_weights_rtf(rtf, psd_n, reference_channel, diagonal_loading, diag_eps, eps)
+ spectrum_enhanced = F.apply_beamforming(w_mvdr, specgram)
+ return spectrum_enhanced
+
+
+class SoudenMVDR(torch.nn.Module):
+ r"""Minimum Variance Distortionless Response (*MVDR* :cite:`capon1969high`) module
+ based on the method proposed by *Souden et, al.* :cite:`souden2009optimal`.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Given the multi-channel complex-valued spectrum :math:`\textbf{Y}`, the power spectral density (PSD) matrix
+ of target speech :math:`\bf{\Phi}_{\textbf{SS}}`, the PSD matrix of noise :math:`\bf{\Phi}_{\textbf{NN}}`, and
+ a one-hot vector that represents the reference channel :math:`\bf{u}`, the module computes the single-channel
+ complex-valued spectrum of the enhanced speech :math:`\hat{\textbf{S}}`. The formula is defined as:
+
+ .. math::
+ \hat{\textbf{S}}(f) = \textbf{w}_{\text{bf}}(f)^{\mathsf{H}} \textbf{Y}(f)
+
+ where :math:`\textbf{w}_{\text{bf}}(f)` is the MVDR beamforming weight for the :math:`f`-th frequency bin.
+
+ The beamforming weight is computed by:
+
+ .. math::
+ \textbf{w}_{\text{MVDR}}(f) =
+ \frac{{{\bf{\Phi}_{\textbf{NN}}^{-1}}(f){\bf{\Phi}_{\textbf{SS}}}}(f)}
+ {\text{Trace}({{{\bf{\Phi}_{\textbf{NN}}^{-1}}(f) \bf{\Phi}_{\textbf{SS}}}(f))}}\bm{u}
+ """
+
+ def forward(
+ self,
+ specgram: Tensor,
+ psd_s: Tensor,
+ psd_n: Tensor,
+ reference_channel: Union[int, Tensor],
+ diagonal_loading: bool = True,
+ diag_eps: float = 1e-7,
+ eps: float = 1e-8,
+ ) -> torch.Tensor:
+ """
+ Args:
+ specgram (torch.Tensor): Multi-channel complex-valued spectrum.
+ Tensor with dimensions `(..., channel, freq, time)`.
+ psd_s (torch.Tensor): The complex-valued power spectral density (PSD) matrix of target speech.
+ Tensor with dimensions `(..., freq, channel, channel)`.
+ psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise.
+ Tensor with dimensions `(..., freq, channel, channel)`.
+ reference_channel (int or torch.Tensor): Specifies the reference channel.
+ If the dtype is ``int``, it represents the reference channel index.
+ If the dtype is ``torch.Tensor``, its shape is `(..., channel)`, where the ``channel`` dimension
+ is one-hot.
+ diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to ``psd_n``.
+ (Default: ``True``)
+ diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading.
+ It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``)
+ eps (float, optional): Value to add to the denominator in the beamforming weight formula.
+ (Default: ``1e-8``)
+
+ Returns:
+ torch.Tensor: Single-channel complex-valued enhanced spectrum with dimensions `(..., freq, time)`.
+ """
+ w_mvdr = F.mvdr_weights_souden(psd_s, psd_n, reference_channel, diagonal_loading, diag_eps, eps)
+ spectrum_enhanced = F.apply_beamforming(w_mvdr, specgram)
+ return spectrum_enhanced
diff --git a/lib/python3.12/site-packages/torchaudio/transforms/_transforms.py b/lib/python3.12/site-packages/torchaudio/transforms/_transforms.py
new file mode 100644
index 0000000000000000000000000000000000000000..802cbd3d7783d6117fa2d5361d55db567aba43d5
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/transforms/_transforms.py
@@ -0,0 +1,2137 @@
+# -*- coding: utf-8 -*-
+
+import math
+import warnings
+from typing import Callable, Optional, Sequence, Tuple, Union
+
+import torch
+from torch import Tensor
+from torch.nn.modules.lazy import LazyModuleMixin
+from torch.nn.parameter import UninitializedParameter
+
+from torchaudio import functional as F
+from torchaudio.functional.functional import (
+ _apply_sinc_resample_kernel,
+ _check_convolve_mode,
+ _fix_waveform_shape,
+ _get_sinc_resample_kernel,
+ _stretch_waveform,
+)
+
+__all__ = []
+
+
+class Spectrogram(torch.nn.Module):
+ r"""Create a spectrogram from a audio signal.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ n_fft (int, optional): Size of FFT, creates ``n_fft // 2 + 1`` bins. (Default: ``400``)
+ win_length (int or None, optional): Window size. (Default: ``n_fft``)
+ hop_length (int or None, optional): Length of hop between STFT windows. (Default: ``win_length // 2``)
+ pad (int, optional): Two sided padding of signal. (Default: ``0``)
+ window_fn (Callable[..., Tensor], optional): A function to create a window tensor
+ that is applied/multiplied to each frame/window. (Default: ``torch.hann_window``)
+ power (float or None, optional): Exponent for the magnitude spectrogram,
+ (must be > 0) e.g., 1 for magnitude, 2 for power, etc.
+ If None, then the complex spectrum is returned instead. (Default: ``2``)
+ normalized (bool or str, optional): Whether to normalize by magnitude after stft. If input is str, choices are
+ ``"window"`` and ``"frame_length"``, if specific normalization type is desirable. ``True`` maps to
+ ``"window"``. (Default: ``False``)
+ wkwargs (dict or None, optional): Arguments for window function. (Default: ``None``)
+ center (bool, optional): whether to pad :attr:`waveform` on both sides so
+ that the :math:`t`-th frame is centered at time :math:`t \times \text{hop\_length}`.
+ (Default: ``True``)
+ pad_mode (string, optional): controls the padding method used when
+ :attr:`center` is ``True``. (Default: ``"reflect"``)
+ onesided (bool, optional): controls whether to return half of results to
+ avoid redundancy (Default: ``True``)
+ return_complex (bool, optional):
+ Deprecated and not used.
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = torchaudio.transforms.Spectrogram(n_fft=800)
+ >>> spectrogram = transform(waveform)
+
+ """
+ __constants__ = ["n_fft", "win_length", "hop_length", "pad", "power", "normalized"]
+
+ def __init__(
+ self,
+ n_fft: int = 400,
+ win_length: Optional[int] = None,
+ hop_length: Optional[int] = None,
+ pad: int = 0,
+ window_fn: Callable[..., Tensor] = torch.hann_window,
+ power: Optional[float] = 2.0,
+ normalized: Union[bool, str] = False,
+ wkwargs: Optional[dict] = None,
+ center: bool = True,
+ pad_mode: str = "reflect",
+ onesided: bool = True,
+ return_complex: Optional[bool] = None,
+ ) -> None:
+ super(Spectrogram, self).__init__()
+ torch._C._log_api_usage_once("torchaudio.transforms.Spectrogram")
+ self.n_fft = n_fft
+ # number of FFT bins. the returned STFT result will have n_fft // 2 + 1
+ # number of frequencies due to onesided=True in torch.stft
+ self.win_length = win_length if win_length is not None else n_fft
+ self.hop_length = hop_length if hop_length is not None else self.win_length // 2
+ window = window_fn(self.win_length) if wkwargs is None else window_fn(self.win_length, **wkwargs)
+ self.register_buffer("window", window)
+ self.pad = pad
+ self.power = power
+ self.normalized = normalized
+ self.center = center
+ self.pad_mode = pad_mode
+ self.onesided = onesided
+ if return_complex is not None:
+ warnings.warn(
+ "`return_complex` argument is now deprecated and is not effective."
+ "`torchaudio.transforms.Spectrogram(power=None)` always returns a tensor with "
+ "complex dtype. Please remove the argument in the function call."
+ )
+
+ def forward(self, waveform: Tensor) -> Tensor:
+ r"""
+ Args:
+ waveform (Tensor): Tensor of audio of dimension (..., time).
+
+ Returns:
+ Tensor: Dimension (..., freq, time), where freq is
+ ``n_fft // 2 + 1`` where ``n_fft`` is the number of
+ Fourier bins, and time is the number of window hops (n_frame).
+ """
+ return F.spectrogram(
+ waveform,
+ self.pad,
+ self.window,
+ self.n_fft,
+ self.hop_length,
+ self.win_length,
+ self.power,
+ self.normalized,
+ self.center,
+ self.pad_mode,
+ self.onesided,
+ )
+
+
+class InverseSpectrogram(torch.nn.Module):
+ r"""Create an inverse spectrogram to recover an audio signal from a spectrogram.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ n_fft (int, optional): Size of FFT, creates ``n_fft // 2 + 1`` bins. (Default: ``400``)
+ win_length (int or None, optional): Window size. (Default: ``n_fft``)
+ hop_length (int or None, optional): Length of hop between STFT windows. (Default: ``win_length // 2``)
+ pad (int, optional): Two sided padding of signal. (Default: ``0``)
+ window_fn (Callable[..., Tensor], optional): A function to create a window tensor
+ that is applied/multiplied to each frame/window. (Default: ``torch.hann_window``)
+ normalized (bool or str, optional): Whether the stft output was normalized by magnitude. If input is str,
+ choices are ``"window"`` and ``"frame_length"``, dependent on normalization mode. ``True`` maps to
+ ``"window"``. (Default: ``False``)
+ wkwargs (dict or None, optional): Arguments for window function. (Default: ``None``)
+ center (bool, optional): whether the signal in spectrogram was padded on both sides so
+ that the :math:`t`-th frame is centered at time :math:`t \times \text{hop\_length}`.
+ (Default: ``True``)
+ pad_mode (string, optional): controls the padding method used when
+ :attr:`center` is ``True``. (Default: ``"reflect"``)
+ onesided (bool, optional): controls whether spectrogram was used to return half of results to
+ avoid redundancy (Default: ``True``)
+
+ Example
+ >>> batch, freq, time = 2, 257, 100
+ >>> length = 25344
+ >>> spectrogram = torch.randn(batch, freq, time, dtype=torch.cdouble)
+ >>> transform = transforms.InverseSpectrogram(n_fft=512)
+ >>> waveform = transform(spectrogram, length)
+ """
+ __constants__ = ["n_fft", "win_length", "hop_length", "pad", "power", "normalized"]
+
+ def __init__(
+ self,
+ n_fft: int = 400,
+ win_length: Optional[int] = None,
+ hop_length: Optional[int] = None,
+ pad: int = 0,
+ window_fn: Callable[..., Tensor] = torch.hann_window,
+ normalized: Union[bool, str] = False,
+ wkwargs: Optional[dict] = None,
+ center: bool = True,
+ pad_mode: str = "reflect",
+ onesided: bool = True,
+ ) -> None:
+ super(InverseSpectrogram, self).__init__()
+ self.n_fft = n_fft
+ # number of FFT bins. the returned STFT result will have n_fft // 2 + 1
+ # number of frequencies due to onesided=True in torch.stft
+ self.win_length = win_length if win_length is not None else n_fft
+ self.hop_length = hop_length if hop_length is not None else self.win_length // 2
+ window = window_fn(self.win_length) if wkwargs is None else window_fn(self.win_length, **wkwargs)
+ self.register_buffer("window", window)
+ self.pad = pad
+ self.normalized = normalized
+ self.center = center
+ self.pad_mode = pad_mode
+ self.onesided = onesided
+
+ def forward(self, spectrogram: Tensor, length: Optional[int] = None) -> Tensor:
+ r"""
+ Args:
+ spectrogram (Tensor): Complex tensor of audio of dimension (..., freq, time).
+ length (int or None, optional): The output length of the waveform.
+
+ Returns:
+ Tensor: Dimension (..., time), Least squares estimation of the original signal.
+ """
+ return F.inverse_spectrogram(
+ spectrogram,
+ length,
+ self.pad,
+ self.window,
+ self.n_fft,
+ self.hop_length,
+ self.win_length,
+ self.normalized,
+ self.center,
+ self.pad_mode,
+ self.onesided,
+ )
+
+
+class GriffinLim(torch.nn.Module):
+ r"""Compute waveform from a linear scale magnitude spectrogram using the Griffin-Lim transformation.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Implementation ported from
+ *librosa* :cite:`brian_mcfee-proc-scipy-2015`, *A fast Griffin-Lim algorithm* :cite:`6701851`
+ and *Signal estimation from modified short-time Fourier transform* :cite:`1172092`.
+
+ Args:
+ n_fft (int, optional): Size of FFT, creates ``n_fft // 2 + 1`` bins. (Default: ``400``)
+ n_iter (int, optional): Number of iteration for phase recovery process. (Default: ``32``)
+ win_length (int or None, optional): Window size. (Default: ``n_fft``)
+ hop_length (int or None, optional): Length of hop between STFT windows. (Default: ``win_length // 2``)
+ window_fn (Callable[..., Tensor], optional): A function to create a window tensor
+ that is applied/multiplied to each frame/window. (Default: ``torch.hann_window``)
+ power (float, optional): Exponent for the magnitude spectrogram,
+ (must be > 0) e.g., 1 for magnitude, 2 for power, etc. (Default: ``2``)
+ wkwargs (dict or None, optional): Arguments for window function. (Default: ``None``)
+ momentum (float, optional): The momentum parameter for fast Griffin-Lim.
+ Setting this to 0 recovers the original Griffin-Lim method.
+ Values near 1 can lead to faster convergence, but above 1 may not converge. (Default: ``0.99``)
+ length (int, optional): Array length of the expected output. (Default: ``None``)
+ rand_init (bool, optional): Initializes phase randomly if True and to zero otherwise. (Default: ``True``)
+
+ Example
+ >>> batch, freq, time = 2, 257, 100
+ >>> spectrogram = torch.randn(batch, freq, time)
+ >>> transform = transforms.GriffinLim(n_fft=512)
+ >>> waveform = transform(spectrogram)
+ """
+ __constants__ = ["n_fft", "n_iter", "win_length", "hop_length", "power", "length", "momentum", "rand_init"]
+
+ def __init__(
+ self,
+ n_fft: int = 400,
+ n_iter: int = 32,
+ win_length: Optional[int] = None,
+ hop_length: Optional[int] = None,
+ window_fn: Callable[..., Tensor] = torch.hann_window,
+ power: float = 2.0,
+ wkwargs: Optional[dict] = None,
+ momentum: float = 0.99,
+ length: Optional[int] = None,
+ rand_init: bool = True,
+ ) -> None:
+ super(GriffinLim, self).__init__()
+
+ if not (0 <= momentum < 1):
+ raise ValueError("momentum must be in the range [0, 1). Found: {}".format(momentum))
+
+ self.n_fft = n_fft
+ self.n_iter = n_iter
+ self.win_length = win_length if win_length is not None else n_fft
+ self.hop_length = hop_length if hop_length is not None else self.win_length // 2
+ window = window_fn(self.win_length) if wkwargs is None else window_fn(self.win_length, **wkwargs)
+ self.register_buffer("window", window)
+ self.length = length
+ self.power = power
+ self.momentum = momentum
+ self.rand_init = rand_init
+
+ def forward(self, specgram: Tensor) -> Tensor:
+ r"""
+ Args:
+ specgram (Tensor):
+ A magnitude-only STFT spectrogram of dimension (..., freq, frames)
+ where freq is ``n_fft // 2 + 1``.
+
+ Returns:
+ Tensor: waveform of (..., time), where time equals the ``length`` parameter if given.
+ """
+ return F.griffinlim(
+ specgram,
+ self.window,
+ self.n_fft,
+ self.hop_length,
+ self.win_length,
+ self.power,
+ self.n_iter,
+ self.momentum,
+ self.length,
+ self.rand_init,
+ )
+
+
+class AmplitudeToDB(torch.nn.Module):
+ r"""Turn a tensor from the power/amplitude scale to the decibel scale.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ This output depends on the maximum value in the input tensor, and so
+ may return different values for an audio clip split into snippets vs. a
+ a full clip.
+
+ Args:
+ stype (str, optional): scale of input tensor (``"power"`` or ``"magnitude"``). The
+ power being the elementwise square of the magnitude. (Default: ``"power"``)
+ top_db (float or None, optional): minimum negative cut-off in decibels. A reasonable
+ number is 80. (Default: ``None``)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = transforms.AmplitudeToDB(stype="amplitude", top_db=80)
+ >>> waveform_db = transform(waveform)
+ """
+ __constants__ = ["multiplier", "amin", "ref_value", "db_multiplier"]
+
+ def __init__(self, stype: str = "power", top_db: Optional[float] = None) -> None:
+ super(AmplitudeToDB, self).__init__()
+ self.stype = stype
+ if top_db is not None and top_db < 0:
+ raise ValueError("top_db must be positive value")
+ self.top_db = top_db
+ self.multiplier = 10.0 if stype == "power" else 20.0
+ self.amin = 1e-10
+ self.ref_value = 1.0
+ self.db_multiplier = math.log10(max(self.amin, self.ref_value))
+
+ def forward(self, x: Tensor) -> Tensor:
+ r"""Numerically stable implementation from Librosa.
+
+ https://librosa.org/doc/latest/generated/librosa.amplitude_to_db.html
+
+ Args:
+ x (Tensor): Input tensor before being converted to decibel scale.
+
+ Returns:
+ Tensor: Output tensor in decibel scale.
+ """
+ return F.amplitude_to_DB(x, self.multiplier, self.amin, self.db_multiplier, self.top_db)
+
+
+class MelScale(torch.nn.Module):
+ r"""Turn a normal STFT into a mel frequency STFT with triangular filter banks.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ n_mels (int, optional): Number of mel filterbanks. (Default: ``128``)
+ sample_rate (int, optional): Sample rate of audio signal. (Default: ``16000``)
+ f_min (float, optional): Minimum frequency. (Default: ``0.``)
+ f_max (float or None, optional): Maximum frequency. (Default: ``sample_rate // 2``)
+ n_stft (int, optional): Number of bins in STFT. See ``n_fft`` in :class:`Spectrogram`. (Default: ``201``)
+ norm (str or None, optional): If ``"slaney"``, divide the triangular mel weights by the width of the mel band
+ (area normalization). (Default: ``None``)
+ mel_scale (str, optional): Scale to use: ``htk`` or ``slaney``. (Default: ``htk``)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> spectrogram_transform = transforms.Spectrogram(n_fft=1024)
+ >>> spectrogram = spectrogram_transform(waveform)
+ >>> melscale_transform = transforms.MelScale(sample_rate=sample_rate, n_stft=1024 // 2 + 1)
+ >>> melscale_spectrogram = melscale_transform(spectrogram)
+
+ See also:
+ :py:func:`torchaudio.functional.melscale_fbanks` - The function used to
+ generate the filter banks.
+ """
+ __constants__ = ["n_mels", "sample_rate", "f_min", "f_max"]
+
+ def __init__(
+ self,
+ n_mels: int = 128,
+ sample_rate: int = 16000,
+ f_min: float = 0.0,
+ f_max: Optional[float] = None,
+ n_stft: int = 201,
+ norm: Optional[str] = None,
+ mel_scale: str = "htk",
+ ) -> None:
+ super(MelScale, self).__init__()
+ self.n_mels = n_mels
+ self.sample_rate = sample_rate
+ self.f_max = f_max if f_max is not None else float(sample_rate // 2)
+ self.f_min = f_min
+ self.norm = norm
+ self.mel_scale = mel_scale
+
+ if f_min > self.f_max:
+ raise ValueError("Require f_min: {} <= f_max: {}".format(f_min, self.f_max))
+
+ fb = F.melscale_fbanks(n_stft, self.f_min, self.f_max, self.n_mels, self.sample_rate, self.norm, self.mel_scale)
+ self.register_buffer("fb", fb)
+
+ def forward(self, specgram: Tensor) -> Tensor:
+ r"""
+ Args:
+ specgram (Tensor): A spectrogram STFT of dimension (..., freq, time).
+
+ Returns:
+ Tensor: Mel frequency spectrogram of size (..., ``n_mels``, time).
+ """
+
+ # (..., time, freq) dot (freq, n_mels) -> (..., n_mels, time)
+ mel_specgram = torch.matmul(specgram.transpose(-1, -2), self.fb).transpose(-1, -2)
+
+ return mel_specgram
+
+
+class InverseMelScale(torch.nn.Module):
+ r"""Estimate a STFT in normal frequency domain from mel frequency domain.
+
+ .. devices:: CPU CUDA
+
+ It minimizes the euclidian norm between the input mel-spectrogram and the product between
+ the estimated spectrogram and the filter banks using `torch.linalg.lstsq`.
+
+ Args:
+ n_stft (int): Number of bins in STFT. See ``n_fft`` in :class:`Spectrogram`.
+ n_mels (int, optional): Number of mel filterbanks. (Default: ``128``)
+ sample_rate (int, optional): Sample rate of audio signal. (Default: ``16000``)
+ f_min (float, optional): Minimum frequency. (Default: ``0.``)
+ f_max (float or None, optional): Maximum frequency. (Default: ``sample_rate // 2``)
+ norm (str or None, optional): If "slaney", divide the triangular mel weights by the width of the mel band
+ (area normalization). (Default: ``None``)
+ mel_scale (str, optional): Scale to use: ``htk`` or ``slaney``. (Default: ``htk``)
+ driver (str, optional): Name of the LAPACK/MAGMA method to be used for `torch.lstsq`.
+ For CPU inputs the valid values are ``"gels"``, ``"gelsy"``, ``"gelsd"``, ``"gelss"``.
+ For CUDA input, the only valid driver is ``"gels"``, which assumes that A is full-rank.
+ (Default: ``"gels``)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> mel_spectrogram_transform = transforms.MelSpectrogram(sample_rate, n_fft=1024)
+ >>> mel_spectrogram = mel_spectrogram_transform(waveform)
+ >>> inverse_melscale_transform = transforms.InverseMelScale(n_stft=1024 // 2 + 1)
+ >>> spectrogram = inverse_melscale_transform(mel_spectrogram)
+ """
+ __constants__ = [
+ "n_stft",
+ "n_mels",
+ "sample_rate",
+ "f_min",
+ "f_max",
+ ]
+
+ def __init__(
+ self,
+ n_stft: int,
+ n_mels: int = 128,
+ sample_rate: int = 16000,
+ f_min: float = 0.0,
+ f_max: Optional[float] = None,
+ norm: Optional[str] = None,
+ mel_scale: str = "htk",
+ driver: str = "gels",
+ ) -> None:
+ super(InverseMelScale, self).__init__()
+ self.n_mels = n_mels
+ self.sample_rate = sample_rate
+ self.f_max = f_max or float(sample_rate // 2)
+ self.f_min = f_min
+ self.driver = driver
+
+ if f_min > self.f_max:
+ raise ValueError("Require f_min: {} <= f_max: {}".format(f_min, self.f_max))
+
+ if driver not in ["gels", "gelsy", "gelsd", "gelss"]:
+ raise ValueError(f'driver must be one of ["gels", "gelsy", "gelsd", "gelss"]. Found {driver}.')
+
+ fb = F.melscale_fbanks(n_stft, self.f_min, self.f_max, self.n_mels, self.sample_rate, norm, mel_scale)
+ self.register_buffer("fb", fb)
+
+ def forward(self, melspec: Tensor) -> Tensor:
+ r"""
+ Args:
+ melspec (Tensor): A Mel frequency spectrogram of dimension (..., ``n_mels``, time)
+
+ Returns:
+ Tensor: Linear scale spectrogram of size (..., freq, time)
+ """
+ # pack batch
+ shape = melspec.size()
+ melspec = melspec.view(-1, shape[-2], shape[-1])
+
+ n_mels, time = shape[-2], shape[-1]
+ freq, _ = self.fb.size() # (freq, n_mels)
+ if self.n_mels != n_mels:
+ raise ValueError("Expected an input with {} mel bins. Found: {}".format(self.n_mels, n_mels))
+
+ specgram = torch.relu(torch.linalg.lstsq(self.fb.transpose(-1, -2)[None], melspec, driver=self.driver).solution)
+
+ # unpack batch
+ specgram = specgram.view(shape[:-2] + (freq, time))
+ return specgram
+
+
+class MelSpectrogram(torch.nn.Module):
+ r"""Create MelSpectrogram for a raw audio signal.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ This is a composition of :py:func:`torchaudio.transforms.Spectrogram`
+ and :py:func:`torchaudio.transforms.MelScale`.
+
+ Sources
+ * https://gist.github.com/kastnerkyle/179d6e9a88202ab0a2fe
+ * https://timsainb.github.io/spectrograms-mfccs-and-inversion-in-python.html
+ * http://haythamfayek.com/2016/04/21/speech-processing-for-machine-learning.html
+
+ Args:
+ sample_rate (int, optional): Sample rate of audio signal. (Default: ``16000``)
+ n_fft (int, optional): Size of FFT, creates ``n_fft // 2 + 1`` bins. (Default: ``400``)
+ win_length (int or None, optional): Window size. (Default: ``n_fft``)
+ hop_length (int or None, optional): Length of hop between STFT windows. (Default: ``win_length // 2``)
+ f_min (float, optional): Minimum frequency. (Default: ``0.``)
+ f_max (float or None, optional): Maximum frequency. (Default: ``None``)
+ pad (int, optional): Two sided padding of signal. (Default: ``0``)
+ n_mels (int, optional): Number of mel filterbanks. (Default: ``128``)
+ window_fn (Callable[..., Tensor], optional): A function to create a window tensor
+ that is applied/multiplied to each frame/window. (Default: ``torch.hann_window``)
+ power (float, optional): Exponent for the magnitude spectrogram,
+ (must be > 0) e.g., 1 for magnitude, 2 for power, etc. (Default: ``2``)
+ normalized (bool, optional): Whether to normalize by magnitude after stft. (Default: ``False``)
+ wkwargs (Dict[..., ...] or None, optional): Arguments for window function. (Default: ``None``)
+ center (bool, optional): whether to pad :attr:`waveform` on both sides so
+ that the :math:`t`-th frame is centered at time :math:`t \times \text{hop\_length}`.
+ (Default: ``True``)
+ pad_mode (string, optional): controls the padding method used when
+ :attr:`center` is ``True``. (Default: ``"reflect"``)
+ onesided: Deprecated and unused.
+ norm (str or None, optional): If "slaney", divide the triangular mel weights by the width of the mel band
+ (area normalization). (Default: ``None``)
+ mel_scale (str, optional): Scale to use: ``htk`` or ``slaney``. (Default: ``htk``)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = transforms.MelSpectrogram(sample_rate)
+ >>> mel_specgram = transform(waveform) # (channel, n_mels, time)
+
+ See also:
+ :py:func:`torchaudio.functional.melscale_fbanks` - The function used to
+ generate the filter banks.
+ """
+ __constants__ = ["sample_rate", "n_fft", "win_length", "hop_length", "pad", "n_mels", "f_min"]
+
+ def __init__(
+ self,
+ sample_rate: int = 16000,
+ n_fft: int = 400,
+ win_length: Optional[int] = None,
+ hop_length: Optional[int] = None,
+ f_min: float = 0.0,
+ f_max: Optional[float] = None,
+ pad: int = 0,
+ n_mels: int = 128,
+ window_fn: Callable[..., Tensor] = torch.hann_window,
+ power: float = 2.0,
+ normalized: bool = False,
+ wkwargs: Optional[dict] = None,
+ center: bool = True,
+ pad_mode: str = "reflect",
+ onesided: Optional[bool] = None,
+ norm: Optional[str] = None,
+ mel_scale: str = "htk",
+ ) -> None:
+ super(MelSpectrogram, self).__init__()
+ torch._C._log_api_usage_once("torchaudio.transforms.MelSpectrogram")
+
+ if onesided is not None:
+ warnings.warn(
+ "Argument 'onesided' has been deprecated and has no influence on the behavior of this module."
+ )
+
+ self.sample_rate = sample_rate
+ self.n_fft = n_fft
+ self.win_length = win_length if win_length is not None else n_fft
+ self.hop_length = hop_length if hop_length is not None else self.win_length // 2
+ self.pad = pad
+ self.power = power
+ self.normalized = normalized
+ self.n_mels = n_mels # number of mel frequency bins
+ self.f_max = f_max
+ self.f_min = f_min
+ self.spectrogram = Spectrogram(
+ n_fft=self.n_fft,
+ win_length=self.win_length,
+ hop_length=self.hop_length,
+ pad=self.pad,
+ window_fn=window_fn,
+ power=self.power,
+ normalized=self.normalized,
+ wkwargs=wkwargs,
+ center=center,
+ pad_mode=pad_mode,
+ onesided=True,
+ )
+ self.mel_scale = MelScale(
+ self.n_mels, self.sample_rate, self.f_min, self.f_max, self.n_fft // 2 + 1, norm, mel_scale
+ )
+
+ def forward(self, waveform: Tensor) -> Tensor:
+ r"""
+ Args:
+ waveform (Tensor): Tensor of audio of dimension (..., time).
+
+ Returns:
+ Tensor: Mel frequency spectrogram of size (..., ``n_mels``, time).
+ """
+ specgram = self.spectrogram(waveform)
+ mel_specgram = self.mel_scale(specgram)
+ return mel_specgram
+
+
+class MFCC(torch.nn.Module):
+ r"""Create the Mel-frequency cepstrum coefficients from an audio signal.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ By default, this calculates the MFCC on the DB-scaled Mel spectrogram.
+ This is not the textbook implementation, but is implemented here to
+ give consistency with librosa.
+
+ This output depends on the maximum value in the input spectrogram, and so
+ may return different values for an audio clip split into snippets vs. a
+ a full clip.
+
+ Args:
+ sample_rate (int, optional): Sample rate of audio signal. (Default: ``16000``)
+ n_mfcc (int, optional): Number of mfc coefficients to retain. (Default: ``40``)
+ dct_type (int, optional): type of DCT (discrete cosine transform) to use. (Default: ``2``)
+ norm (str, optional): norm to use. (Default: ``"ortho"``)
+ log_mels (bool, optional): whether to use log-mel spectrograms instead of db-scaled. (Default: ``False``)
+ melkwargs (dict or None, optional): arguments for MelSpectrogram. (Default: ``None``)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = transforms.MFCC(
+ >>> sample_rate=sample_rate,
+ >>> n_mfcc=13,
+ >>> melkwargs={"n_fft": 400, "hop_length": 160, "n_mels": 23, "center": False},
+ >>> )
+ >>> mfcc = transform(waveform)
+
+ See also:
+ :py:func:`torchaudio.functional.melscale_fbanks` - The function used to
+ generate the filter banks.
+ """
+ __constants__ = ["sample_rate", "n_mfcc", "dct_type", "top_db", "log_mels"]
+
+ def __init__(
+ self,
+ sample_rate: int = 16000,
+ n_mfcc: int = 40,
+ dct_type: int = 2,
+ norm: str = "ortho",
+ log_mels: bool = False,
+ melkwargs: Optional[dict] = None,
+ ) -> None:
+ super(MFCC, self).__init__()
+ supported_dct_types = [2]
+ if dct_type not in supported_dct_types:
+ raise ValueError("DCT type not supported: {}".format(dct_type))
+ self.sample_rate = sample_rate
+ self.n_mfcc = n_mfcc
+ self.dct_type = dct_type
+ self.norm = norm
+ self.top_db = 80.0
+ self.amplitude_to_DB = AmplitudeToDB("power", self.top_db)
+
+ melkwargs = melkwargs or {}
+ self.MelSpectrogram = MelSpectrogram(sample_rate=self.sample_rate, **melkwargs)
+
+ if self.n_mfcc > self.MelSpectrogram.n_mels:
+ raise ValueError("Cannot select more MFCC coefficients than # mel bins")
+ dct_mat = F.create_dct(self.n_mfcc, self.MelSpectrogram.n_mels, self.norm)
+ self.register_buffer("dct_mat", dct_mat)
+ self.log_mels = log_mels
+
+ def forward(self, waveform: Tensor) -> Tensor:
+ r"""
+ Args:
+ waveform (Tensor): Tensor of audio of dimension (..., time).
+
+ Returns:
+ Tensor: specgram_mel_db of size (..., ``n_mfcc``, time).
+ """
+ mel_specgram = self.MelSpectrogram(waveform)
+ if self.log_mels:
+ log_offset = 1e-6
+ mel_specgram = torch.log(mel_specgram + log_offset)
+ else:
+ mel_specgram = self.amplitude_to_DB(mel_specgram)
+
+ # (..., time, n_mels) dot (n_mels, n_mfcc) -> (..., n_nfcc, time)
+ mfcc = torch.matmul(mel_specgram.transpose(-1, -2), self.dct_mat).transpose(-1, -2)
+ return mfcc
+
+
+class LFCC(torch.nn.Module):
+ r"""Create the linear-frequency cepstrum coefficients from an audio signal.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ By default, this calculates the LFCC on the DB-scaled linear filtered spectrogram.
+ This is not the textbook implementation, but is implemented here to
+ give consistency with librosa.
+
+ This output depends on the maximum value in the input spectrogram, and so
+ may return different values for an audio clip split into snippets vs. a
+ a full clip.
+
+ Args:
+ sample_rate (int, optional): Sample rate of audio signal. (Default: ``16000``)
+ n_filter (int, optional): Number of linear filters to apply. (Default: ``128``)
+ n_lfcc (int, optional): Number of lfc coefficients to retain. (Default: ``40``)
+ f_min (float, optional): Minimum frequency. (Default: ``0.``)
+ f_max (float or None, optional): Maximum frequency. (Default: ``None``)
+ dct_type (int, optional): type of DCT (discrete cosine transform) to use. (Default: ``2``)
+ norm (str, optional): norm to use. (Default: ``"ortho"``)
+ log_lf (bool, optional): whether to use log-lf spectrograms instead of db-scaled. (Default: ``False``)
+ speckwargs (dict or None, optional): arguments for Spectrogram. (Default: ``None``)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = transforms.LFCC(
+ >>> sample_rate=sample_rate,
+ >>> n_lfcc=13,
+ >>> speckwargs={"n_fft": 400, "hop_length": 160, "center": False},
+ >>> )
+ >>> lfcc = transform(waveform)
+
+ See also:
+ :py:func:`torchaudio.functional.linear_fbanks` - The function used to
+ generate the filter banks.
+ """
+ __constants__ = ["sample_rate", "n_filter", "n_lfcc", "dct_type", "top_db", "log_lf"]
+
+ def __init__(
+ self,
+ sample_rate: int = 16000,
+ n_filter: int = 128,
+ f_min: float = 0.0,
+ f_max: Optional[float] = None,
+ n_lfcc: int = 40,
+ dct_type: int = 2,
+ norm: str = "ortho",
+ log_lf: bool = False,
+ speckwargs: Optional[dict] = None,
+ ) -> None:
+ super(LFCC, self).__init__()
+ supported_dct_types = [2]
+ if dct_type not in supported_dct_types:
+ raise ValueError("DCT type not supported: {}".format(dct_type))
+ self.sample_rate = sample_rate
+ self.f_min = f_min
+ self.f_max = f_max if f_max is not None else float(sample_rate // 2)
+ self.n_filter = n_filter
+ self.n_lfcc = n_lfcc
+ self.dct_type = dct_type
+ self.norm = norm
+ self.top_db = 80.0
+ self.amplitude_to_DB = AmplitudeToDB("power", self.top_db)
+
+ speckwargs = speckwargs or {}
+ self.Spectrogram = Spectrogram(**speckwargs)
+
+ if self.n_lfcc > self.Spectrogram.n_fft:
+ raise ValueError("Cannot select more LFCC coefficients than # fft bins")
+
+ filter_mat = F.linear_fbanks(
+ n_freqs=self.Spectrogram.n_fft // 2 + 1,
+ f_min=self.f_min,
+ f_max=self.f_max,
+ n_filter=self.n_filter,
+ sample_rate=self.sample_rate,
+ )
+ self.register_buffer("filter_mat", filter_mat)
+
+ dct_mat = F.create_dct(self.n_lfcc, self.n_filter, self.norm)
+ self.register_buffer("dct_mat", dct_mat)
+ self.log_lf = log_lf
+
+ def forward(self, waveform: Tensor) -> Tensor:
+ r"""
+ Args:
+ waveform (Tensor): Tensor of audio of dimension (..., time).
+
+ Returns:
+ Tensor: Linear Frequency Cepstral Coefficients of size (..., ``n_lfcc``, time).
+ """
+ specgram = self.Spectrogram(waveform)
+
+ # (..., time, freq) dot (freq, n_filter) -> (..., n_filter, time)
+ specgram = torch.matmul(specgram.transpose(-1, -2), self.filter_mat).transpose(-1, -2)
+
+ if self.log_lf:
+ log_offset = 1e-6
+ specgram = torch.log(specgram + log_offset)
+ else:
+ specgram = self.amplitude_to_DB(specgram)
+
+ # (..., time, n_filter) dot (n_filter, n_lfcc) -> (..., n_lfcc, time)
+ lfcc = torch.matmul(specgram.transpose(-1, -2), self.dct_mat).transpose(-1, -2)
+ return lfcc
+
+
+class MuLawEncoding(torch.nn.Module):
+ r"""Encode signal based on mu-law companding.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: TorchScript
+
+ For more info see the
+ `Wikipedia Entry `_
+
+ This algorithm assumes the signal has been scaled to between -1 and 1 and
+ returns a signal encoded with values from 0 to quantization_channels - 1
+
+ Args:
+ quantization_channels (int, optional): Number of channels. (Default: ``256``)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = torchaudio.transforms.MuLawEncoding(quantization_channels=512)
+ >>> mulawtrans = transform(waveform)
+
+ """
+ __constants__ = ["quantization_channels"]
+
+ def __init__(self, quantization_channels: int = 256) -> None:
+ super(MuLawEncoding, self).__init__()
+ self.quantization_channels = quantization_channels
+
+ def forward(self, x: Tensor) -> Tensor:
+ r"""
+ Args:
+ x (Tensor): A signal to be encoded.
+
+ Returns:
+ Tensor: An encoded signal.
+ """
+ return F.mu_law_encoding(x, self.quantization_channels)
+
+
+class MuLawDecoding(torch.nn.Module):
+ r"""Decode mu-law encoded signal.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: TorchScript
+
+ For more info see the
+ `Wikipedia Entry `_
+
+ This expects an input with values between 0 and ``quantization_channels - 1``
+ and returns a signal scaled between -1 and 1.
+
+ Args:
+ quantization_channels (int, optional): Number of channels. (Default: ``256``)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = torchaudio.transforms.MuLawDecoding(quantization_channels=512)
+ >>> mulawtrans = transform(waveform)
+ """
+ __constants__ = ["quantization_channels"]
+
+ def __init__(self, quantization_channels: int = 256) -> None:
+ super(MuLawDecoding, self).__init__()
+ self.quantization_channels = quantization_channels
+
+ def forward(self, x_mu: Tensor) -> Tensor:
+ r"""
+ Args:
+ x_mu (Tensor): A mu-law encoded signal which needs to be decoded.
+
+ Returns:
+ Tensor: The signal decoded.
+ """
+ return F.mu_law_decoding(x_mu, self.quantization_channels)
+
+
+class Resample(torch.nn.Module):
+ r"""Resample a signal from one frequency to another. A resampling method can be given.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Note:
+ If resampling on waveforms of higher precision than float32, there may be a small loss of precision
+ because the kernel is cached once as float32. If high precision resampling is important for your application,
+ the functional form will retain higher precision, but run slower because it does not cache the kernel.
+ Alternatively, you could rewrite a transform that caches a higher precision kernel.
+
+ Args:
+ orig_freq (int, optional): The original frequency of the signal. (Default: ``16000``)
+ new_freq (int, optional): The desired frequency. (Default: ``16000``)
+ resampling_method (str, optional): The resampling method to use.
+ Options: [``sinc_interp_hann``, ``sinc_interp_kaiser``] (Default: ``"sinc_interp_hann"``)
+ lowpass_filter_width (int, optional): Controls the sharpness of the filter, more == sharper
+ but less efficient. (Default: ``6``)
+ rolloff (float, optional): The roll-off frequency of the filter, as a fraction of the Nyquist.
+ Lower values reduce anti-aliasing, but also reduce some of the highest frequencies. (Default: ``0.99``)
+ beta (float or None, optional): The shape parameter used for kaiser window.
+ dtype (torch.device, optional):
+ Determnines the precision that resampling kernel is pre-computed and cached. If not provided,
+ kernel is computed with ``torch.float64`` then cached as ``torch.float32``.
+ If you need higher precision, provide ``torch.float64``, and the pre-computed kernel is computed and
+ cached as ``torch.float64``. If you use resample with lower precision, then instead of providing this
+ providing this argument, please use ``Resample.to(dtype)``, so that the kernel generation is still
+ carried out on ``torch.float64``.
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = transforms.Resample(sample_rate, sample_rate/10)
+ >>> waveform = transform(waveform)
+ """
+
+ def __init__(
+ self,
+ orig_freq: int = 16000,
+ new_freq: int = 16000,
+ resampling_method: str = "sinc_interp_hann",
+ lowpass_filter_width: int = 6,
+ rolloff: float = 0.99,
+ beta: Optional[float] = None,
+ *,
+ dtype: Optional[torch.dtype] = None,
+ ) -> None:
+ super().__init__()
+
+ self.orig_freq = orig_freq
+ self.new_freq = new_freq
+ self.gcd = math.gcd(int(self.orig_freq), int(self.new_freq))
+ self.resampling_method = resampling_method
+ self.lowpass_filter_width = lowpass_filter_width
+ self.rolloff = rolloff
+ self.beta = beta
+
+ if self.orig_freq != self.new_freq:
+ kernel, self.width = _get_sinc_resample_kernel(
+ self.orig_freq,
+ self.new_freq,
+ self.gcd,
+ self.lowpass_filter_width,
+ self.rolloff,
+ self.resampling_method,
+ beta,
+ dtype=dtype,
+ )
+ self.register_buffer("kernel", kernel)
+
+ def forward(self, waveform: Tensor) -> Tensor:
+ r"""
+ Args:
+ waveform (Tensor): Tensor of audio of dimension (..., time).
+
+ Returns:
+ Tensor: Output signal of dimension (..., time).
+ """
+ if self.orig_freq == self.new_freq:
+ return waveform
+ return _apply_sinc_resample_kernel(waveform, self.orig_freq, self.new_freq, self.gcd, self.kernel, self.width)
+
+
+class ComputeDeltas(torch.nn.Module):
+ r"""Compute delta coefficients of a tensor, usually a spectrogram.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ See `torchaudio.functional.compute_deltas` for more details.
+
+ Args:
+ win_length (int, optional): The window length used for computing delta. (Default: ``5``)
+ mode (str, optional): Mode parameter passed to padding. (Default: ``"replicate"``)
+ """
+ __constants__ = ["win_length"]
+
+ def __init__(self, win_length: int = 5, mode: str = "replicate") -> None:
+ super(ComputeDeltas, self).__init__()
+ self.win_length = win_length
+ self.mode = mode
+
+ def forward(self, specgram: Tensor) -> Tensor:
+ r"""
+ Args:
+ specgram (Tensor): Tensor of audio of dimension (..., freq, time).
+
+ Returns:
+ Tensor: Tensor of deltas of dimension (..., freq, time).
+ """
+ return F.compute_deltas(specgram, win_length=self.win_length, mode=self.mode)
+
+
+class TimeStretch(torch.nn.Module):
+ r"""Stretch stft in time without modifying pitch for a given rate.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Proposed in *SpecAugment* :cite:`specaugment`.
+
+ Args:
+ hop_length (int or None, optional): Length of hop between STFT windows.
+ (Default: ``n_fft // 2``, where ``n_fft == (n_freq - 1) * 2``)
+ n_freq (int, optional): number of filter banks from stft. (Default: ``201``)
+ fixed_rate (float or None, optional): rate to speed up or slow down by.
+ If None is provided, rate must be passed to the forward method. (Default: ``None``)
+
+ .. note::
+
+ The expected input is raw, complex-valued spectrogram.
+
+ Example
+ >>> spectrogram = torchaudio.transforms.Spectrogram(power=None)
+ >>> stretch = torchaudio.transforms.TimeStretch()
+ >>>
+ >>> original = spectrogram(waveform)
+ >>> stretched_1_2 = stretch(original, 1.2)
+ >>> stretched_0_9 = stretch(original, 0.9)
+
+ .. image:: https://download.pytorch.org/torchaudio/doc-assets/specaugment_time_stretch.png
+ :width: 600
+ :alt: The visualization of stretched spectrograms.
+ """
+ __constants__ = ["fixed_rate"]
+
+ def __init__(self, hop_length: Optional[int] = None, n_freq: int = 201, fixed_rate: Optional[float] = None) -> None:
+ super(TimeStretch, self).__init__()
+
+ self.fixed_rate = fixed_rate
+
+ n_fft = (n_freq - 1) * 2
+ hop_length = hop_length if hop_length is not None else n_fft // 2
+ self.register_buffer("phase_advance", torch.linspace(0, math.pi * hop_length, n_freq)[..., None])
+
+ def forward(self, complex_specgrams: Tensor, overriding_rate: Optional[float] = None) -> Tensor:
+ r"""
+ Args:
+ complex_specgrams (Tensor):
+ A tensor of dimension `(..., freq, num_frame)` with complex dtype.
+ overriding_rate (float or None, optional): speed up to apply to this batch.
+ If no rate is passed, use ``self.fixed_rate``. (Default: ``None``)
+
+ Returns:
+ Tensor:
+ Stretched spectrogram. The resulting tensor is of the corresponding complex dtype
+ as the input spectrogram, and the number of frames is changed to ``ceil(num_frame / rate)``.
+ """
+ if not torch.is_complex(complex_specgrams):
+ warnings.warn(
+ "The input to TimeStretch must be complex type. "
+ "Providing non-complex tensor produces invalid results.",
+ stacklevel=4,
+ )
+
+ if overriding_rate is None:
+ if self.fixed_rate is None:
+ raise ValueError("If no fixed_rate is specified, must pass a valid rate to the forward method.")
+ rate = self.fixed_rate
+ else:
+ rate = overriding_rate
+ return F.phase_vocoder(complex_specgrams, rate, self.phase_advance)
+
+
+class Fade(torch.nn.Module):
+ r"""Add a fade in and/or fade out to an waveform.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ fade_in_len (int, optional): Length of fade-in (time frames). (Default: ``0``)
+ fade_out_len (int, optional): Length of fade-out (time frames). (Default: ``0``)
+ fade_shape (str, optional): Shape of fade. Must be one of: "quarter_sine",
+ ``"half_sine"``, ``"linear"``, ``"logarithmic"``, ``"exponential"``.
+ (Default: ``"linear"``)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = transforms.Fade(fade_in_len=sample_rate, fade_out_len=2 * sample_rate, fade_shape="linear")
+ >>> faded_waveform = transform(waveform)
+ """
+
+ def __init__(self, fade_in_len: int = 0, fade_out_len: int = 0, fade_shape: str = "linear") -> None:
+ super(Fade, self).__init__()
+ self.fade_in_len = fade_in_len
+ self.fade_out_len = fade_out_len
+ self.fade_shape = fade_shape
+
+ def forward(self, waveform: Tensor) -> Tensor:
+ r"""
+ Args:
+ waveform (Tensor): Tensor of audio of dimension `(..., time)`.
+
+ Returns:
+ Tensor: Tensor of audio of dimension `(..., time)`.
+ """
+ waveform_length = waveform.size()[-1]
+ device = waveform.device
+ return self._fade_in(waveform_length, device) * self._fade_out(waveform_length, device) * waveform
+
+ def _fade_in(self, waveform_length: int, device: torch.device) -> Tensor:
+ fade = torch.linspace(0, 1, self.fade_in_len, device=device)
+ ones = torch.ones(waveform_length - self.fade_in_len, device=device)
+
+ if self.fade_shape == "linear":
+ fade = fade
+
+ if self.fade_shape == "exponential":
+ fade = torch.pow(2, (fade - 1)) * fade
+
+ if self.fade_shape == "logarithmic":
+ fade = torch.log10(0.1 + fade) + 1
+
+ if self.fade_shape == "quarter_sine":
+ fade = torch.sin(fade * math.pi / 2)
+
+ if self.fade_shape == "half_sine":
+ fade = torch.sin(fade * math.pi - math.pi / 2) / 2 + 0.5
+
+ return torch.cat((fade, ones)).clamp_(0, 1)
+
+ def _fade_out(self, waveform_length: int, device: torch.device) -> Tensor:
+ fade = torch.linspace(0, 1, self.fade_out_len, device=device)
+ ones = torch.ones(waveform_length - self.fade_out_len, device=device)
+
+ if self.fade_shape == "linear":
+ fade = -fade + 1
+
+ if self.fade_shape == "exponential":
+ fade = torch.pow(2, -fade) * (1 - fade)
+
+ if self.fade_shape == "logarithmic":
+ fade = torch.log10(1.1 - fade) + 1
+
+ if self.fade_shape == "quarter_sine":
+ fade = torch.sin(fade * math.pi / 2 + math.pi / 2)
+
+ if self.fade_shape == "half_sine":
+ fade = torch.sin(fade * math.pi + math.pi / 2) / 2 + 0.5
+
+ return torch.cat((ones, fade)).clamp_(0, 1)
+
+
+class _AxisMasking(torch.nn.Module):
+ r"""Apply masking to a spectrogram.
+
+ Args:
+ mask_param (int): Maximum possible length of the mask.
+ axis (int): What dimension the mask is applied on (assuming the tensor is 3D).
+ For frequency masking, axis = 1.
+ For time masking, axis = 2.
+ iid_masks (bool): Applies iid masks to each of the examples in the batch dimension.
+ This option is applicable only when the dimension of the input tensor is >= 3.
+ p (float, optional): maximum proportion of columns that can be masked. (Default: 1.0)
+ """
+ __constants__ = ["mask_param", "axis", "iid_masks", "p"]
+
+ def __init__(self, mask_param: int, axis: int, iid_masks: bool, p: float = 1.0) -> None:
+ super(_AxisMasking, self).__init__()
+ self.mask_param = mask_param
+ self.axis = axis
+ self.iid_masks = iid_masks
+ self.p = p
+
+ def forward(self, specgram: Tensor, mask_value: float = 0.0) -> Tensor:
+ r"""
+ Args:
+ specgram (Tensor): Tensor of dimension `(..., freq, time)`.
+ mask_value (float): Value to assign to the masked columns.
+
+ Returns:
+ Tensor: Masked spectrogram of dimensions `(..., freq, time)`.
+ """
+ # if iid_masks flag marked and specgram has a batch dimension
+ # self.axis + specgram.dim() - 3 gives the time/frequency dimension (last two dimensions)
+ # for input tensor for which the dimension is not 3.
+ if self.iid_masks:
+ return F.mask_along_axis_iid(
+ specgram, self.mask_param, mask_value, self.axis + specgram.dim() - 3, p=self.p
+ )
+ else:
+ return F.mask_along_axis(specgram, self.mask_param, mask_value, self.axis + specgram.dim() - 3, p=self.p)
+
+
+class FrequencyMasking(_AxisMasking):
+ r"""Apply masking to a spectrogram in the frequency domain.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Proposed in *SpecAugment* :cite:`specaugment`.
+
+ Args:
+ freq_mask_param (int): maximum possible length of the mask.
+ Indices uniformly sampled from [0, freq_mask_param).
+ iid_masks (bool, optional): whether to apply different masks to each
+ example/channel in the batch. (Default: ``False``)
+ This option is applicable only when the input tensor >= 3D.
+
+ Example
+ >>> spectrogram = torchaudio.transforms.Spectrogram()
+ >>> masking = torchaudio.transforms.FrequencyMasking(freq_mask_param=80)
+ >>>
+ >>> original = spectrogram(waveform)
+ >>> masked = masking(original)
+
+ .. image:: https://download.pytorch.org/torchaudio/doc-assets/specaugment_freq_masking1.png
+ :alt: The original spectrogram
+
+ .. image:: https://download.pytorch.org/torchaudio/doc-assets/specaugment_freq_masking2.png
+ :alt: The spectrogram masked along frequency axis
+ """
+
+ def __init__(self, freq_mask_param: int, iid_masks: bool = False) -> None:
+ super(FrequencyMasking, self).__init__(freq_mask_param, 1, iid_masks)
+
+
+class TimeMasking(_AxisMasking):
+ r"""Apply masking to a spectrogram in the time domain.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Proposed in *SpecAugment* :cite:`specaugment`.
+
+ Args:
+ time_mask_param (int): maximum possible length of the mask.
+ Indices uniformly sampled from [0, time_mask_param).
+ iid_masks (bool, optional): whether to apply different masks to each
+ example/channel in the batch. (Default: ``False``)
+ This option is applicable only when the input tensor >= 3D.
+ p (float, optional): maximum proportion of time steps that can be masked.
+ Must be within range [0.0, 1.0]. (Default: 1.0)
+
+ Example
+ >>> spectrogram = torchaudio.transforms.Spectrogram()
+ >>> masking = torchaudio.transforms.TimeMasking(time_mask_param=80)
+ >>>
+ >>> original = spectrogram(waveform)
+ >>> masked = masking(original)
+
+ .. image:: https://download.pytorch.org/torchaudio/doc-assets/specaugment_time_masking1.png
+ :alt: The original spectrogram
+
+ .. image:: https://download.pytorch.org/torchaudio/doc-assets/specaugment_time_masking2.png
+ :alt: The spectrogram masked along time axis
+ """
+
+ def __init__(self, time_mask_param: int, iid_masks: bool = False, p: float = 1.0) -> None:
+ if not 0.0 <= p <= 1.0:
+ raise ValueError(f"The value of p must be between 0.0 and 1.0 ({p} given).")
+ super(TimeMasking, self).__init__(time_mask_param, 2, iid_masks, p=p)
+
+
+class SpecAugment(torch.nn.Module):
+ r"""Apply time and frequency masking to a spectrogram.
+ Args:
+ n_time_masks (int): Number of time masks. If its value is zero, no time masking will be applied.
+ time_mask_param (int): Maximum possible length of the time mask.
+ n_freq_masks (int): Number of frequency masks. If its value is zero, no frequency masking will be applied.
+ freq_mask_param (int): Maximum possible length of the frequency mask.
+ iid_masks (bool, optional): Applies iid masks to each of the examples in the batch dimension.
+ This option is applicable only when the input tensor is 4D. (Default: ``True``)
+ p (float, optional): maximum proportion of time steps that can be masked.
+ Must be within range [0.0, 1.0]. (Default: 1.0)
+ zero_masking (bool, optional): If ``True``, use 0 as the mask value,
+ else use mean of the input tensor. (Default: ``False``)
+ """
+ __constants__ = [
+ "n_time_masks",
+ "time_mask_param",
+ "n_freq_masks",
+ "freq_mask_param",
+ "iid_masks",
+ "p",
+ "zero_masking",
+ ]
+
+ def __init__(
+ self,
+ n_time_masks: int,
+ time_mask_param: int,
+ n_freq_masks: int,
+ freq_mask_param: int,
+ iid_masks: bool = True,
+ p: float = 1.0,
+ zero_masking: bool = False,
+ ) -> None:
+ super(SpecAugment, self).__init__()
+ self.n_time_masks = n_time_masks
+ self.time_mask_param = time_mask_param
+ self.n_freq_masks = n_freq_masks
+ self.freq_mask_param = freq_mask_param
+ self.iid_masks = iid_masks
+ self.p = p
+ self.zero_masking = zero_masking
+
+ def forward(self, specgram: Tensor) -> Tensor:
+ r"""
+ Args:
+ specgram (Tensor): Tensor of shape `(..., freq, time)`.
+ Returns:
+ Tensor: Masked spectrogram of shape `(..., freq, time)`.
+ """
+ if self.zero_masking:
+ mask_value = 0.0
+ else:
+ mask_value = specgram.mean()
+ time_dim = specgram.dim() - 1
+ freq_dim = time_dim - 1
+
+ if specgram.dim() > 2 and self.iid_masks is True:
+ for _ in range(self.n_time_masks):
+ specgram = F.mask_along_axis_iid(specgram, self.time_mask_param, mask_value, time_dim, p=self.p)
+ for _ in range(self.n_freq_masks):
+ specgram = F.mask_along_axis_iid(specgram, self.freq_mask_param, mask_value, freq_dim, p=self.p)
+ else:
+ for _ in range(self.n_time_masks):
+ specgram = F.mask_along_axis(specgram, self.time_mask_param, mask_value, time_dim, p=self.p)
+ for _ in range(self.n_freq_masks):
+ specgram = F.mask_along_axis(specgram, self.freq_mask_param, mask_value, freq_dim, p=self.p)
+
+ return specgram
+
+
+class Loudness(torch.nn.Module):
+ r"""Measure audio loudness according to the ITU-R BS.1770-4 recommendation.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: TorchScript
+
+ Args:
+ sample_rate (int): Sample rate of audio signal.
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = transforms.Loudness(sample_rate)
+ >>> loudness = transform(waveform)
+
+ Reference:
+ - https://www.itu.int/rec/R-REC-BS.1770-4-201510-I/en
+ """
+ __constants__ = ["sample_rate"]
+
+ def __init__(self, sample_rate: int):
+ super(Loudness, self).__init__()
+ self.sample_rate = sample_rate
+
+ def forward(self, wavefrom: Tensor):
+ r"""
+ Args:
+ waveform(torch.Tensor): audio waveform of dimension `(..., channels, time)`
+
+ Returns:
+ Tensor: loudness estimates (LKFS)
+ """
+ return F.loudness(wavefrom, self.sample_rate)
+
+
+class Vol(torch.nn.Module):
+ r"""Adjust volume of waveform.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ gain (float): Interpreted according to the given gain_type:
+ If ``gain_type`` = ``amplitude``, ``gain`` is a positive amplitude ratio.
+ If ``gain_type`` = ``power``, ``gain`` is a power (voltage squared).
+ If ``gain_type`` = ``db``, ``gain`` is in decibels.
+ gain_type (str, optional): Type of gain. One of: ``amplitude``, ``power``, ``db`` (Default: ``amplitude``)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = transforms.Vol(gain=0.5, gain_type="amplitude")
+ >>> quieter_waveform = transform(waveform)
+ """
+
+ def __init__(self, gain: float, gain_type: str = "amplitude"):
+ super(Vol, self).__init__()
+ self.gain = gain
+ self.gain_type = gain_type
+
+ if gain_type in ["amplitude", "power"] and gain < 0:
+ raise ValueError("If gain_type = amplitude or power, gain must be positive.")
+
+ def forward(self, waveform: Tensor) -> Tensor:
+ r"""
+ Args:
+ waveform (Tensor): Tensor of audio of dimension `(..., time)`.
+
+ Returns:
+ Tensor: Tensor of audio of dimension `(..., time)`.
+ """
+ if self.gain_type == "amplitude":
+ waveform = waveform * self.gain
+
+ if self.gain_type == "db":
+ waveform = F.gain(waveform, self.gain)
+
+ if self.gain_type == "power":
+ waveform = F.gain(waveform, 10 * math.log10(self.gain))
+
+ return torch.clamp(waveform, -1, 1)
+
+
+class SlidingWindowCmn(torch.nn.Module):
+ r"""
+ Apply sliding-window cepstral mean (and optionally variance) normalization per utterance.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ cmn_window (int, optional): Window in frames for running average CMN computation (int, default = 600)
+ min_cmn_window (int, optional): Minimum CMN window used at start of decoding (adds latency only at start).
+ Only applicable if center == false, ignored if center==true (int, default = 100)
+ center (bool, optional): If true, use a window centered on the current frame
+ (to the extent possible, modulo end effects). If false, window is to the left. (bool, default = false)
+ norm_vars (bool, optional): If true, normalize variance to one. (bool, default = false)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = transforms.SlidingWindowCmn(cmn_window=1000)
+ >>> cmn_waveform = transform(waveform)
+ """
+
+ def __init__(
+ self, cmn_window: int = 600, min_cmn_window: int = 100, center: bool = False, norm_vars: bool = False
+ ) -> None:
+ super().__init__()
+ self.cmn_window = cmn_window
+ self.min_cmn_window = min_cmn_window
+ self.center = center
+ self.norm_vars = norm_vars
+
+ def forward(self, specgram: Tensor) -> Tensor:
+ r"""
+ Args:
+ specgram (Tensor): Tensor of spectrogram of dimension `(..., time, freq)`.
+
+ Returns:
+ Tensor: Tensor of spectrogram of dimension `(..., time, freq)`.
+ """
+ cmn_specgram = F.sliding_window_cmn(specgram, self.cmn_window, self.min_cmn_window, self.center, self.norm_vars)
+ return cmn_specgram
+
+
+class Vad(torch.nn.Module):
+ r"""Voice Activity Detector. Similar to SoX implementation.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: TorchScript
+
+ Attempts to trim silence and quiet background sounds from the ends of recordings of speech.
+ The algorithm currently uses a simple cepstral power measurement to detect voice,
+ so may be fooled by other things, especially music.
+
+ The effect can trim only from the front of the audio,
+ so in order to trim from the back, the reverse effect must also be used.
+
+ Args:
+ sample_rate (int): Sample rate of audio signal.
+ trigger_level (float, optional): The measurement level used to trigger activity detection.
+ This may need to be changed depending on the noise level, signal level,
+ and other characteristics of the input audio. (Default: 7.0)
+ trigger_time (float, optional): The time constant (in seconds)
+ used to help ignore short bursts of sound. (Default: 0.25)
+ search_time (float, optional): The amount of audio (in seconds)
+ to search for quieter/shorter bursts of audio to include prior
+ to the detected trigger point. (Default: 1.0)
+ allowed_gap (float, optional): The allowed gap (in seconds) between
+ quiteter/shorter bursts of audio to include prior
+ to the detected trigger point. (Default: 0.25)
+ pre_trigger_time (float, optional): The amount of audio (in seconds) to preserve
+ before the trigger point and any found quieter/shorter bursts. (Default: 0.0)
+ boot_time (float, optional) The algorithm (internally) uses adaptive noise
+ estimation/reduction in order to detect the start of the wanted audio.
+ This option sets the time for the initial noise estimate. (Default: 0.35)
+ noise_up_time (float, optional) Time constant used by the adaptive noise estimator
+ for when the noise level is increasing. (Default: 0.1)
+ noise_down_time (float, optional) Time constant used by the adaptive noise estimator
+ for when the noise level is decreasing. (Default: 0.01)
+ noise_reduction_amount (float, optional) Amount of noise reduction to use in
+ the detection algorithm (e.g. 0, 0.5, ...). (Default: 1.35)
+ measure_freq (float, optional) Frequency of the algorithm’s
+ processing/measurements. (Default: 20.0)
+ measure_duration: (float or None, optional) Measurement duration.
+ (Default: Twice the measurement period; i.e. with overlap.)
+ measure_smooth_time (float, optional) Time constant used to smooth
+ spectral measurements. (Default: 0.4)
+ hp_filter_freq (float, optional) "Brick-wall" frequency of high-pass filter applied
+ at the input to the detector algorithm. (Default: 50.0)
+ lp_filter_freq (float, optional) "Brick-wall" frequency of low-pass filter applied
+ at the input to the detector algorithm. (Default: 6000.0)
+ hp_lifter_freq (float, optional) "Brick-wall" frequency of high-pass lifter used
+ in the detector algorithm. (Default: 150.0)
+ lp_lifter_freq (float, optional) "Brick-wall" frequency of low-pass lifter used
+ in the detector algorithm. (Default: 2000.0)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> waveform_reversed, sample_rate = apply_effects_tensor(waveform, sample_rate, [["reverse"]])
+ >>> transform = transforms.Vad(sample_rate=sample_rate, trigger_level=7.5)
+ >>> waveform_reversed_front_trim = transform(waveform_reversed)
+ >>> waveform_end_trim, sample_rate = apply_effects_tensor(
+ >>> waveform_reversed_front_trim, sample_rate, [["reverse"]]
+ >>> )
+
+ Reference:
+ - http://sox.sourceforge.net/sox.html
+ """
+
+ def __init__(
+ self,
+ sample_rate: int,
+ trigger_level: float = 7.0,
+ trigger_time: float = 0.25,
+ search_time: float = 1.0,
+ allowed_gap: float = 0.25,
+ pre_trigger_time: float = 0.0,
+ boot_time: float = 0.35,
+ noise_up_time: float = 0.1,
+ noise_down_time: float = 0.01,
+ noise_reduction_amount: float = 1.35,
+ measure_freq: float = 20.0,
+ measure_duration: Optional[float] = None,
+ measure_smooth_time: float = 0.4,
+ hp_filter_freq: float = 50.0,
+ lp_filter_freq: float = 6000.0,
+ hp_lifter_freq: float = 150.0,
+ lp_lifter_freq: float = 2000.0,
+ ) -> None:
+ super().__init__()
+
+ self.sample_rate = sample_rate
+ self.trigger_level = trigger_level
+ self.trigger_time = trigger_time
+ self.search_time = search_time
+ self.allowed_gap = allowed_gap
+ self.pre_trigger_time = pre_trigger_time
+ self.boot_time = boot_time
+ self.noise_up_time = noise_up_time
+ self.noise_down_time = noise_down_time
+ self.noise_reduction_amount = noise_reduction_amount
+ self.measure_freq = measure_freq
+ self.measure_duration = measure_duration
+ self.measure_smooth_time = measure_smooth_time
+ self.hp_filter_freq = hp_filter_freq
+ self.lp_filter_freq = lp_filter_freq
+ self.hp_lifter_freq = hp_lifter_freq
+ self.lp_lifter_freq = lp_lifter_freq
+
+ def forward(self, waveform: Tensor) -> Tensor:
+ r"""
+ Args:
+ waveform (Tensor): Tensor of audio of dimension `(channels, time)` or `(time)`
+ Tensor of shape `(channels, time)` is treated as a multi-channel recording
+ of the same event and the resulting output will be trimmed to the earliest
+ voice activity in any channel.
+ """
+ return F.vad(
+ waveform=waveform,
+ sample_rate=self.sample_rate,
+ trigger_level=self.trigger_level,
+ trigger_time=self.trigger_time,
+ search_time=self.search_time,
+ allowed_gap=self.allowed_gap,
+ pre_trigger_time=self.pre_trigger_time,
+ boot_time=self.boot_time,
+ noise_up_time=self.noise_up_time,
+ noise_down_time=self.noise_down_time,
+ noise_reduction_amount=self.noise_reduction_amount,
+ measure_freq=self.measure_freq,
+ measure_duration=self.measure_duration,
+ measure_smooth_time=self.measure_smooth_time,
+ hp_filter_freq=self.hp_filter_freq,
+ lp_filter_freq=self.lp_filter_freq,
+ hp_lifter_freq=self.hp_lifter_freq,
+ lp_lifter_freq=self.lp_lifter_freq,
+ )
+
+
+class SpectralCentroid(torch.nn.Module):
+ r"""Compute the spectral centroid for each channel along the time axis.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ The spectral centroid is defined as the weighted average of the
+ frequency values, weighted by their magnitude.
+
+ Args:
+ sample_rate (int): Sample rate of audio signal.
+ n_fft (int, optional): Size of FFT, creates ``n_fft // 2 + 1`` bins. (Default: ``400``)
+ win_length (int or None, optional): Window size. (Default: ``n_fft``)
+ hop_length (int or None, optional): Length of hop between STFT windows. (Default: ``win_length // 2``)
+ pad (int, optional): Two sided padding of signal. (Default: ``0``)
+ window_fn (Callable[..., Tensor], optional): A function to create a window tensor
+ that is applied/multiplied to each frame/window. (Default: ``torch.hann_window``)
+ wkwargs (dict or None, optional): Arguments for window function. (Default: ``None``)
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = transforms.SpectralCentroid(sample_rate)
+ >>> spectral_centroid = transform(waveform) # (channel, time)
+ """
+ __constants__ = ["sample_rate", "n_fft", "win_length", "hop_length", "pad"]
+
+ def __init__(
+ self,
+ sample_rate: int,
+ n_fft: int = 400,
+ win_length: Optional[int] = None,
+ hop_length: Optional[int] = None,
+ pad: int = 0,
+ window_fn: Callable[..., Tensor] = torch.hann_window,
+ wkwargs: Optional[dict] = None,
+ ) -> None:
+ super(SpectralCentroid, self).__init__()
+ self.sample_rate = sample_rate
+ self.n_fft = n_fft
+ self.win_length = win_length if win_length is not None else n_fft
+ self.hop_length = hop_length if hop_length is not None else self.win_length // 2
+ window = window_fn(self.win_length) if wkwargs is None else window_fn(self.win_length, **wkwargs)
+ self.register_buffer("window", window)
+ self.pad = pad
+
+ def forward(self, waveform: Tensor) -> Tensor:
+ r"""
+ Args:
+ waveform (Tensor): Tensor of audio of dimension `(..., time)`.
+
+ Returns:
+ Tensor: Spectral Centroid of size `(..., time)`.
+ """
+
+ return F.spectral_centroid(
+ waveform, self.sample_rate, self.pad, self.window, self.n_fft, self.hop_length, self.win_length
+ )
+
+
+class PitchShift(LazyModuleMixin, torch.nn.Module):
+ r"""Shift the pitch of a waveform by ``n_steps`` steps.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: TorchScript
+
+ Args:
+ waveform (Tensor): The input waveform of shape `(..., time)`.
+ sample_rate (int): Sample rate of `waveform`.
+ n_steps (int): The (fractional) steps to shift `waveform`.
+ bins_per_octave (int, optional): The number of steps per octave (Default : ``12``).
+ n_fft (int, optional): Size of FFT, creates ``n_fft // 2 + 1`` bins (Default: ``512``).
+ win_length (int or None, optional): Window size. If None, then ``n_fft`` is used. (Default: ``None``).
+ hop_length (int or None, optional): Length of hop between STFT windows. If None, then ``win_length // 4``
+ is used (Default: ``None``).
+ window (Tensor or None, optional): Window tensor that is applied/multiplied to each frame/window.
+ If None, then ``torch.hann_window(win_length)`` is used (Default: ``None``).
+
+ Example
+ >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True)
+ >>> transform = transforms.PitchShift(sample_rate, 4)
+ >>> waveform_shift = transform(waveform) # (channel, time)
+ """
+ __constants__ = ["sample_rate", "n_steps", "bins_per_octave", "n_fft", "win_length", "hop_length"]
+
+ kernel: UninitializedParameter
+ width: int
+
+ def __init__(
+ self,
+ sample_rate: int,
+ n_steps: int,
+ bins_per_octave: int = 12,
+ n_fft: int = 512,
+ win_length: Optional[int] = None,
+ hop_length: Optional[int] = None,
+ window_fn: Callable[..., Tensor] = torch.hann_window,
+ wkwargs: Optional[dict] = None,
+ ) -> None:
+ super().__init__()
+ self.n_steps = n_steps
+ self.bins_per_octave = bins_per_octave
+ self.sample_rate = sample_rate
+ self.n_fft = n_fft
+ self.win_length = win_length if win_length is not None else n_fft
+ self.hop_length = hop_length if hop_length is not None else self.win_length // 4
+ window = window_fn(self.win_length) if wkwargs is None else window_fn(self.win_length, **wkwargs)
+ self.register_buffer("window", window)
+ rate = 2.0 ** (-float(n_steps) / bins_per_octave)
+ self.orig_freq = int(sample_rate / rate)
+ self.gcd = math.gcd(int(self.orig_freq), int(sample_rate))
+
+ if self.orig_freq != sample_rate:
+ self.width = -1
+ self.kernel = UninitializedParameter(device=None, dtype=None)
+
+ def initialize_parameters(self, input):
+ if self.has_uninitialized_params():
+ if self.orig_freq != self.sample_rate:
+ with torch.no_grad():
+ kernel, self.width = _get_sinc_resample_kernel(
+ self.orig_freq,
+ self.sample_rate,
+ self.gcd,
+ dtype=input.dtype,
+ device=input.device,
+ )
+ self.kernel.materialize(kernel.shape)
+ self.kernel.copy_(kernel)
+
+ def forward(self, waveform: Tensor) -> Tensor:
+ r"""
+ Args:
+ waveform (Tensor): Tensor of audio of dimension `(..., time)`.
+
+ Returns:
+ Tensor: The pitch-shifted audio of shape `(..., time)`.
+ """
+ shape = waveform.size()
+
+ waveform_stretch = _stretch_waveform(
+ waveform,
+ self.n_steps,
+ self.bins_per_octave,
+ self.n_fft,
+ self.win_length,
+ self.hop_length,
+ self.window,
+ )
+
+ if self.orig_freq != self.sample_rate:
+ waveform_shift = _apply_sinc_resample_kernel(
+ waveform_stretch,
+ self.orig_freq,
+ self.sample_rate,
+ self.gcd,
+ self.kernel,
+ self.width,
+ )
+ else:
+ waveform_shift = waveform_stretch
+
+ return _fix_waveform_shape(
+ waveform_shift,
+ shape,
+ )
+
+
+class RNNTLoss(torch.nn.Module):
+ """Compute the RNN Transducer loss from *Sequence Transduction with Recurrent Neural Networks*
+ :cite:`graves2012sequence`.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ The RNN Transducer loss extends the CTC loss by defining a distribution over output
+ sequences of all lengths, and by jointly modelling both input-output and output-output
+ dependencies.
+
+ Args:
+ blank (int, optional): blank label (Default: ``-1``)
+ clamp (float, optional): clamp for gradients (Default: ``-1``)
+ reduction (string, optional): Specifies the reduction to apply to the output:
+ ``"none"`` | ``"mean"`` | ``"sum"``. (Default: ``"mean"``)
+ fused_log_softmax (bool): set to False if calling log_softmax outside of loss (Default: ``True``)
+
+ Example
+ >>> # Hypothetical values
+ >>> logits = torch.tensor([[[[0.1, 0.6, 0.1, 0.1, 0.1],
+ >>> [0.1, 0.1, 0.6, 0.1, 0.1],
+ >>> [0.1, 0.1, 0.2, 0.8, 0.1]],
+ >>> [[0.1, 0.6, 0.1, 0.1, 0.1],
+ >>> [0.1, 0.1, 0.2, 0.1, 0.1],
+ >>> [0.7, 0.1, 0.2, 0.1, 0.1]]]],
+ >>> dtype=torch.float32,
+ >>> requires_grad=True)
+ >>> targets = torch.tensor([[1, 2]], dtype=torch.int)
+ >>> logit_lengths = torch.tensor([2], dtype=torch.int)
+ >>> target_lengths = torch.tensor([2], dtype=torch.int)
+ >>> transform = transforms.RNNTLoss(blank=0)
+ >>> loss = transform(logits, targets, logit_lengths, target_lengths)
+ >>> loss.backward()
+ """
+
+ def __init__(
+ self,
+ blank: int = -1,
+ clamp: float = -1.0,
+ reduction: str = "mean",
+ fused_log_softmax: bool = True,
+ ):
+ super().__init__()
+ self.blank = blank
+ self.clamp = clamp
+ self.reduction = reduction
+ self.fused_log_softmax = fused_log_softmax
+
+ def forward(
+ self,
+ logits: Tensor,
+ targets: Tensor,
+ logit_lengths: Tensor,
+ target_lengths: Tensor,
+ ):
+ """
+ Args:
+ logits (Tensor): Tensor of dimension `(batch, max seq length, max target length + 1, class)`
+ containing output from joiner
+ targets (Tensor): Tensor of dimension `(batch, max target length)` containing targets with zero padded
+ logit_lengths (Tensor): Tensor of dimension `(batch)` containing lengths of each sequence from encoder
+ target_lengths (Tensor): Tensor of dimension `(batch)` containing lengths of targets for each sequence
+ Returns:
+ Tensor: Loss with the reduction option applied. If ``reduction`` is ``"none"``, then size (batch),
+ otherwise scalar.
+ """
+ return F.rnnt_loss(
+ logits,
+ targets,
+ logit_lengths,
+ target_lengths,
+ self.blank,
+ self.clamp,
+ self.reduction,
+ self.fused_log_softmax,
+ )
+
+
+class Convolve(torch.nn.Module):
+ r"""
+ Convolves inputs along their last dimension using the direct method.
+ Note that, in contrast to :class:`torch.nn.Conv1d`, which actually applies the valid cross-correlation
+ operator, this module applies the true `convolution`_ operator.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ mode (str, optional): Must be one of ("full", "valid", "same").
+
+ * "full": Returns the full convolution result, with shape `(..., N + M - 1)`, where
+ `N` and `M` are the trailing dimensions of the two inputs. (Default)
+ * "valid": Returns the segment of the full convolution result corresponding to where
+ the two inputs overlap completely, with shape `(..., max(N, M) - min(N, M) + 1)`.
+ * "same": Returns the center segment of the full convolution result, with shape `(..., N)`.
+
+ .. _convolution:
+ https://en.wikipedia.org/wiki/Convolution
+ """
+
+ def __init__(self, mode: str = "full") -> None:
+ _check_convolve_mode(mode)
+
+ super().__init__()
+ self.mode = mode
+
+ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+ r"""
+ Args:
+ x (torch.Tensor): First convolution operand, with shape `(..., N)`.
+ y (torch.Tensor): Second convolution operand, with shape `(..., M)`
+ (leading dimensions must be broadcast-able with those of ``x``).
+
+ Returns:
+ torch.Tensor: Result of convolving ``x`` and ``y``, with shape `(..., L)`, where
+ the leading dimensions match those of ``x`` and `L` is dictated by ``mode``.
+ """
+ return F.convolve(x, y, mode=self.mode)
+
+
+class FFTConvolve(torch.nn.Module):
+ r"""
+ Convolves inputs along their last dimension using FFT. For inputs with large last dimensions, this module
+ is generally much faster than :class:`Convolve`.
+ Note that, in contrast to :class:`torch.nn.Conv1d`, which actually applies the valid cross-correlation
+ operator, this module applies the true `convolution`_ operator.
+ Also note that this module can only output float tensors (int tensor inputs will be cast to float).
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ mode (str, optional): Must be one of ("full", "valid", "same").
+
+ * "full": Returns the full convolution result, with shape `(..., N + M - 1)`, where
+ `N` and `M` are the trailing dimensions of the two inputs. (Default)
+ * "valid": Returns the segment of the full convolution result corresponding to where
+ the two inputs overlap completely, with shape `(..., max(N, M) - min(N, M) + 1)`.
+ * "same": Returns the center segment of the full convolution result, with shape `(..., N)`.
+
+ .. _convolution:
+ https://en.wikipedia.org/wiki/Convolution
+ """
+
+ def __init__(self, mode: str = "full") -> None:
+ _check_convolve_mode(mode)
+
+ super().__init__()
+ self.mode = mode
+
+ def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+ r"""
+ Args:
+ x (torch.Tensor): First convolution operand, with shape `(..., N)`.
+ y (torch.Tensor): Second convolution operand, with shape `(..., M)`
+ (leading dimensions must be broadcast-able with those of ``x``).
+
+ Returns:
+ torch.Tensor: Result of convolving ``x`` and ``y``, with shape `(..., L)`, where
+ the leading dimensions match those of ``x`` and `L` is dictated by ``mode``.
+ """
+ return F.fftconvolve(x, y, mode=self.mode)
+
+
+def _source_target_sample_rate(orig_freq: int, speed: float) -> Tuple[int, int]:
+ source_sample_rate = int(speed * orig_freq)
+ target_sample_rate = int(orig_freq)
+ gcd = math.gcd(source_sample_rate, target_sample_rate)
+ return source_sample_rate // gcd, target_sample_rate // gcd
+
+
+class Speed(torch.nn.Module):
+ r"""Adjusts waveform speed.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ orig_freq (int): Original frequency of the signals in ``waveform``.
+ factor (float): Factor by which to adjust speed of input. Values greater than 1.0
+ compress ``waveform`` in time, whereas values less than 1.0 stretch ``waveform`` in time.
+ """
+
+ def __init__(self, orig_freq, factor) -> None:
+ super().__init__()
+
+ self.orig_freq = orig_freq
+ self.factor = factor
+
+ self.source_sample_rate, self.target_sample_rate = _source_target_sample_rate(orig_freq, factor)
+ self.resampler = Resample(orig_freq=self.source_sample_rate, new_freq=self.target_sample_rate)
+
+ def forward(self, waveform, lengths: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ r"""
+ Args:
+ waveform (torch.Tensor): Input signals, with shape `(..., time)`.
+ lengths (torch.Tensor or None, optional): Valid lengths of signals in ``waveform``, with shape `(...)`.
+ If ``None``, all elements in ``waveform`` are treated as valid. (Default: ``None``)
+
+ Returns:
+ (torch.Tensor, torch.Tensor or None):
+ torch.Tensor
+ Speed-adjusted waveform, with shape `(..., new_time).`
+ torch.Tensor or None
+ If ``lengths`` is not ``None``, valid lengths of signals in speed-adjusted waveform,
+ with shape `(...)`; otherwise, ``None``.
+ """
+
+ if lengths is None:
+ out_lengths = None
+ else:
+ out_lengths = torch.ceil(lengths * self.target_sample_rate / self.source_sample_rate).to(lengths.dtype)
+
+ return self.resampler(waveform), out_lengths
+
+
+class SpeedPerturbation(torch.nn.Module):
+ r"""Applies the speed perturbation augmentation introduced in
+ *Audio augmentation for speech recognition* :cite:`ko15_interspeech`. For a given input,
+ the module samples a speed-up factor from ``factors`` uniformly at random and adjusts
+ the speed of the input by that factor.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ orig_freq (int): Original frequency of the signals in ``waveform``.
+ factors (Sequence[float]): Factors by which to adjust speed of input. Values greater than 1.0
+ compress ``waveform`` in time, whereas values less than 1.0 stretch ``waveform`` in time.
+
+ Example
+ >>> speed_perturb = SpeedPerturbation(16000, [0.9, 1.1, 1.0, 1.0, 1.0])
+ >>> # waveform speed will be adjusted by factor 0.9 with 20% probability,
+ >>> # 1.1 with 20% probability, and 1.0 (i.e. kept the same) with 60% probability.
+ >>> speed_perturbed_waveform = speed_perturb(waveform, lengths)
+ """
+
+ def __init__(self, orig_freq: int, factors: Sequence[float]) -> None:
+ super().__init__()
+
+ self.speeders = torch.nn.ModuleList([Speed(orig_freq=orig_freq, factor=factor) for factor in factors])
+
+ def forward(
+ self, waveform: torch.Tensor, lengths: Optional[torch.Tensor] = None
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ r"""
+ Args:
+ waveform (torch.Tensor): Input signals, with shape `(..., time)`.
+ lengths (torch.Tensor or None, optional): Valid lengths of signals in ``waveform``, with shape `(...)`.
+ If ``None``, all elements in ``waveform`` are treated as valid. (Default: ``None``)
+
+ Returns:
+ (torch.Tensor, torch.Tensor or None):
+ torch.Tensor
+ Speed-adjusted waveform, with shape `(..., new_time).`
+ torch.Tensor or None
+ If ``lengths`` is not ``None``, valid lengths of signals in speed-adjusted waveform,
+ with shape `(...)`; otherwise, ``None``.
+ """
+
+ idx = int(torch.randint(len(self.speeders), ()))
+ # NOTE: we do this because TorchScript doesn't allow for
+ # indexing ModuleList instances with non-literals.
+ for speeder_idx, speeder in enumerate(self.speeders):
+ if idx == speeder_idx:
+ return speeder(waveform, lengths)
+ raise RuntimeError("Speeder not found; execution should have never reached here.")
+
+
+class AddNoise(torch.nn.Module):
+ r"""Scales and adds noise to waveform per signal-to-noise ratio.
+ See :meth:`torchaudio.functional.add_noise` for more details.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+ """
+
+ def forward(
+ self, waveform: torch.Tensor, noise: torch.Tensor, snr: torch.Tensor, lengths: Optional[torch.Tensor] = None
+ ) -> torch.Tensor:
+ r"""
+ Args:
+ waveform (torch.Tensor): Input waveform, with shape `(..., L)`.
+ noise (torch.Tensor): Noise, with shape `(..., L)` (same shape as ``waveform``).
+ snr (torch.Tensor): Signal-to-noise ratios in dB, with shape `(...,)`.
+ lengths (torch.Tensor or None, optional): Valid lengths of signals in ``waveform`` and ``noise``,
+ with shape `(...,)` (leading dimensions must match those of ``waveform``). If ``None``, all
+ elements in ``waveform`` and ``noise`` are treated as valid. (Default: ``None``)
+
+ Returns:
+ torch.Tensor: Result of scaling and adding ``noise`` to ``waveform``, with shape `(..., L)`
+ (same shape as ``waveform``).
+ """
+ return F.add_noise(waveform, noise, snr, lengths)
+
+
+class Preemphasis(torch.nn.Module):
+ r"""Pre-emphasizes a waveform along its last dimension.
+ See :meth:`torchaudio.functional.preemphasis` for more details.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ coeff (float, optional): Pre-emphasis coefficient. Typically between 0.0 and 1.0.
+ (Default: 0.97)
+ """
+
+ def __init__(self, coeff: float = 0.97) -> None:
+ super().__init__()
+ self.coeff = coeff
+
+ def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+ r"""
+ Args:
+ waveform (torch.Tensor): Waveform, with shape `(..., N)`.
+
+ Returns:
+ torch.Tensor: Pre-emphasized waveform, with shape `(..., N)`.
+ """
+ return F.preemphasis(waveform, coeff=self.coeff)
+
+
+class Deemphasis(torch.nn.Module):
+ r"""De-emphasizes a waveform along its last dimension.
+ See :meth:`torchaudio.functional.deemphasis` for more details.
+
+ .. devices:: CPU CUDA
+
+ .. properties:: Autograd TorchScript
+
+ Args:
+ coeff (float, optional): De-emphasis coefficient. Typically between 0.0 and 1.0.
+ (Default: 0.97)
+ """
+
+ def __init__(self, coeff: float = 0.97) -> None:
+ super().__init__()
+ self.coeff = coeff
+
+ def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+ r"""
+ Args:
+ waveform (torch.Tensor): Waveform, with shape `(..., N)`.
+
+ Returns:
+ torch.Tensor: De-emphasized waveform, with shape `(..., N)`.
+ """
+ return F.deemphasis(waveform, coeff=self.coeff)
diff --git a/lib/python3.12/site-packages/torchaudio/utils/__init__.py b/lib/python3.12/site-packages/torchaudio/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..89bffaa34d61fbb12cfafbe7287af0b92139b19c
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/utils/__init__.py
@@ -0,0 +1,11 @@
+from torio.utils import ffmpeg_utils
+
+from . import sox_utils
+from .download import download_asset
+
+
+__all__ = [
+ "download_asset",
+ "sox_utils",
+ "ffmpeg_utils",
+]
diff --git a/lib/python3.12/site-packages/torchaudio/utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/utils/__pycache__/__init__.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ae2fcebffaa2cfe6c49d8ebc5caabd52a2da8f2c
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/utils/__pycache__/__init__.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/utils/__pycache__/download.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/utils/__pycache__/download.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..471f3bb980a36e3e6e03dd211e3872a64aa1b411
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/utils/__pycache__/download.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/utils/__pycache__/ffmpeg_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/utils/__pycache__/ffmpeg_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f4d0ed62998437f305dbfcb1e541354f2753ade4
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/utils/__pycache__/ffmpeg_utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/utils/__pycache__/sox_utils.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/utils/__pycache__/sox_utils.cpython-312.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cdfef3854479d785044b335ed863873b7fef4157
Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/utils/__pycache__/sox_utils.cpython-312.pyc differ
diff --git a/lib/python3.12/site-packages/torchaudio/utils/download.py b/lib/python3.12/site-packages/torchaudio/utils/download.py
new file mode 100644
index 0000000000000000000000000000000000000000..2081877d15a13e91a6fcb87905634addd23cc712
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/utils/download.py
@@ -0,0 +1,89 @@
+import hashlib
+import logging
+from os import PathLike
+from pathlib import Path
+from typing import Union
+
+import torch
+from torchaudio._internal import download_url_to_file
+
+_LG = logging.getLogger(__name__)
+
+
+def _get_local_path(key):
+ path = Path(torch.hub.get_dir()) / "torchaudio" / Path(key)
+ path.parent.mkdir(parents=True, exist_ok=True)
+ return path
+
+
+def _download(key, path, progress):
+ url = f"https://download.pytorch.org/torchaudio/{key}"
+ download_url_to_file(url, path, progress=progress)
+
+
+def _get_hash(path, hash, chunk_size=1028):
+ m = hashlib.sha256()
+ with open(path, "rb") as file:
+ data = file.read(chunk_size)
+ while data:
+ m.update(data)
+ data = file.read(chunk_size)
+ return m.hexdigest()
+
+
+def download_asset(
+ key: str,
+ hash: str = "",
+ path: Union[str, PathLike] = "",
+ *,
+ progress: bool = True,
+) -> str:
+ """Download and store torchaudio assets to local file system.
+
+ If a file exists at the download path, then that path is returned with or without
+ hash validation.
+
+ Args:
+ key (str): The asset identifier.
+ hash (str, optional):
+ The value of SHA256 hash of the asset. If provided, it is used to verify
+ the downloaded / cached object. If not provided, then no hash validation
+ is performed. This means if a file exists at the download path, then the path
+ is returned as-is without verifying the identity of the file.
+ path (path-like object, optional):
+ By default, the downloaded asset is saved in a directory under
+ :py:func:`torch.hub.get_dir` and intermediate directories based on the given `key`
+ are created.
+ This argument can be used to overwrite the target location.
+ When this argument is provided, all the intermediate directories have to be
+ created beforehand.
+ progress (bool): Whether to show progress bar for downloading. Default: ``True``.
+
+ Note:
+ Currently the valid key values are the route on ``download.pytorch.org/torchaudio``,
+ but this is an implementation detail.
+
+ Returns:
+ str: The path to the asset on the local file system.
+ """
+ path = path or _get_local_path(key)
+
+ if path.exists():
+ _LG.info("The local file (%s) exists. Skipping the download.", path)
+ else:
+ _LG.info("Downloading %s to %s", key, path)
+ _download(key, path, progress=progress)
+
+ if hash:
+ _LG.info("Verifying the hash value.")
+ digest = _get_hash(path, hash)
+
+ if digest != hash:
+ raise ValueError(
+ f"The hash value of the downloaded file ({path}), '{digest}' does not match "
+ f"the provided hash value, '{hash}'."
+ )
+
+ _LG.info("Hash validated.")
+
+ return str(path)
diff --git a/lib/python3.12/site-packages/torchaudio/utils/ffmpeg_utils.py b/lib/python3.12/site-packages/torchaudio/utils/ffmpeg_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..385596edc1491e45ecd4aae14a07b2c0e64ecd22
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/utils/ffmpeg_utils.py
@@ -0,0 +1,11 @@
+"""Module to change the configuration of FFmpeg libraries (such as libavformat).
+
+It affects functionalities in :py:mod:`torchaudio.io` (and indirectly :py:func:`torchaudio.load`).
+"""
+
+
+# This file is just for BC.
+def __getattr__(item):
+ from torio.utils import ffmpeg_utils
+
+ return getattr(ffmpeg_utils, item)
diff --git a/lib/python3.12/site-packages/torchaudio/utils/sox_utils.py b/lib/python3.12/site-packages/torchaudio/utils/sox_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..5212b77ea9d5ae0e58741322db7c9852a4ddafff
--- /dev/null
+++ b/lib/python3.12/site-packages/torchaudio/utils/sox_utils.py
@@ -0,0 +1,99 @@
+"""Module to change the configuration of libsox, which is used by I/O functions like
+:py:mod:`~torchaudio.backend.sox_io_backend` and :py:mod:`~torchaudio.sox_effects`.
+"""
+
+from typing import Dict, List
+
+import torchaudio
+
+sox_ext = torchaudio._extension.lazy_import_sox_ext()
+
+
+def set_seed(seed: int):
+ """Set libsox's PRNG
+
+ Args:
+ seed (int): seed value. valid range is int32.
+
+ See Also:
+ http://sox.sourceforge.net/sox.html
+ """
+ sox_ext.set_seed(seed)
+
+
+def set_verbosity(verbosity: int):
+ """Set libsox's verbosity
+
+ Args:
+ verbosity (int): Set verbosity level of libsox.
+
+ * ``1`` failure messages
+ * ``2`` warnings
+ * ``3`` details of processing
+ * ``4``-``6`` increasing levels of debug messages
+
+ See Also:
+ http://sox.sourceforge.net/sox.html
+ """
+ sox_ext.set_verbosity(verbosity)
+
+
+def set_buffer_size(buffer_size: int):
+ """Set buffer size for sox effect chain
+
+ Args:
+ buffer_size (int): Set the size in bytes of the buffers used for processing audio.
+
+ See Also:
+ http://sox.sourceforge.net/sox.html
+ """
+ sox_ext.set_buffer_size(buffer_size)
+
+
+def set_use_threads(use_threads: bool):
+ """Set multithread option for sox effect chain
+
+ Args:
+ use_threads (bool): When ``True``, enables ``libsox``'s parallel effects channels processing.
+ To use mutlithread, the underlying ``libsox`` has to be compiled with OpenMP support.
+
+ See Also:
+ http://sox.sourceforge.net/sox.html
+ """
+ sox_ext.set_use_threads(use_threads)
+
+
+def list_effects() -> Dict[str, str]:
+ """List the available sox effect names
+
+ Returns:
+ Dict[str, str]: Mapping from ``effect name`` to ``usage``
+ """
+ return dict(sox_ext.list_effects())
+
+
+def list_read_formats() -> List[str]:
+ """List the supported audio formats for read
+
+ Returns:
+ List[str]: List of supported audio formats
+ """
+ return sox_ext.list_read_formats()
+
+
+def list_write_formats() -> List[str]:
+ """List the supported audio formats for write
+
+ Returns:
+ List[str]: List of supported audio formats
+ """
+ return sox_ext.list_write_formats()
+
+
+def get_buffer_size() -> int:
+ """Get buffer size for sox effect chain
+
+ Returns:
+ int: size in bytes of buffers used for processing audio.
+ """
+ return sox_ext.get_buffer_size()