Skip to content

GitLab

Commit 5859923a authored by Joao Gomes's avatar Joao Gomes Committed by Facebook GitHub Bot
Browse files

Apply arc lint to pytorch audio (#2096) 

Summary:
Pull Request resolved: https://github.com/pytorch/audio/pull/2096

run: `arc lint --apply-patches --paths-cmd 'hg files -I "./**/*.py"'`

Reviewed By: mthrok

Differential Revision: D33297351

fbshipit-source-id: 7bf5956edf0717c5ca90219f72414ff4eeaf5aa8
parent 0e5913d5
Hide whitespace changes
Inline Side-by-side
Showing with 800 additions and 926 deletions
torchaudio/pipelines/_tts/__init__.py
View file @ 5859923a
from .interface import Tacotron2TTSBundle
from .impl import (
TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH,
TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH,
TACOTRON2_WAVERNN_CHAR_LJSPEECH,
TACOTRON2_WAVERNN_PHONE_LJSPEECH,
)
from .interface import Tacotron2TTSBundle
__all__ = [
'Tacotron2TTSBundle',
'TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH',
'TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH',
'TACOTRON2_WAVERNN_CHAR_LJSPEECH',
'TACOTRON2_WAVERNN_PHONE_LJSPEECH',
"Tacotron2TTSBundle",
"TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH",
"TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH",
"TACOTRON2_WAVERNN_CHAR_LJSPEECH",
"TACOTRON2_WAVERNN_PHONE_LJSPEECH",
]
torchaudio/pipelines/_tts/impl.py
View file @ 5859923a
from dataclasses import dataclass
import re
from dataclasses import dataclass
from typing import Union, Optional, Dict, Any, Tuple, List
import torch
from torch import Tensor
from torchaudio._internal import load_state_dict_from_url
from torchaudio.models import Tacotron2, WaveRNN
from torchaudio.functional import mu_law_decoding
from torchaudio.models import Tacotron2, WaveRNN
from torchaudio.transforms import InverseMelScale, GriffinLim
from . import utils
from .interface import Tacotron2TTSBundle
__all__ = []
_BASE_URL = 'https://download.pytorch.org/torchaudio/models'
_BASE_URL = "https://download.pytorch.org/torchaudio/models"
################################################################################
......@@ -44,8 +44,7 @@ class _EnglishPhoneProcessor(Tacotron2TTSBundle.TextProcessor):
super().__init__()
self._tokens = utils._get_phones()
self._mapping = {p: i for i, p in enumerate(self._tokens)}
self._phonemizer = utils._load_phonemizer(
'en_us_cmudict_forward.pt', dl_kwargs=dl_kwargs)
self._phonemizer = utils._load_phonemizer("en_us_cmudict_forward.pt", dl_kwargs=dl_kwargs)
self._pattern = r"(\[[A-Z]+?\]|[_!'(),.:;? -])"
@property
......@@ -57,9 +56,9 @@ class _EnglishPhoneProcessor(Tacotron2TTSBundle.TextProcessor):
texts = [texts]
indices = []
for phones in self._phonemizer(texts, lang='en_us'):
for phones in self._phonemizer(texts, lang="en_us"):
# '[F][UW][B][AA][R]!' -> ['F', 'UW', 'B', 'AA', 'R', '!']
ret = [re.sub(r'[\[\]]', '', r) for r in re.findall(self._pattern, phones)]
ret = [re.sub(r"[\[\]]", "", r) for r in re.findall(self._pattern, phones)]
indices.append([self._mapping[p] for p in ret])
return utils._to_tensor(indices)
......@@ -68,12 +67,9 @@ class _EnglishPhoneProcessor(Tacotron2TTSBundle.TextProcessor):
# Pipeline implementation - Vocoder
################################################################################
class _WaveRNNVocoder(torch.nn.Module, Tacotron2TTSBundle.Vocoder):
def __init__(
self,
model: WaveRNN,
min_level_db: Optional[float] = -100
):
def __init__(self, model: WaveRNN, min_level_db: Optional[float] = -100):
super().__init__()
self._sample_rate = 22050
self._model = model
......@@ -104,10 +100,10 @@ class _GriffinLimVocoder(torch.nn.Module, Tacotron2TTSBundle.Vocoder):
n_stft=(1024 // 2 + 1),
n_mels=80,
sample_rate=self.sample_rate,
f_min=0.,
f_max=8000.,
f_min=0.0,
f_max=8000.0,
mel_scale="slaney",
norm='slaney',
norm="slaney",
)
self._griffin_lim = GriffinLim(
n_fft=1024,
......@@ -151,7 +147,7 @@ class _Tacotron2Mixin:
def get_tacotron2(self, *, dl_kwargs=None) -> Tacotron2:
model = Tacotron2(**self._tacotron2_params)
url = f'{_BASE_URL}/{self._tacotron2_path}'
url = f"{_BASE_URL}/{self._tacotron2_path}"
dl_kwargs = {} if dl_kwargs is None else dl_kwargs
state_dict = load_state_dict_from_url(url, **dl_kwargs)
model.load_state_dict(state_dict)
......@@ -170,7 +166,7 @@ class _WaveRNNMixin:
def _get_wavernn(self, *, dl_kwargs=None):
model = WaveRNN(**self._wavernn_params)
url = f'{_BASE_URL}/{self._wavernn_path}'
url = f"{_BASE_URL}/{self._wavernn_path}"
dl_kwargs = {} if dl_kwargs is None else dl_kwargs
state_dict = load_state_dict_from_url(url, **dl_kwargs)
model.load_state_dict(state_dict)
......@@ -214,11 +210,10 @@ class _Tacotron2GriffinLimPhoneBundle(_GriffinLimMixin, _Tacotron2Mixin, _PhoneM
TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH = _Tacotron2GriffinLimCharBundle(
_tacotron2_path='tacotron2_english_characters_1500_epochs_ljspeech.pth',
_tacotron2_path="tacotron2_english_characters_1500_epochs_ljspeech.pth",
_tacotron2_params=utils._get_taco_params(n_symbols=38),
)
TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH.__doc__ = (
'''Character-based TTS pipeline with :py:class:`torchaudio.models.Tacotron2` and
TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH.__doc__ = """Character-based TTS pipeline with :py:class:`torchaudio.models.Tacotron2` and
:py:class:`torchaudio.transforms.GriffinLim`.
The text processor encodes the input texts character-by-character.
......@@ -254,14 +249,13 @@ Example - "The examination and testimony of the experts enabled the Commission t
<source src="https://download.pytorch.org/torchaudio/doc-assets/TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH_v2.wav" type="audio/wav">
Your browser does not support the <code>audio</code> element.
</audio>
''') # noqa: E501
""" # noqa: E501
TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH = _Tacotron2GriffinLimPhoneBundle(
_tacotron2_path='tacotron2_english_phonemes_1500_epochs_ljspeech.pth',
_tacotron2_path="tacotron2_english_phonemes_1500_epochs_ljspeech.pth",
_tacotron2_params=utils._get_taco_params(n_symbols=96),
)
TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH.__doc__ = (
'''Phoneme-based TTS pipeline with :py:class:`torchaudio.models.Tacotron2` and
TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH.__doc__ = """Phoneme-based TTS pipeline with :py:class:`torchaudio.models.Tacotron2` and
:py:class:`torchaudio.transforms.GriffinLim`.
The text processor encodes the input texts based on phoneme.
......@@ -302,16 +296,15 @@ Example - "The examination and testimony of the experts enabled the Commission t
Your browser does not support the <code>audio</code> element.
</audio>
''') # noqa: E501
""" # noqa: E501
TACOTRON2_WAVERNN_CHAR_LJSPEECH = _Tacotron2WaveRNNCharBundle(
_tacotron2_path='tacotron2_english_characters_1500_epochs_wavernn_ljspeech.pth',
_tacotron2_path="tacotron2_english_characters_1500_epochs_wavernn_ljspeech.pth",
_tacotron2_params=utils._get_taco_params(n_symbols=38),
_wavernn_path='wavernn_10k_epochs_8bits_ljspeech.pth',
_wavernn_path="wavernn_10k_epochs_8bits_ljspeech.pth",
_wavernn_params=utils._get_wrnn_params(),
)
TACOTRON2_WAVERNN_CHAR_LJSPEECH.__doc__ = (
'''Character-based TTS pipeline with :py:class:`torchaudio.models.Tacotron2` and
TACOTRON2_WAVERNN_CHAR_LJSPEECH.__doc__ = """Character-based TTS pipeline with :py:class:`torchaudio.models.Tacotron2` and
:py:class:`torchaudio.models.WaveRNN`.
The text processor encodes the input texts character-by-character.
......@@ -350,16 +343,15 @@ Example - "The examination and testimony of the experts enabled the Commission t
<source src="https://download.pytorch.org/torchaudio/doc-assets/TACOTRON2_WAVERNN_CHAR_LJSPEECH_v2.wav" type="audio/wav">
Your browser does not support the <code>audio</code> element.
</audio>
''') # noqa: E501
""" # noqa: E501
TACOTRON2_WAVERNN_PHONE_LJSPEECH = _Tacotron2WaveRNNPhoneBundle(
_tacotron2_path='tacotron2_english_phonemes_1500_epochs_wavernn_ljspeech.pth',
_tacotron2_path="tacotron2_english_phonemes_1500_epochs_wavernn_ljspeech.pth",
_tacotron2_params=utils._get_taco_params(n_symbols=96),
_wavernn_path='wavernn_10k_epochs_8bits_ljspeech.pth',
_wavernn_path="wavernn_10k_epochs_8bits_ljspeech.pth",
_wavernn_params=utils._get_wrnn_params(),
)
TACOTRON2_WAVERNN_PHONE_LJSPEECH.__doc__ = (
'''Phoneme-based TTS pipeline with :py:class:`torchaudio.models.Tacotron2` and
TACOTRON2_WAVERNN_PHONE_LJSPEECH.__doc__ = """Phoneme-based TTS pipeline with :py:class:`torchaudio.models.Tacotron2` and
:py:class:`torchaudio.models.WaveRNN`.
The text processor encodes the input texts based on phoneme.
......@@ -403,4 +395,4 @@ Example - "The examination and testimony of the experts enabled the Commission t
<source src="https://download.pytorch.org/torchaudio/doc-assets/TACOTRON2_WAVERNN_PHONE_LJSPEECH_v2.wav" type="audio/wav">
Your browser does not support the <code>audio</code> element.
</audio>
''') # noqa: E501
""" # noqa: E501
torchaudio/pipelines/_tts/interface.py
View file @ 5859923a
......@@ -2,7 +2,6 @@ from abc import ABC, abstractmethod
from typing import Union, List, Tuple, Optional
from torch import Tensor
from torchaudio.models import Tacotron2
......
torchaudio/pipelines/_tts/utils.py
View file @ 5859923a
import os
import logging
import os
import torch
from torchaudio._internal import (
download_url_to_file,
module_utils as _mod_utils,
......@@ -11,44 +10,44 @@ from torchaudio._internal import (
def _get_chars():
return (
'_',
'-',
'!',
"_",
"-",
"!",
"'",
'(',
')',
',',
'.',
':',
';',
'?',
' ',
'a',
'b',
'c',
'd',
'e',
'f',
'g',
'h',
'i',
'j',
'k',
'l',
'm',
'n',
'o',
'p',
'q',
'r',
's',
't',
'u',
'v',
'w',
'x',
'y',
'z',
"(",
")",
",",
".",
":",
";",
"?",
" ",
"a",
"b",
"c",
"d",
"e",
"f",
"g",
"h",
"i",
"j",
"k",
"l",
"m",
"n",
"o",
"p",
"q",
"r",
"s",
"t",
"u",
"v",
"w",
"x",
"y",
"z",
)
......@@ -149,7 +148,7 @@ def _get_phones():
"W",
"Y",
"Z",
"ZH"
"ZH",
)
......@@ -161,18 +160,18 @@ def _to_tensor(indices):
def _load_phonemizer(file, dl_kwargs):
if not _mod_utils.is_module_available('dp'):
raise RuntimeError('DeepPhonemizer is not installed. Please install it.')
if not _mod_utils.is_module_available("dp"):
raise RuntimeError("DeepPhonemizer is not installed. Please install it.")
from dp.phonemizer import Phonemizer
# By default, dp issues DEBUG level log.
logger = logging.getLogger('dp')
logger = logging.getLogger("dp")
orig_level = logger.level
logger.setLevel(logging.INFO)
try:
url = f'https://public-asai-dl-models.s3.eu-central-1.amazonaws.com/DeepPhonemizer/{file}'
directory = os.path.join(torch.hub.get_dir(), 'checkpoints')
url = f"https://public-asai-dl-models.s3.eu-central-1.amazonaws.com/DeepPhonemizer/{file}"
directory = os.path.join(torch.hub.get_dir(), "checkpoints")
os.makedirs(directory, exist_ok=True)
path = os.path.join(directory, file)
if not os.path.exists(path):
......@@ -192,41 +191,41 @@ def _unnormalize_waveform(waveform: torch.Tensor, bits: int) -> torch.Tensor:
def _get_taco_params(n_symbols):
return {
'mask_padding': False,
'n_mels': 80,
'n_frames_per_step': 1,
'symbol_embedding_dim': 512,
'encoder_embedding_dim': 512,
'encoder_n_convolution': 3,
'encoder_kernel_size': 5,
'decoder_rnn_dim': 1024,
'decoder_max_step': 2000,
'decoder_dropout': 0.1,
'decoder_early_stopping': True,
'attention_rnn_dim': 1024,
'attention_hidden_dim': 128,
'attention_location_n_filter': 32,
'attention_location_kernel_size': 31,
'attention_dropout': 0.1,
'prenet_dim': 256,
'postnet_n_convolution': 5,
'postnet_kernel_size': 5,
'postnet_embedding_dim': 512,
'gate_threshold': 0.5,
'n_symbol': n_symbols,
"mask_padding": False,
"n_mels": 80,
"n_frames_per_step": 1,
"symbol_embedding_dim": 512,
"encoder_embedding_dim": 512,
"encoder_n_convolution": 3,
"encoder_kernel_size": 5,
"decoder_rnn_dim": 1024,
"decoder_max_step": 2000,
"decoder_dropout": 0.1,
"decoder_early_stopping": True,
"attention_rnn_dim": 1024,
"attention_hidden_dim": 128,
"attention_location_n_filter": 32,
"attention_location_kernel_size": 31,
"attention_dropout": 0.1,
"prenet_dim": 256,
"postnet_n_convolution": 5,
"postnet_kernel_size": 5,
"postnet_embedding_dim": 512,
"gate_threshold": 0.5,
"n_symbol": n_symbols,
}
def _get_wrnn_params():
return {
'upsample_scales': [5, 5, 11],
'n_classes': 2 ** 8, # n_bits = 8
'hop_length': 275,
'n_res_block': 10,
'n_rnn': 512,
'n_fc': 512,
'kernel_size': 5,
'n_freq': 80,
'n_hidden': 128,
'n_output': 128
"upsample_scales": [5, 5, 11],
"n_classes": 2 ** 8, # n_bits = 8
"hop_length": 275,
"n_res_block": 10,
"n_rnn": 512,
"n_fc": 512,
"kernel_size": 5,
"n_freq": 80,
"n_hidden": 128,
"n_output": 128,
}
torchaudio/pipelines/_wav2vec2/impl.py
View file @ 5859923a
......@@ -2,9 +2,9 @@ from dataclasses import dataclass
from typing import Dict, Tuple, Any
import torch
from torchaudio._internal import load_state_dict_from_url
from torchaudio.models import wav2vec2_model, Wav2Vec2Model
from . import utils
......@@ -43,6 +43,7 @@ class Wav2Vec2Bundle:
>>> # Extract acoustic features
>>> features, _ = model.extract_features(waveform)
""" # noqa: E501
_path: str
_params: Dict[str, Any]
_sample_rate: float
......@@ -56,7 +57,7 @@ class Wav2Vec2Bundle:
return self._sample_rate
def _get_state_dict(self, dl_kwargs):
url = f'https://download.pytorch.org/torchaudio/models/{self._path}'
url = f"https://download.pytorch.org/torchaudio/models/{self._path}"
dl_kwargs = {} if dl_kwargs is None else dl_kwargs
state_dict = load_state_dict_from_url(url, **dl_kwargs)
return state_dict
......@@ -120,13 +121,14 @@ class Wav2Vec2ASRBundle(Wav2Vec2Bundle):
>>> # `ctc_decode` is for illustration purpose only
>>> transcripts = ctc_decode(emissions, labels)
""" # noqa: E501
_labels: Tuple[str]
_remove_aux_axis: Tuple[int] = (1, 2, 3)
def get_labels(
self,
*,
blank: str = '-',
self,
*,
blank: str = "-",
) -> Tuple[str]:
"""The output class labels (only applicable to fine-tuned bundles)
......@@ -161,17 +163,17 @@ class Wav2Vec2ASRBundle(Wav2Vec2Bundle):
# that resembles mistake.
# The label `1` shows up in the training dataset of German (1 out of 16M),
# English (1 / 28M), Spanish (1 / 9.4M), Romanian (1 / 4.7M) and Polish (6 / 5.8M)
for key in ['aux.weight', 'aux.bias']:
for key in ["aux.weight", "aux.bias"]:
t = state_dict[key]
state_dict[key] = torch.stack([t[i] for i in range(t.size(0)) if i not in self._remove_aux_axis])
return state_dict
WAV2VEC2_BASE = Wav2Vec2Bundle(
_path='wav2vec2_fairseq_base_ls960.pth',
_path="wav2vec2_fairseq_base_ls960.pth",
_params={
'extractor_mode': 'group_norm',
'extractor_conv_layer_config': [
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
(512, 10, 5),
(512, 3, 2),
(512, 3, 2),
......@@ -180,19 +182,19 @@ WAV2VEC2_BASE = Wav2Vec2Bundle(
(512, 2, 2),
(512, 2, 2),
],
'extractor_conv_bias': False,
'encoder_embed_dim': 768,
'encoder_projection_dropout': 0.1,
'encoder_pos_conv_kernel': 128,
'encoder_pos_conv_groups': 16,
'encoder_num_layers': 12,
'encoder_num_heads': 12,
'encoder_attention_dropout': 0.1,
'encoder_ff_interm_features': 3072,
'encoder_ff_interm_dropout': 0.0,
'encoder_dropout': 0.1,
'encoder_layer_norm_first': False,
'encoder_layer_drop': 0.05,
"extractor_conv_bias": False,
"encoder_embed_dim": 768,
"encoder_projection_dropout": 0.1,
"encoder_pos_conv_kernel": 128,
"encoder_pos_conv_groups": 16,
"encoder_num_layers": 12,
"encoder_num_heads": 12,
"encoder_attention_dropout": 0.1,
"encoder_ff_interm_features": 3072,
"encoder_ff_interm_dropout": 0.0,
"encoder_dropout": 0.1,
"encoder_layer_norm_first": False,
"encoder_layer_drop": 0.05,
"aux_num_out": None,
},
_sample_rate=16000,
......@@ -212,10 +214,10 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage.
""" # noqa: E501
WAV2VEC2_ASR_BASE_10M = Wav2Vec2ASRBundle(
_path='wav2vec2_fairseq_base_ls960_asr_ll10m.pth',
_path="wav2vec2_fairseq_base_ls960_asr_ll10m.pth",
_params={
'extractor_mode': 'group_norm',
'extractor_conv_layer_config': [
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
(512, 10, 5),
(512, 3, 2),
(512, 3, 2),
......@@ -224,19 +226,19 @@ WAV2VEC2_ASR_BASE_10M = Wav2Vec2ASRBundle(
(512, 2, 2),
(512, 2, 2),
],
'extractor_conv_bias': False,
'encoder_embed_dim': 768,
'encoder_projection_dropout': 0.1,
'encoder_pos_conv_kernel': 128,
'encoder_pos_conv_groups': 16,
'encoder_num_layers': 12,
'encoder_num_heads': 12,
'encoder_attention_dropout': 0.1,
'encoder_ff_interm_features': 3072,
'encoder_ff_interm_dropout': 0.0,
'encoder_dropout': 0.1,
'encoder_layer_norm_first': False,
'encoder_layer_drop': 0.05,
"extractor_conv_bias": False,
"encoder_embed_dim": 768,
"encoder_projection_dropout": 0.1,
"encoder_pos_conv_kernel": 128,
"encoder_pos_conv_groups": 16,
"encoder_num_layers": 12,
"encoder_num_heads": 12,
"encoder_attention_dropout": 0.1,
"encoder_ff_interm_features": 3072,
"encoder_ff_interm_dropout": 0.0,
"encoder_dropout": 0.1,
"encoder_layer_norm_first": False,
"encoder_layer_drop": 0.05,
"aux_num_out": 29,
},
_labels=utils._get_en_labels(),
......@@ -258,10 +260,10 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
""" # noqa: E501
WAV2VEC2_ASR_BASE_100H = Wav2Vec2ASRBundle(
'wav2vec2_fairseq_base_ls960_asr_ls100.pth',
"wav2vec2_fairseq_base_ls960_asr_ls100.pth",
{
'extractor_mode': 'group_norm',
'extractor_conv_layer_config': [
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
(512, 10, 5),
(512, 3, 2),
(512, 3, 2),
......@@ -270,19 +272,19 @@ WAV2VEC2_ASR_BASE_100H = Wav2Vec2ASRBundle(
(512, 2, 2),
(512, 2, 2),
],
'extractor_conv_bias': False,
'encoder_embed_dim': 768,
'encoder_projection_dropout': 0.1,
'encoder_pos_conv_kernel': 128,
'encoder_pos_conv_groups': 16,
'encoder_num_layers': 12,
'encoder_num_heads': 12,
'encoder_attention_dropout': 0.1,
'encoder_ff_interm_features': 3072,
'encoder_ff_interm_dropout': 0.0,
'encoder_dropout': 0.1,
'encoder_layer_norm_first': False,
'encoder_layer_drop': 0.05,
"extractor_conv_bias": False,
"encoder_embed_dim": 768,
"encoder_projection_dropout": 0.1,
"encoder_pos_conv_kernel": 128,
"encoder_pos_conv_groups": 16,
"encoder_num_layers": 12,
"encoder_num_heads": 12,
"encoder_attention_dropout": 0.1,
"encoder_ff_interm_features": 3072,
"encoder_ff_interm_dropout": 0.0,
"encoder_dropout": 0.1,
"encoder_layer_norm_first": False,
"encoder_layer_drop": 0.05,
"aux_num_out": 29,
},
_labels=utils._get_en_labels(),
......@@ -304,7 +306,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
""" # noqa: E501
WAV2VEC2_ASR_BASE_960H = Wav2Vec2ASRBundle(
'wav2vec2_fairseq_base_ls960_asr_ls960.pth',
"wav2vec2_fairseq_base_ls960_asr_ls960.pth",
{
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
......@@ -349,7 +351,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
""" # noqa: E501
WAV2VEC2_LARGE = Wav2Vec2Bundle(
'wav2vec2_fairseq_large_ls960.pth',
"wav2vec2_fairseq_large_ls960.pth",
{
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
......@@ -393,7 +395,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage.
""" # noqa: E501
WAV2VEC2_ASR_LARGE_10M = Wav2Vec2ASRBundle(
'wav2vec2_fairseq_large_ls960_asr_ll10m.pth',
"wav2vec2_fairseq_large_ls960_asr_ll10m.pth",
{
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
......@@ -439,7 +441,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
""" # noqa: E501
WAV2VEC2_ASR_LARGE_100H = Wav2Vec2ASRBundle(
'wav2vec2_fairseq_large_ls960_asr_ls100.pth',
"wav2vec2_fairseq_large_ls960_asr_ls100.pth",
{
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
......@@ -485,7 +487,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
""" # noqa: E501
WAV2VEC2_ASR_LARGE_960H = Wav2Vec2ASRBundle(
'wav2vec2_fairseq_large_ls960_asr_ls960.pth',
"wav2vec2_fairseq_large_ls960_asr_ls960.pth",
{
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
......@@ -530,7 +532,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
""" # noqa: E501
WAV2VEC2_LARGE_LV60K = Wav2Vec2Bundle(
'wav2vec2_fairseq_large_lv60k.pth',
"wav2vec2_fairseq_large_lv60k.pth",
{
"extractor_mode": "layer_norm",
"extractor_conv_layer_config": [
......@@ -574,7 +576,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage.
""" # noqa: E501
WAV2VEC2_ASR_LARGE_LV60K_10M = Wav2Vec2ASRBundle(
'wav2vec2_fairseq_large_lv60k_asr_ll10m.pth',
"wav2vec2_fairseq_large_lv60k_asr_ll10m.pth",
{
"extractor_mode": "layer_norm",
"extractor_conv_layer_config": [
......@@ -620,7 +622,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
""" # noqa: E501
WAV2VEC2_ASR_LARGE_LV60K_100H = Wav2Vec2ASRBundle(
'wav2vec2_fairseq_large_lv60k_asr_ls100.pth',
"wav2vec2_fairseq_large_lv60k_asr_ls100.pth",
{
"extractor_mode": "layer_norm",
"extractor_conv_layer_config": [
......@@ -666,7 +668,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
""" # noqa: E501
WAV2VEC2_ASR_LARGE_LV60K_960H = Wav2Vec2ASRBundle(
'wav2vec2_fairseq_large_lv60k_asr_ls960.pth',
"wav2vec2_fairseq_large_lv60k_asr_ls960.pth",
{
"extractor_mode": "layer_norm",
"extractor_conv_layer_config": [
......@@ -713,7 +715,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
""" # noqa: E501
WAV2VEC2_XLSR53 = Wav2Vec2Bundle(
'wav2vec2_fairseq_large_xlsr53.pth',
"wav2vec2_fairseq_large_xlsr53.pth",
{
"extractor_mode": "layer_norm",
"extractor_conv_layer_config": [
......@@ -760,10 +762,10 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage.
""" # noqa: E501
HUBERT_BASE = Wav2Vec2Bundle(
'hubert_fairseq_base_ls960.pth',
"hubert_fairseq_base_ls960.pth",
{
'extractor_mode': 'group_norm',
'extractor_conv_layer_config': [
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
(512, 10, 5),
(512, 3, 2),
(512, 3, 2),
......@@ -772,20 +774,20 @@ HUBERT_BASE = Wav2Vec2Bundle(
(512, 2, 2),
(512, 2, 2),
],
'extractor_conv_bias': False,
'encoder_embed_dim': 768,
'encoder_projection_dropout': 0.1,
'encoder_pos_conv_kernel': 128,
'encoder_pos_conv_groups': 16,
'encoder_num_layers': 12,
'encoder_num_heads': 12,
'encoder_attention_dropout': 0.1,
'encoder_ff_interm_features': 3072,
'encoder_ff_interm_dropout': 0.0,
'encoder_dropout': 0.1,
'encoder_layer_norm_first': False,
'encoder_layer_drop': 0.05,
'aux_num_out': None,
"extractor_conv_bias": False,
"encoder_embed_dim": 768,
"encoder_projection_dropout": 0.1,
"encoder_pos_conv_kernel": 128,
"encoder_pos_conv_groups": 16,
"encoder_num_layers": 12,
"encoder_num_heads": 12,
"encoder_attention_dropout": 0.1,
"encoder_ff_interm_features": 3072,
"encoder_ff_interm_dropout": 0.0,
"encoder_dropout": 0.1,
"encoder_layer_norm_first": False,
"encoder_layer_drop": 0.05,
"aux_num_out": None,
},
_sample_rate=16000,
)
......@@ -804,10 +806,10 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage.
""" # noqa: E501
HUBERT_LARGE = Wav2Vec2Bundle(
'hubert_fairseq_large_ll60k.pth',
"hubert_fairseq_large_ll60k.pth",
{
'extractor_mode': 'layer_norm',
'extractor_conv_layer_config': [
"extractor_mode": "layer_norm",
"extractor_conv_layer_config": [
(512, 10, 5),
(512, 3, 2),
(512, 3, 2),
......@@ -816,20 +818,20 @@ HUBERT_LARGE = Wav2Vec2Bundle(
(512, 2, 2),
(512, 2, 2),
],
'extractor_conv_bias': False,
'encoder_embed_dim': 1024,
'encoder_projection_dropout': 0.0,
'encoder_pos_conv_kernel': 128,
'encoder_pos_conv_groups': 16,
'encoder_num_layers': 24,
'encoder_num_heads': 16,
'encoder_attention_dropout': 0.0,
'encoder_ff_interm_features': 4096,
'encoder_ff_interm_dropout': 0.0,
'encoder_dropout': 0.0,
'encoder_layer_norm_first': True,
'encoder_layer_drop': 0.0,
'aux_num_out': None,
"extractor_conv_bias": False,
"encoder_embed_dim": 1024,
"encoder_projection_dropout": 0.0,
"encoder_pos_conv_kernel": 128,
"encoder_pos_conv_groups": 16,
"encoder_num_layers": 24,
"encoder_num_heads": 16,
"encoder_attention_dropout": 0.0,
"encoder_ff_interm_features": 4096,
"encoder_ff_interm_dropout": 0.0,
"encoder_dropout": 0.0,
"encoder_layer_norm_first": True,
"encoder_layer_drop": 0.0,
"aux_num_out": None,
},
_sample_rate=16000,
)
......@@ -848,10 +850,10 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage.
""" # noqa: E501
HUBERT_XLARGE = Wav2Vec2Bundle(
'hubert_fairseq_xlarge_ll60k.pth',
"hubert_fairseq_xlarge_ll60k.pth",
{
'extractor_mode': 'layer_norm',
'extractor_conv_layer_config': [
"extractor_mode": "layer_norm",
"extractor_conv_layer_config": [
(512, 10, 5),
(512, 3, 2),
(512, 3, 2),
......@@ -860,20 +862,20 @@ HUBERT_XLARGE = Wav2Vec2Bundle(
(512, 2, 2),
(512, 2, 2),
],
'extractor_conv_bias': False,
'encoder_embed_dim': 1280,
'encoder_projection_dropout': 0.0,
'encoder_pos_conv_kernel': 128,
'encoder_pos_conv_groups': 16,
'encoder_num_layers': 48,
'encoder_num_heads': 16,
'encoder_attention_dropout': 0.0,
'encoder_ff_interm_features': 5120,
'encoder_ff_interm_dropout': 0.0,
'encoder_dropout': 0.0,
'encoder_layer_norm_first': True,
'encoder_layer_drop': 0.0,
'aux_num_out': None,
"extractor_conv_bias": False,
"encoder_embed_dim": 1280,
"encoder_projection_dropout": 0.0,
"encoder_pos_conv_kernel": 128,
"encoder_pos_conv_groups": 16,
"encoder_num_layers": 48,
"encoder_num_heads": 16,
"encoder_attention_dropout": 0.0,
"encoder_ff_interm_features": 5120,
"encoder_ff_interm_dropout": 0.0,
"encoder_dropout": 0.0,
"encoder_layer_norm_first": True,
"encoder_layer_drop": 0.0,
"aux_num_out": None,
},
_sample_rate=16000,
)
......@@ -892,10 +894,10 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage.
""" # noqa: E501
HUBERT_ASR_LARGE = Wav2Vec2ASRBundle(
'hubert_fairseq_large_ll60k_asr_ls960.pth',
"hubert_fairseq_large_ll60k_asr_ls960.pth",
{
'extractor_mode': 'layer_norm',
'extractor_conv_layer_config': [
"extractor_mode": "layer_norm",
"extractor_conv_layer_config": [
(512, 10, 5),
(512, 3, 2),
(512, 3, 2),
......@@ -904,20 +906,20 @@ HUBERT_ASR_LARGE = Wav2Vec2ASRBundle(
(512, 2, 2),
(512, 2, 2),
],
'extractor_conv_bias': False,
'encoder_embed_dim': 1024,
'encoder_projection_dropout': 0.0,
'encoder_pos_conv_kernel': 128,
'encoder_pos_conv_groups': 16,
'encoder_num_layers': 24,
'encoder_num_heads': 16,
'encoder_attention_dropout': 0.0,
'encoder_ff_interm_features': 4096,
'encoder_ff_interm_dropout': 0.1,
'encoder_dropout': 0.0,
'encoder_layer_norm_first': True,
'encoder_layer_drop': 0.1,
'aux_num_out': 29,
"extractor_conv_bias": False,
"encoder_embed_dim": 1024,
"encoder_projection_dropout": 0.0,
"encoder_pos_conv_kernel": 128,
"encoder_pos_conv_groups": 16,
"encoder_num_layers": 24,
"encoder_num_heads": 16,
"encoder_attention_dropout": 0.0,
"encoder_ff_interm_features": 4096,
"encoder_ff_interm_dropout": 0.1,
"encoder_dropout": 0.0,
"encoder_layer_norm_first": True,
"encoder_layer_drop": 0.1,
"aux_num_out": 29,
},
_labels=utils._get_en_labels(),
_sample_rate=16000,
......@@ -939,10 +941,10 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
""" # noqa: E501
HUBERT_ASR_XLARGE = Wav2Vec2ASRBundle(
'hubert_fairseq_xlarge_ll60k_asr_ls960.pth',
"hubert_fairseq_xlarge_ll60k_asr_ls960.pth",
{
'extractor_mode': 'layer_norm',
'extractor_conv_layer_config': [
"extractor_mode": "layer_norm",
"extractor_conv_layer_config": [
(512, 10, 5),
(512, 3, 2),
(512, 3, 2),
......@@ -951,20 +953,20 @@ HUBERT_ASR_XLARGE = Wav2Vec2ASRBundle(
(512, 2, 2),
(512, 2, 2),
],
'extractor_conv_bias': False,
'encoder_embed_dim': 1280,
'encoder_projection_dropout': 0.0,
'encoder_pos_conv_kernel': 128,
'encoder_pos_conv_groups': 16,
'encoder_num_layers': 48,
'encoder_num_heads': 16,
'encoder_attention_dropout': 0.0,
'encoder_ff_interm_features': 5120,
'encoder_ff_interm_dropout': 0.1,
'encoder_dropout': 0.0,
'encoder_layer_norm_first': True,
'encoder_layer_drop': 0.1,
'aux_num_out': 29,
"extractor_conv_bias": False,
"encoder_embed_dim": 1280,
"encoder_projection_dropout": 0.0,
"encoder_pos_conv_kernel": 128,
"encoder_pos_conv_groups": 16,
"encoder_num_layers": 48,
"encoder_num_heads": 16,
"encoder_attention_dropout": 0.0,
"encoder_ff_interm_features": 5120,
"encoder_ff_interm_dropout": 0.1,
"encoder_dropout": 0.0,
"encoder_layer_norm_first": True,
"encoder_layer_drop": 0.1,
"aux_num_out": 29,
},
_labels=utils._get_en_labels(),
_sample_rate=16000,
......@@ -987,7 +989,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
VOXPOPULI_ASR_BASE_10K_DE = Wav2Vec2ASRBundle(
'wav2vec2_voxpopuli_base_10k_asr_de.pt',
"wav2vec2_voxpopuli_base_10k_asr_de.pt",
{
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
......@@ -1034,7 +1036,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
VOXPOPULI_ASR_BASE_10K_EN = Wav2Vec2ASRBundle(
'wav2vec2_voxpopuli_base_10k_asr_en.pt',
"wav2vec2_voxpopuli_base_10k_asr_en.pt",
{
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
......@@ -1059,7 +1061,7 @@ VOXPOPULI_ASR_BASE_10K_EN = Wav2Vec2ASRBundle(
"encoder_dropout": 0.0,
"encoder_layer_norm_first": False,
"encoder_layer_drop": 0.1,
"aux_num_out": 28
"aux_num_out": 28,
},
_labels=utils._get_vp_en_labels(),
_sample_rate=16000,
......@@ -1081,7 +1083,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
VOXPOPULI_ASR_BASE_10K_ES = Wav2Vec2ASRBundle(
'wav2vec2_voxpopuli_base_10k_asr_es.pt',
"wav2vec2_voxpopuli_base_10k_asr_es.pt",
{
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
......@@ -1106,7 +1108,7 @@ VOXPOPULI_ASR_BASE_10K_ES = Wav2Vec2ASRBundle(
"encoder_dropout": 0.0,
"encoder_layer_norm_first": False,
"encoder_layer_drop": 0.1,
"aux_num_out": 35
"aux_num_out": 35,
},
_labels=utils._get_es_labels(),
_sample_rate=16000,
......@@ -1127,7 +1129,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
""" # noqa: E501
VOXPOPULI_ASR_BASE_10K_FR = Wav2Vec2ASRBundle(
'wav2vec2_voxpopuli_base_10k_asr_fr.pt',
"wav2vec2_voxpopuli_base_10k_asr_fr.pt",
{
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
......@@ -1152,7 +1154,7 @@ VOXPOPULI_ASR_BASE_10K_FR = Wav2Vec2ASRBundle(
"encoder_dropout": 0.0,
"encoder_layer_norm_first": False,
"encoder_layer_drop": 0.1,
"aux_num_out": 43
"aux_num_out": 43,
},
_labels=utils._get_fr_labels(),
_sample_rate=16000,
......@@ -1173,7 +1175,7 @@ Please refer to :func:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage.
VOXPOPULI_ASR_BASE_10K_IT = Wav2Vec2ASRBundle(
'wav2vec2_voxpopuli_base_10k_asr_it.pt',
"wav2vec2_voxpopuli_base_10k_asr_it.pt",
{
"extractor_mode": "group_norm",
"extractor_conv_layer_config": [
......
torchaudio/pipelines/_wav2vec2/utils.py
View file @ 5859923a
def _get_en_labels():
return (
'|',
'E',
'T',
'A',
'O',
'N',
'I',
'H',
'S',
'R',
'D',
'L',
'U',
'M',
'W',
'C',
'F',
'G',
'Y',
'P',
'B',
'V',
'K',
"|",
"E",
"T",
"A",
"O",
"N",
"I",
"H",
"S",
"R",
"D",
"L",
"U",
"M",
"W",
"C",
"F",
"G",
"Y",
"P",
"B",
"V",
"K",
"'",
'X',
'J',
'Q',
'Z',
"X",
"J",
"Q",
"Z",
)
......
torchaudio/prototype/conformer.py
View file @ 5859923a
import math
import torch
from typing import List, Optional, Tuple
import torch
__all__ = ["Conformer"]
......@@ -12,9 +13,9 @@ PADDING_IDX = 1
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)
padding_mask = torch.arange(max_length, device=lengths.device, dtype=lengths.dtype).expand(
batch_size, max_length
) >= lengths.unsqueeze(1)
return padding_mask
......@@ -31,12 +32,8 @@ def _get_sinusoidal_embeddings(
from the description in Section 3.5 of "Attention Is All You Need".
"""
half_dim = embedding_dim // 2
t = (
torch.arange(half_dim, dtype=torch.float) * -math.log(10000) / (half_dim - 1)
).exp()
embedding_t = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(
1
) * t.unsqueeze(0)
t = (torch.arange(half_dim, dtype=torch.float) * -math.log(10000) / (half_dim - 1)).exp()
embedding_t = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * t.unsqueeze(0)
embeddings = torch.cat([embedding_t.sin(), embedding_t.cos()], dim=1)
if embedding_dim % 2 == 1:
embeddings = torch.cat([embeddings, torch.zeros(num_embeddings, 1)], dim=1)
......@@ -64,13 +61,16 @@ class ConvolutionModule(torch.nn.Module):
dropout: float = 0.0,
) -> None:
super().__init__()
assert (
depthwise_kernel_size - 1
) % 2 == 0, "depthwise_kernel_size must be odd to achieve 'SAME' padding."
assert (depthwise_kernel_size - 1) % 2 == 0, "depthwise_kernel_size must be odd to achieve 'SAME' padding."
self.layer_norm = torch.nn.LayerNorm(input_dim)
self.sequential = torch.nn.Sequential(
torch.nn.Conv1d(
input_dim, 2 * num_channels, 1, stride=1, padding=0, bias=bias,
input_dim,
2 * num_channels,
1,
stride=1,
padding=0,
bias=bias,
),
torch.nn.GLU(dim=1),
torch.nn.Conv1d(
......@@ -85,7 +85,12 @@ class ConvolutionModule(torch.nn.Module):
torch.nn.BatchNorm1d(num_channels),
torch.nn.SiLU(),
torch.nn.Conv1d(
num_channels, input_dim, kernel_size=1, stride=1, padding=0, bias=bias,
num_channels,
input_dim,
kernel_size=1,
stride=1,
padding=0,
bias=bias,
),
torch.nn.Dropout(dropout),
)
......@@ -159,9 +164,7 @@ class ConformerLayer(torch.nn.Module):
self.ffn1 = FeedForwardModule(input_dim, ffn_dim, dropout=dropout)
self.self_attn_layer_norm = torch.nn.LayerNorm(input_dim)
self.self_attn = torch.nn.MultiheadAttention(
input_dim, num_attention_heads, dropout=dropout
)
self.self_attn = torch.nn.MultiheadAttention(input_dim, num_attention_heads, dropout=dropout)
self.self_attn_dropout = torch.nn.Dropout(dropout)
self.conv_module = ConvolutionModule(
......@@ -173,9 +176,7 @@ class ConformerLayer(torch.nn.Module):
self.ffn2 = FeedForwardModule(input_dim, ffn_dim, dropout=dropout)
self.final_layer_norm = torch.nn.LayerNorm(input_dim)
def forward(
self, input: torch.Tensor, key_padding_mask: Optional[torch.Tensor]
) -> torch.Tensor:
def forward(self, input: torch.Tensor, key_padding_mask: Optional[torch.Tensor]) -> torch.Tensor:
r"""
Args:
input (torch.Tensor): input, with shape `(T, B, D)`.
......@@ -256,9 +257,7 @@ class Conv1dSubsampler(torch.nn.Module):
out = ((out.float() - 1) / 2 + 1).floor().long()
return out.to(torch.int32)
def forward(
self, input: torch.Tensor, lengths: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
def forward(self, input: torch.Tensor, lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
r"""
Args:
input (torch.Tensor): input frames, with shape `(B, T_in, in_channels)`.
......@@ -289,15 +288,11 @@ class SinusoidalPositionalEmbedding(torch.nn.Module):
init_size (int, optional): initial embedding count. (Default: 1024)
"""
def __init__(
self, embedding_dim: int, padding_idx: int = 0, init_size: int = 1024
) -> None:
def __init__(self, embedding_dim: int, padding_idx: int = 0, init_size: int = 1024) -> None:
super().__init__()
self.embedding_dim = embedding_dim
self.padding_idx = padding_idx
self.embeddings = _get_sinusoidal_embeddings(
init_size, embedding_dim, padding_idx
)
self.embeddings = _get_sinusoidal_embeddings(init_size, embedding_dim, padding_idx)
def forward(self, input: torch.Tensor) -> torch.Tensor:
r"""
......@@ -310,14 +305,10 @@ class SinusoidalPositionalEmbedding(torch.nn.Module):
B, T = input.shape
max_pos = self.padding_idx + 1 + T
if max_pos > self.embeddings.size(0):
self.embeddings = _get_sinusoidal_embeddings(
max_pos, self.embedding_dim, self.padding_idx
)
self.embeddings = _get_sinusoidal_embeddings(max_pos, self.embedding_dim, self.padding_idx)
self.embeddings = self.embeddings.to(input)
positions = _make_positions(input, self.padding_idx)
return (
self.embeddings.index_select(0, positions.view(-1)).view(B, T, -1).detach()
)
return self.embeddings.index_select(0, positions.view(-1)).view(B, T, -1).detach()
class Conformer(torch.nn.Module):
......@@ -370,16 +361,17 @@ class Conformer(torch.nn.Module):
super().__init__()
self.subsample = Conv1dSubsampler(
input_dim, conv_channels, conformer_layer_input_dim, conv_kernel_sizes,
input_dim,
conv_channels,
conformer_layer_input_dim,
conv_kernel_sizes,
)
self.position_embedding = SinusoidalPositionalEmbedding(
conformer_layer_input_dim,
padding_idx=PADDING_IDX,
init_size=max_source_positions + PADDING_IDX + 1,
)
self.linear = torch.nn.Linear(
conformer_layer_input_dim, conformer_layer_input_dim
)
self.linear = torch.nn.Linear(conformer_layer_input_dim, conformer_layer_input_dim)
self.dropout = torch.nn.Dropout(dropout)
self.conformer_layers = torch.nn.ModuleList(
[
......@@ -394,9 +386,7 @@ class Conformer(torch.nn.Module):
]
)
def forward(
self, input: torch.Tensor, lengths: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
def forward(self, input: torch.Tensor, lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
r"""
Args:
input (torch.Tensor): with shape `(B, T_in, input_dim)`.
......
torchaudio/prototype/emformer.py
View file @ 5859923a
......@@ -10,9 +10,9 @@ __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)
padding_mask = torch.arange(max_length, device=lengths.device, dtype=lengths.dtype).expand(
batch_size, max_length
) >= lengths.unsqueeze(1)
return padding_mask
......@@ -30,15 +30,8 @@ def _gen_padding_mask(
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
)
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
......@@ -54,9 +47,7 @@ def _get_activation_module(activation: str) -> torch.nn.Module:
raise ValueError(f"Unsupported activation {activation}")
def _get_weight_init_gains(
weight_init_scale_strategy: Optional[str], num_layers: int
) -> List[Optional[float]]:
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":
......@@ -64,17 +55,13 @@ def _get_weight_init_gains(
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}"
)
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:
assert len(col_widths) == len(
col_mask
), "Length of col_widths must match that of col_mask"
assert len(col_widths) == len(col_mask), "Length of col_widths must match that of col_mask"
mask_block = [
torch.ones(num_rows, col_width, device=device)
......@@ -110,9 +97,7 @@ class _EmformerAttention(torch.nn.Module):
super().__init__()
if input_dim % num_heads != 0:
raise ValueError(
f"input_dim ({input_dim}) is not a multiple of num_heads ({num_heads})."
)
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
......@@ -127,23 +112,15 @@ class _EmformerAttention(torch.nn.Module):
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
)
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]:
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
)
key, value = self.emb_to_key_value(mems_right_ctx_utterance_block).chunk(chunks=2, dim=2)
return key, value
def _gen_attention_probs(
......@@ -153,27 +130,17 @@ class _EmformerAttention(torch.nn.Module):
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
)
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.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
)
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,
......@@ -193,9 +160,7 @@ class _EmformerAttention(torch.nn.Module):
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)
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)
......@@ -203,37 +168,29 @@ class _EmformerAttention(torch.nn.Module):
[
key[: mems.size(0) + right_context_blocks_length],
left_context_key,
key[mems.size(0) + right_context_blocks_length:],
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:],
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)
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)
)
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
)
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
)
attention_probs = self._gen_attention_probs(attention_weights, attention_mask, padding_mask)
# Compute attention.
attention = torch.bmm(attention_probs, reshaped_value)
......@@ -249,7 +206,7 @@ class _EmformerAttention(torch.nn.Module):
summary_length = summary.size(0)
output_right_context = output_right_context_mems[: T - summary_length]
output_mems = 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:
......@@ -291,9 +248,7 @@ class _EmformerAttention(torch.nn.Module):
Tensor
updated memory elements, with shape `(M, B, D)`.
"""
output, output_mems, _, _ = self._forward_impl(
utterance, lengths, right_context, summary, mems, attention_mask
)
output, output_mems, _, _ = self._forward_impl(utterance, lengths, right_context, summary, mems, attention_mask)
return output, output_mems[:-1]
@torch.jit.export
......@@ -338,15 +293,8 @@ class _EmformerAttention(torch.nn.Module):
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
)
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,
......@@ -361,8 +309,8 @@ class _EmformerAttention(torch.nn.Module):
return (
output,
output_mems,
key[mems.size(0) + right_context.size(0):],
value[mems.size(0) + right_context.size(0):],
key[mems.size(0) + right_context.size(0) :],
value[mems.size(0) + right_context.size(0) :],
)
......@@ -410,9 +358,7 @@ class _EmformerLayer(torch.nn.Module):
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
)
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(
......@@ -433,18 +379,10 @@ class _EmformerLayer(torch.nn.Module):
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
)
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]
......@@ -453,12 +391,10 @@ class _EmformerLayer(torch.nn.Module):
) -> 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:]
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(
......@@ -471,9 +407,9 @@ class _EmformerLayer(torch.nn.Module):
) -> 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[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
......@@ -493,7 +429,7 @@ class _EmformerLayer(torch.nn.Module):
) -> 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) :],
layer_norm_input[: right_context.size(0)],
)
......@@ -501,7 +437,7 @@ class _EmformerLayer(torch.nn.Module):
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)]
return rc_output[right_context.size(0) :], rc_output[: right_context.size(0)]
def _apply_attention_forward(
self,
......@@ -512,9 +448,7 @@ class _EmformerLayer(torch.nn.Module):
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"
)
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)
......@@ -602,9 +536,7 @@ class _EmformerLayer(torch.nn.Module):
mems,
attention_mask,
)
output_utterance, output_right_context = self._apply_post_attention_ffn(
rc_output, utterance, right_context
)
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
......@@ -652,9 +584,7 @@ class _EmformerLayer(torch.nn.Module):
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
)
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
......@@ -708,12 +638,12 @@ class Emformer(torch.nn.Module):
self.use_mem = max_memory_size > 0
self.memory_op = torch.nn.AvgPool1d(
kernel_size=segment_length, stride=segment_length, ceil_mode=True,
kernel_size=segment_length,
stride=segment_length,
ceil_mode=True,
)
weight_init_gains = _get_weight_init_gains(
weight_init_scale_strategy, num_layers
)
weight_init_gains = _get_weight_init_gains(weight_init_scale_strategy, num_layers)
self.emformer_layers = torch.nn.ModuleList(
[
_EmformerLayer(
......@@ -747,12 +677,10 @@ class Emformer(torch.nn.Module):
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:])
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]:
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
......@@ -830,19 +758,13 @@ class Emformer(torch.nn.Module):
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_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)
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]:
def forward(self, input: torch.Tensor, lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
r"""Forward pass for training.
B: batch size;
......@@ -874,9 +796,7 @@ class Emformer(torch.nn.Module):
)
output = utterance
for layer in self.emformer_layers:
output, right_context, mems = layer(
output, lengths, right_context, mems, attention_mask
)
output, right_context, mems = layer(output, lengths, right_context, mems, attention_mask)
return output.permute(1, 0, 2), lengths
@torch.jit.export
......
torchaudio/prototype/rnnt.py
View file @ 5859923a
......@@ -15,13 +15,12 @@ class _TimeReduction(torch.nn.Module):
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]:
def forward(self, input: torch.Tensor, lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
r"""Forward pass.
B: batch size;
......@@ -64,6 +63,7 @@ class _CustomLSTM(torch.nn.Module):
layer_norm_epsilon (float, optional): value of epsilon to use in
layer normalization layers (Default: 1e-5)
"""
def __init__(
self,
input_dim: int,
......@@ -179,7 +179,9 @@ class _Transcriber(torch.nn.Module):
) -> None:
super().__init__()
self.input_linear = torch.nn.Linear(
input_dim, time_reduction_input_dim, bias=False,
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
......@@ -200,9 +202,7 @@ class _Transcriber(torch.nn.Module):
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]:
def forward(self, input: torch.Tensor, lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
r"""Forward pass for training.
B: batch size;
......@@ -225,12 +225,8 @@ class _Transcriber(torch.nn.Module):
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
)
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
......@@ -271,9 +267,7 @@ class _Transcriber(torch.nn.Module):
of ``infer``.
"""
input_linear_out = self.input_linear(input)
time_reduction_out, time_reduction_lengths = self.time_reduction(
input_linear_out, lengths
)
time_reduction_out, time_reduction_lengths = self.time_reduction(input_linear_out, lengths)
(
transformer_out,
transformer_lengths,
......@@ -299,6 +293,7 @@ class _Predictor(torch.nn.Module):
lstm_dropout (float, optional): LSTM dropout probability. (Default: 0.0)
"""
def __init__(
self,
num_symbols: int,
......@@ -368,9 +363,7 @@ class _Predictor(torch.nn.Module):
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, 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)
......@@ -426,10 +419,7 @@ class _Joiner(torch.nn.Module):
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()
)
joint_encodings = source_encodings.unsqueeze(2).contiguous() + target_encodings.unsqueeze(1).contiguous()
relu_out = self.relu(joint_encodings)
output = self.linear(relu_out)
return output, source_lengths, target_lengths
......@@ -447,9 +437,7 @@ class RNNT(torch.nn.Module):
joiner (torch.nn.Module): joint network.
"""
def __init__(
self, transcriber: _Transcriber, predictor: _Predictor, joiner: _Joiner
) -> None:
def __init__(self, transcriber: _Transcriber, predictor: _Predictor, joiner: _Joiner) -> None:
super().__init__()
self.transcriber = transcriber
self.predictor = predictor
......@@ -500,10 +488,13 @@ class RNNT(torch.nn.Module):
of ``forward``.
"""
source_encodings, source_lengths = self.transcriber(
input=sources, lengths=source_lengths,
input=sources,
lengths=source_lengths,
)
target_encodings, target_lengths, predictor_state = self.predictor(
input=targets, lengths=target_lengths, state=predictor_state,
input=targets,
lengths=target_lengths,
state=predictor_state,
)
output, source_lengths, target_lengths = self.joiner(
source_encodings=source_encodings,
......@@ -558,7 +549,9 @@ class RNNT(torch.nn.Module):
@torch.jit.export
def transcribe(
self, sources: torch.Tensor, source_lengths: torch.Tensor,
self,
sources: torch.Tensor,
source_lengths: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
r"""Applies transcription network to sources in non-streaming mode.
......
torchaudio/prototype/rnnt_decoder.py
View file @ 5859923a
......@@ -35,21 +35,14 @@ def _batch_state(hypos: List[Hypothesis]) -> List[List[torch.Tensor]]:
for i in range(len(hypos[0].state)):
batched_state_components: List[torch.Tensor] = []
for j in range(len(hypos[0].state[i])):
batched_state_components.append(
torch.cat([hypo.state[i][j] for hypo in hypos])
)
batched_state_components.append(torch.cat([hypo.state[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]]:
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
]
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:
......@@ -57,18 +50,14 @@ def _default_hypo_sort_key(hypo: Hypothesis) -> float:
def _compute_updated_scores(
hypos: List[Hypothesis], next_token_probs: torch.Tensor, beam_width: int,
hypos: List[Hypothesis],
next_token_probs: torch.Tensor,
beam_width: int,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
hypo_scores = torch.tensor([h.score 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_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
......@@ -114,9 +103,7 @@ class RNNTBeamSearch(torch.nn.Module):
self.step_max_tokens = step_max_tokens
def _init_b_hypos(
self, hypo: Optional[Hypothesis], device: torch.device
) -> List[Hypothesis]:
def _init_b_hypos(self, hypo: Optional[Hypothesis], device: torch.device) -> List[Hypothesis]:
if hypo is not None:
token = hypo.tokens[-1]
state = hypo.state
......@@ -125,9 +112,7 @@ class RNNTBeamSearch(torch.nn.Module):
state = None
one_tensor = torch.tensor([1], device=device)
pred_out, _, pred_state = self.model.predict(
torch.tensor([[token]], device=device), one_tensor, state
)
pred_out, _, pred_state = self.model.predict(torch.tensor([[token]], device=device), one_tensor, state)
init_hypo = Hypothesis(
tokens=[token],
predictor_out=pred_out[0].detach(),
......@@ -150,9 +135,7 @@ class RNNTBeamSearch(torch.nn.Module):
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
)
joined_out = torch.nn.functional.log_softmax(joined_out / self.temperature, dim=3)
joined_out[:, :, :, :4].add_(-99999) # blank out invalid tokens
return joined_out[:, 0, 0]
......@@ -220,9 +203,7 @@ class RNNTBeamSearch(torch.nn.Module):
new_scores.append(score)
if base_hypos:
new_hypos = self._gen_new_hypos(
base_hypos, new_tokens, new_scores, t, device
)
new_hypos = self._gen_new_hypos(base_hypos, new_tokens, new_scores, t, device)
else:
new_hypos: List[Hypothesis] = []
......@@ -239,7 +220,9 @@ class RNNTBeamSearch(torch.nn.Module):
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,
tgt_tokens,
torch.tensor([1] * len(base_hypos), device=device),
states,
)
new_hypos: List[Hypothesis] = []
for i, h_a in enumerate(base_hypos):
......@@ -258,7 +241,10 @@ class RNNTBeamSearch(torch.nn.Module):
return new_hypos
def _search(
self, enc_out: torch.Tensor, hypo: Optional[Hypothesis], beam_width: int,
self,
enc_out: torch.Tensor,
hypo: Optional[Hypothesis],
beam_width: int,
) -> List[Hypothesis]:
n_time_steps = enc_out.shape[1]
device = enc_out.device
......@@ -272,33 +258,35 @@ class RNNTBeamSearch(torch.nn.Module):
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 = 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,
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,
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(hypo) for hypo in b_hypos]
).topk(beam_width)
_, sorted_idx = torch.tensor([self.hypo_sort_key(hypo) for hypo 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]:
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;
......@@ -319,9 +307,7 @@ class RNNTBeamSearch(torch.nn.Module):
if input.dim() == 2:
input = input.unsqueeze(0)
assert length.shape == () or length.shape == (
1,
), "length must be of shape () or (1,)"
assert length.shape == () or length.shape == (1,), "length must be of shape () or (1,)"
if input.dim() == 0:
input = input.unsqueeze(0)
......@@ -367,9 +353,7 @@ class RNNTBeamSearch(torch.nn.Module):
if input.dim() == 2:
input = input.unsqueeze(0)
assert length.shape == () or length.shape == (
1,
), "length must be of shape () or (1,)"
assert length.shape == () or length.shape == (1,), "length must be of shape () or (1,)"
if input.dim() == 0:
input = input.unsqueeze(0)
......
torchaudio/sox_effects/__init__.py
View file @ 5859923a
from torchaudio._internal import module_utils as _mod_utils
from .sox_effects import (
init_sox_effects,
shutdown_sox_effects,
......@@ -10,13 +11,14 @@ from .sox_effects import (
if _mod_utils.is_sox_available():
import atexit
init_sox_effects()
atexit.register(shutdown_sox_effects)
__all__ = [
'init_sox_effects',
'shutdown_sox_effects',
'effect_names',
'apply_effects_tensor',
'apply_effects_file',
"init_sox_effects",
"shutdown_sox_effects",
"effect_names",
"apply_effects_tensor",
"apply_effects_file",
]
torchaudio/sox_effects/sox_effects.py
View file @ 5859923a
......@@ -2,7 +2,6 @@ import os
from typing import List, Tuple, Optional
import torch
import torchaudio
from torchaudio._internal import module_utils as _mod_utils
from torchaudio.utils.sox_utils import list_effects
......@@ -53,10 +52,10 @@ def effect_names() -> List[str]:
@_mod_utils.requires_sox()
def apply_effects_tensor(
tensor: torch.Tensor,
sample_rate: int,
effects: List[List[str]],
channels_first: bool = True,
tensor: torch.Tensor,
sample_rate: int,
effects: List[List[str]],
channels_first: bool = True,
) -> Tuple[torch.Tensor, int]:
"""Apply sox effects to given Tensor
......@@ -149,17 +148,16 @@ def apply_effects_tensor(
>>> waveform, sample_rate = transform(waveform, input_sample_rate)
>>> assert sample_rate == 8000
"""
return torch.ops.torchaudio.sox_effects_apply_effects_tensor(
tensor, sample_rate, effects, channels_first)
return torch.ops.torchaudio.sox_effects_apply_effects_tensor(tensor, sample_rate, effects, channels_first)
@_mod_utils.requires_sox()
def apply_effects_file(
path: str,
effects: List[List[str]],
normalize: bool = True,
channels_first: bool = True,
format: Optional[str] = None,
path: str,
effects: List[List[str]],
normalize: bool = True,
channels_first: bool = True,
format: Optional[str] = None,
) -> Tuple[torch.Tensor, int]:
"""Apply sox effects to the audio file and load the resulting data as Tensor
......@@ -265,9 +263,7 @@ def apply_effects_file(
>>> pass
"""
if not torch.jit.is_scripting():
if hasattr(path, 'read'):
return torchaudio._torchaudio.apply_effects_fileobj(
path, effects, normalize, channels_first, format)
if hasattr(path, "read"):
return torchaudio._torchaudio.apply_effects_fileobj(path, effects, normalize, channels_first, format)
path = os.fspath(path)
return torch.ops.torchaudio.sox_effects_apply_effects_file(
path, effects, normalize, channels_first, format)
return torch.ops.torchaudio.sox_effects_apply_effects_file(path, effects, normalize, channels_first, format)
torchaudio/transforms.py
View file @ 5859923a
......@@ -14,31 +14,31 @@ from .functional.functional import (
)
__all__ = [
'Spectrogram',
'InverseSpectrogram',
'GriffinLim',
'AmplitudeToDB',
'MelScale',
'InverseMelScale',
'MelSpectrogram',
'MFCC',
'LFCC',
'MuLawEncoding',
'MuLawDecoding',
'Resample',
'TimeStretch',
'Fade',
'FrequencyMasking',
'TimeMasking',
'SlidingWindowCmn',
'Vad',
'SpectralCentroid',
'Vol',
'ComputeDeltas',
'PitchShift',
'RNNTLoss',
'PSD',
'MVDR',
"Spectrogram",
"InverseSpectrogram",
"GriffinLim",
"AmplitudeToDB",
"MelScale",
"InverseMelScale",
"MelSpectrogram",
"MFCC",
"LFCC",
"MuLawEncoding",
"MuLawDecoding",
"Resample",
"TimeStretch",
"Fade",
"FrequencyMasking",
"TimeMasking",
"SlidingWindowCmn",
"Vad",
"SpectralCentroid",
"Vol",
"ComputeDeltas",
"PitchShift",
"RNNTLoss",
"PSD",
"MVDR",
]
......@@ -73,21 +73,23 @@ class Spectrogram(torch.nn.Module):
>>> 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.,
normalized: bool = False,
wkwargs: Optional[dict] = None,
center: bool = True,
pad_mode: str = "reflect",
onesided: bool = True,
return_complex: Optional[bool] = None) -> None:
__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: bool = 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__()
self.n_fft = n_fft
# number of FFT bins. the returned STFT result will have n_fft // 2 + 1
......@@ -95,7 +97,7 @@ class Spectrogram(torch.nn.Module):
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.register_buffer("window", window)
self.pad = pad
self.power = power
self.normalized = normalized
......@@ -162,19 +164,21 @@ class InverseSpectrogram(torch.nn.Module):
>>> 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: bool = False,
wkwargs: Optional[dict] = None,
center: bool = True,
pad_mode: str = "reflect",
onesided: bool = True) -> None:
__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: bool = 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
......@@ -182,7 +186,7 @@ class InverseSpectrogram(torch.nn.Module):
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.register_buffer("window", window)
self.pad = pad
self.normalized = normalized
self.center = center
......@@ -242,31 +246,32 @@ class GriffinLim(torch.nn.Module):
>>> 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.,
wkwargs: Optional[dict] = None,
momentum: float = 0.99,
length: Optional[int] = None,
rand_init: bool = True) -> None:
__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__()
assert momentum < 1, 'momentum={} > 1 can be unstable'.format(momentum)
assert momentum >= 0, 'momentum={} < 0'.format(momentum)
assert momentum < 1, "momentum={} > 1 can be unstable".format(momentum)
assert momentum >= 0, "momentum={} < 0".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.register_buffer("window", window)
self.length = length
self.power = power
self.momentum = momentum / (1 + momentum)
......@@ -282,8 +287,18 @@ class GriffinLim(torch.nn.Module):
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)
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):
......@@ -299,15 +314,15 @@ class AmplitudeToDB(torch.nn.Module):
top_db (float or None, optional): minimum negative cut-off in decibels. A reasonable
number is 80. (Default: ``None``)
"""
__constants__ = ['multiplier', 'amin', 'ref_value', 'db_multiplier']
__constants__ = ["multiplier", "amin", "ref_value", "db_multiplier"]
def __init__(self, stype: str = 'power', top_db: Optional[float] = None) -> None:
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')
raise ValueError("top_db must be positive value")
self.top_db = top_db
self.multiplier = 10.0 if stype == 'power' else 20.0
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))
......@@ -344,16 +359,18 @@ class MelScale(torch.nn.Module):
: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.,
f_max: Optional[float] = None,
n_stft: int = 201,
norm: Optional[str] = None,
mel_scale: str = "htk") -> None:
__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
......@@ -362,11 +379,9 @@ class MelScale(torch.nn.Module):
self.norm = norm
self.mel_scale = mel_scale
assert f_min <= self.f_max, '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)
assert f_min <= self.f_max, "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"""
......@@ -404,21 +419,32 @@ class InverseMelScale(torch.nn.Module):
(area normalization). (Default: ``None``)
mel_scale (str, optional): Scale to use: ``htk`` or ``slaney``. (Default: ``htk``)
"""
__constants__ = ['n_stft', 'n_mels', 'sample_rate', 'f_min', 'f_max', 'max_iter', 'tolerance_loss',
'tolerance_change', 'sgdargs']
def __init__(self,
n_stft: int,
n_mels: int = 128,
sample_rate: int = 16000,
f_min: float = 0.,
f_max: Optional[float] = None,
max_iter: int = 100000,
tolerance_loss: float = 1e-5,
tolerance_change: float = 1e-8,
sgdargs: Optional[dict] = None,
norm: Optional[str] = None,
mel_scale: str = "htk") -> None:
__constants__ = [
"n_stft",
"n_mels",
"sample_rate",
"f_min",
"f_max",
"max_iter",
"tolerance_loss",
"tolerance_change",
"sgdargs",
]
def __init__(
self,
n_stft: int,
n_mels: int = 128,
sample_rate: int = 16000,
f_min: float = 0.0,
f_max: Optional[float] = None,
max_iter: int = 100000,
tolerance_loss: float = 1e-5,
tolerance_change: float = 1e-8,
sgdargs: Optional[dict] = None,
norm: Optional[str] = None,
mel_scale: str = "htk",
) -> None:
super(InverseMelScale, self).__init__()
self.n_mels = n_mels
self.sample_rate = sample_rate
......@@ -427,13 +453,12 @@ class InverseMelScale(torch.nn.Module):
self.max_iter = max_iter
self.tolerance_loss = tolerance_loss
self.tolerance_change = tolerance_change
self.sgdargs = sgdargs or {'lr': 0.1, 'momentum': 0.9}
self.sgdargs = sgdargs or {"lr": 0.1, "momentum": 0.9}
assert f_min <= self.f_max, 'Require f_min: {} < f_max: {}'.format(f_min, self.f_max)
assert f_min <= self.f_max, "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,
norm, mel_scale)
self.register_buffer('fb', fb)
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"""
......@@ -452,12 +477,13 @@ class InverseMelScale(torch.nn.Module):
melspec = melspec.transpose(-1, -2)
assert self.n_mels == n_mels
specgram = torch.rand(melspec.size()[0], time, freq, requires_grad=True,
dtype=melspec.dtype, device=melspec.device)
specgram = torch.rand(
melspec.size()[0], time, freq, requires_grad=True, dtype=melspec.dtype, device=melspec.device
)
optim = torch.optim.SGD([specgram], **self.sgdargs)
loss = float('inf')
loss = float("inf")
for _ in range(self.max_iter):
optim.zero_grad()
diff = melspec - specgram.matmul(self.fb)
......@@ -527,26 +553,28 @@ class MelSpectrogram(torch.nn.Module):
: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.,
f_max: Optional[float] = None,
pad: int = 0,
n_mels: int = 128,
window_fn: Callable[..., Tensor] = torch.hann_window,
power: float = 2.,
normalized: bool = False,
wkwargs: Optional[dict] = None,
center: bool = True,
pad_mode: str = "reflect",
onesided: bool = True,
norm: Optional[str] = None,
mel_scale: str = "htk") -> None:
__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: bool = True,
norm: Optional[str] = None,
mel_scale: str = "htk",
) -> None:
super(MelSpectrogram, self).__init__()
self.sample_rate = sample_rate
self.n_fft = n_fft
......@@ -558,19 +586,21 @@ class MelSpectrogram(torch.nn.Module):
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=onesided)
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=onesided,
)
self.mel_scale = MelScale(
self.n_mels,
self.sample_rate,
self.f_min,
self.f_max,
self.n_fft // 2 + 1,
norm,
mel_scale
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:
......@@ -609,33 +639,35 @@ class MFCC(torch.nn.Module):
: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:
__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))
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)
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')
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.register_buffer("dct_mat", dct_mat)
self.log_mels = log_mels
def forward(self, waveform: Tensor) -> Tensor:
......@@ -685,22 +717,24 @@ class LFCC(torch.nn.Module):
: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.,
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:
__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))
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)
......@@ -709,13 +743,13 @@ class LFCC(torch.nn.Module):
self.dct_type = dct_type
self.norm = norm
self.top_db = 80.0
self.amplitude_to_DB = AmplitudeToDB('power', self.top_db)
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')
raise ValueError("Cannot select more LFCC coefficients than # fft bins")
filter_mat = F.linear_fbanks(
n_freqs=self.Spectrogram.n_fft // 2 + 1,
......@@ -727,7 +761,7 @@ class LFCC(torch.nn.Module):
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.register_buffer("dct_mat", dct_mat)
self.log_lf = log_lf
def forward(self, waveform: Tensor) -> Tensor:
......@@ -770,7 +804,7 @@ class MuLawEncoding(torch.nn.Module):
>>> mulawtrans = transform(waveform)
"""
__constants__ = ['quantization_channels']
__constants__ = ["quantization_channels"]
def __init__(self, quantization_channels: int = 256) -> None:
super(MuLawEncoding, self).__init__()
......@@ -802,7 +836,7 @@ class MuLawDecoding(torch.nn.Module):
>>> transform = torchaudio.transforms.MuLawDecoding(quantization_channels=512)
>>> mulawtrans = transform(waveform)
"""
__constants__ = ['quantization_channels']
__constants__ = ["quantization_channels"]
def __init__(self, quantization_channels: int = 256) -> None:
super(MuLawDecoding, self).__init__()
......@@ -853,15 +887,15 @@ class Resample(torch.nn.Module):
"""
def __init__(
self,
orig_freq: int = 16000,
new_freq: int = 16000,
resampling_method: str = 'sinc_interpolation',
lowpass_filter_width: int = 6,
rolloff: float = 0.99,
beta: Optional[float] = None,
*,
dtype: Optional[torch.dtype] = None,
self,
orig_freq: int = 16000,
new_freq: int = 16000,
resampling_method: str = "sinc_interpolation",
lowpass_filter_width: int = 6,
rolloff: float = 0.99,
beta: Optional[float] = None,
*,
dtype: Optional[torch.dtype] = None,
) -> None:
super().__init__()
......@@ -875,10 +909,16 @@ class Resample(torch.nn.Module):
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)
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"""
......@@ -890,9 +930,7 @@ class Resample(torch.nn.Module):
"""
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)
return _apply_sinc_resample_kernel(waveform, self.orig_freq, self.new_freq, self.gcd, self.kernel, self.width)
class ComputeDeltas(torch.nn.Module):
......@@ -904,7 +942,7 @@ class ComputeDeltas(torch.nn.Module):
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']
__constants__ = ["win_length"]
def __init__(self, win_length: int = 5, mode: str = "replicate") -> None:
super(ComputeDeltas, self).__init__()
......@@ -954,19 +992,16 @@ class TimeStretch(torch.nn.Module):
:alt: Spectrogram streched by 0.9
"""
__constants__ = ['fixed_rate']
__constants__ = ["fixed_rate"]
def __init__(self,
hop_length: Optional[int] = None,
n_freq: int = 201,
fixed_rate: Optional[float] = None) -> None:
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])
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"""
......@@ -983,8 +1018,7 @@ class TimeStretch(torch.nn.Module):
"""
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.")
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
......@@ -1007,10 +1041,7 @@ class Fade(torch.nn.Module):
>>> faded_waveform = transform(waveform)
"""
def __init__(self,
fade_in_len: int = 0,
fade_out_len: int = 0,
fade_shape: str = "linear") -> None:
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
......@@ -1026,11 +1057,7 @@ class Fade(torch.nn.Module):
"""
waveform_length = waveform.size()[-1]
device = waveform.device
return (
self._fade_in(waveform_length, device)
* self._fade_out(waveform_length, device)
* waveform
)
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)
......@@ -1043,7 +1070,7 @@ class Fade(torch.nn.Module):
fade = torch.pow(2, (fade - 1)) * fade
if self.fade_shape == "logarithmic":
fade = torch.log10(.1 + fade) + 1
fade = torch.log10(0.1 + fade) + 1
if self.fade_shape == "quarter_sine":
fade = torch.sin(fade * math.pi / 2)
......@@ -1058,10 +1085,10 @@ class Fade(torch.nn.Module):
ones = torch.ones(waveform_length - self.fade_out_len, device=device)
if self.fade_shape == "linear":
fade = - fade + 1
fade = -fade + 1
if self.fade_shape == "exponential":
fade = torch.pow(2, - fade) * (1 - fade)
fade = torch.pow(2, -fade) * (1 - fade)
if self.fade_shape == "logarithmic":
fade = torch.log10(1.1 - fade) + 1
......@@ -1084,7 +1111,7 @@ class _AxisMasking(torch.nn.Module):
iid_masks (bool): Applies iid masks to each of the examples in the batch dimension.
This option is applicable only when the input tensor is 4D.
"""
__constants__ = ['mask_param', 'axis', 'iid_masks']
__constants__ = ["mask_param", "axis", "iid_masks"]
def __init__(self, mask_param: int, axis: int, iid_masks: bool) -> None:
......@@ -1093,7 +1120,7 @@ class _AxisMasking(torch.nn.Module):
self.axis = axis
self.iid_masks = iid_masks
def forward(self, specgram: Tensor, mask_value: float = 0.) -> Tensor:
def forward(self, specgram: Tensor, mask_value: float = 0.0) -> Tensor:
r"""
Args:
specgram (Tensor): Tensor of dimension `(..., freq, time)`.
......@@ -1180,12 +1207,12 @@ class Vol(torch.nn.Module):
gain_type (str, optional): Type of gain. One of: ``amplitude``, ``power``, ``db`` (Default: ``amplitude``)
"""
def __init__(self, gain: float, gain_type: str = 'amplitude'):
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:
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:
......@@ -1221,11 +1248,9 @@ class SlidingWindowCmn(torch.nn.Module):
norm_vars (bool, optional): If true, normalize variance to one. (bool, default = false)
"""
def __init__(self,
cmn_window: int = 600,
min_cmn_window: int = 100,
center: bool = False,
norm_vars: bool = False) -> None:
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
......@@ -1240,8 +1265,7 @@ class SlidingWindowCmn(torch.nn.Module):
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)
cmn_specgram = F.sliding_window_cmn(specgram, self.cmn_window, self.min_cmn_window, self.center, self.norm_vars)
return cmn_specgram
......@@ -1297,24 +1321,26 @@ class Vad(torch.nn.Module):
- 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 = .35,
noise_up_time: float = .1,
noise_down_time: float = .01,
noise_reduction_amount: float = 1.35,
measure_freq: float = 20.0,
measure_duration: Optional[float] = None,
measure_smooth_time: float = .4,
hp_filter_freq: float = 50.,
lp_filter_freq: float = 6000.,
hp_lifter_freq: float = 150.,
lp_lifter_freq: float = 2000.) -> None:
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
......@@ -1386,23 +1412,25 @@ class SpectralCentroid(torch.nn.Module):
>>> 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:
__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.register_buffer("window", window)
self.pad = pad
def forward(self, waveform: Tensor) -> Tensor:
......@@ -1414,8 +1442,9 @@ class SpectralCentroid(torch.nn.Module):
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)
return F.spectral_centroid(
waveform, self.sample_rate, self.pad, self.window, self.n_fft, self.hop_length, self.win_length
)
class PitchShift(torch.nn.Module):
......@@ -1438,17 +1467,19 @@ class PitchShift(torch.nn.Module):
>>> 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']
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:
__constants__ = ["sample_rate", "n_steps", "bins_per_octave", "n_fft", "win_length", "hop_length"]
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(PitchShift, self).__init__()
self.n_steps = n_steps
self.bins_per_octave = bins_per_octave
......@@ -1457,7 +1488,7 @@ class PitchShift(torch.nn.Module):
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)
self.register_buffer("window", window)
def forward(self, waveform: Tensor) -> Tensor:
r"""
......@@ -1468,8 +1499,16 @@ class PitchShift(torch.nn.Module):
Tensor: The pitch-shifted audio of shape `(..., time)`.
"""
return F.pitch_shift(waveform, self.sample_rate, self.n_steps, self.bins_per_octave, self.n_fft,
self.win_length, self.hop_length, self.window)
return F.pitch_shift(
waveform,
self.sample_rate,
self.n_steps,
self.bins_per_octave,
self.n_fft,
self.win_length,
self.hop_length,
self.window,
)
class RNNTLoss(torch.nn.Module):
......@@ -1506,7 +1545,7 @@ class RNNTLoss(torch.nn.Module):
def __init__(
self,
blank: int = -1,
clamp: float = -1.,
clamp: float = -1.0,
reduction: str = "mean",
):
super().__init__()
......@@ -1532,15 +1571,7 @@ class RNNTLoss(torch.nn.Module):
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
)
return F.rnnt_loss(logits, targets, logit_lengths, target_lengths, self.blank, self.clamp, self.reduction)
def _get_mat_trace(input: torch.Tensor, dim1: int = -1, dim2: int = -2) -> torch.Tensor:
......@@ -1557,8 +1588,7 @@ def _get_mat_trace(input: torch.Tensor, dim1: int = -1, dim2: int = -2) -> torch
torch.Tensor: trace of the input Tensor
"""
assert input.ndim >= 2, "The dimension of the tensor must be at least 2."
assert input.shape[dim1] == input.shape[dim2],\
"The size of ``dim1`` and ``dim2`` must be the same."
assert input.shape[dim1] == input.shape[dim2], "The size of ``dim1`` and ``dim2`` must be the same."
input = torch.diagonal(input, 0, dim1=dim1, dim2=dim2)
return input.sum(dim=-1)
......@@ -1684,8 +1714,11 @@ class MVDR(torch.nn.Module):
online: bool = False,
):
super().__init__()
assert solution in ["ref_channel", "stv_evd", "stv_power"],\
"Unknown solution provided. Must be one of [``ref_channel``, ``stv_evd``, ``stv_power``]."
assert solution in [
"ref_channel",
"stv_evd",
"stv_power",
], "Unknown solution provided. Must be one of [``ref_channel``, ``stv_evd``, ``stv_power``]."
self.ref_channel = ref_channel
self.solution = solution
self.multi_mask = multi_mask
......@@ -1698,10 +1731,10 @@ class MVDR(torch.nn.Module):
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)
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,
......@@ -1710,7 +1743,7 @@ class MVDR(torch.nn.Module):
mask_s: torch.Tensor,
mask_n: torch.Tensor,
reference_vector: torch.Tensor,
solution: str = 'ref_channel',
solution: str = "ref_channel",
diagonal_loading: bool = True,
diag_eps: float = 1e-7,
eps: float = 1e-8,
......@@ -1788,7 +1821,7 @@ class MVDR(torch.nn.Module):
psd_s: torch.Tensor,
psd_n: torch.Tensor,
reference_vector: torch.Tensor,
solution: str = 'ref_channel',
solution: str = "ref_channel",
diagonal_loading: bool = True,
diag_eps: float = 1e-7,
eps: float = 1e-8,
......@@ -1851,10 +1884,7 @@ class MVDR(torch.nn.Module):
return stv
def _get_steering_vector_power(
self,
psd_s: torch.Tensor,
psd_n: torch.Tensor,
reference_vector: torch.Tensor
self, psd_s: torch.Tensor, psd_n: torch.Tensor, reference_vector: torch.Tensor
) -> torch.Tensor:
r"""Estimate the steering vector by the power method.
......@@ -1876,11 +1906,7 @@ class MVDR(torch.nn.Module):
stv = torch.matmul(psd_s, stv)
return stv
def _apply_beamforming_vector(
self,
specgram: torch.Tensor,
beamform_vector: torch.Tensor
) -> torch.Tensor:
def _apply_beamforming_vector(self, specgram: torch.Tensor, beamform_vector: torch.Tensor) -> torch.Tensor:
r"""Apply the beamforming weight to the noisy STFT
Args:
specgram (torch.tensor): multi-channel noisy STFT
......@@ -1896,12 +1922,7 @@ class MVDR(torch.nn.Module):
specgram_enhanced = torch.einsum("...fc,...cft->...ft", [beamform_vector.conj(), specgram])
return specgram_enhanced
def _tik_reg(
self,
mat: torch.Tensor,
reg: float = 1e-7,
eps: float = 1e-8
) -> torch.Tensor:
def _tik_reg(self, mat: torch.Tensor, reg: float = 1e-7, eps: float = 1e-8) -> torch.Tensor:
"""Perform Tikhonov regularization (only modifying real part).
Args:
mat (torch.Tensor): input matrix (..., channel, channel)
......@@ -1922,10 +1943,7 @@ class MVDR(torch.nn.Module):
return mat
def forward(
self,
specgram: torch.Tensor,
mask_s: torch.Tensor,
mask_n: Optional[torch.Tensor] = None
self, specgram: torch.Tensor, mask_s: torch.Tensor, mask_n: Optional[torch.Tensor] = None
) -> torch.Tensor:
"""Perform MVDR beamforming.
......@@ -1946,9 +1964,7 @@ class MVDR(torch.nn.Module):
"""
dtype = specgram.dtype
if specgram.ndim < 3:
raise ValueError(
f"Expected at least 3D tensor (..., channel, freq, time). Found: {specgram.shape}"
)
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``.\
......@@ -1958,9 +1974,7 @@ class MVDR(torch.nn.Module):
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``."
)
warnings.warn("``mask_n`` is not provided, use ``1 - mask_s`` as ``mask_n``.")
mask_n = 1 - mask_s
shape = specgram.size()
......@@ -1977,33 +1991,15 @@ class MVDR(torch.nn.Module):
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 = 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
psd_s, psd_n, mask_s, mask_n, u, self.solution, self.diag_loading, self.diag_eps
)
else:
w_mvdr = self._get_mvdr_vector(
psd_s,
psd_n,
u,
self.solution,
self.diag_loading,
self.diag_eps
)
w_mvdr = self._get_mvdr_vector(psd_s, psd_n, u, self.solution, self.diag_loading, self.diag_eps)
specgram_enhanced = self._apply_beamforming_vector(specgram, w_mvdr)
......
torchaudio/utils/__init__.py
View file @ 5859923a
from torchaudio._internal import module_utils as _mod_utils
from . import (
sox_utils,
)
from torchaudio._internal import module_utils as _mod_utils
if _mod_utils.is_sox_available():
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment