torchaudio 2.7.0__cp313-cp313t-manylinux_2_28_x86_64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of torchaudio might be problematic. Click here for more details.

Files changed (148) hide show
  1. torchaudio/__init__.py +53 -0
  2. torchaudio/_backend/__init__.py +61 -0
  3. torchaudio/_backend/backend.py +53 -0
  4. torchaudio/_backend/common.py +52 -0
  5. torchaudio/_backend/ffmpeg.py +334 -0
  6. torchaudio/_backend/soundfile.py +54 -0
  7. torchaudio/_backend/soundfile_backend.py +457 -0
  8. torchaudio/_backend/sox.py +91 -0
  9. torchaudio/_backend/utils.py +317 -0
  10. torchaudio/_extension/__init__.py +74 -0
  11. torchaudio/_extension/utils.py +180 -0
  12. torchaudio/_internal/__init__.py +10 -0
  13. torchaudio/_internal/module_utils.py +113 -0
  14. torchaudio/backend/__init__.py +8 -0
  15. torchaudio/backend/_no_backend.py +25 -0
  16. torchaudio/backend/_sox_io_backend.py +294 -0
  17. torchaudio/backend/common.py +13 -0
  18. torchaudio/backend/no_backend.py +14 -0
  19. torchaudio/backend/soundfile_backend.py +14 -0
  20. torchaudio/backend/sox_io_backend.py +14 -0
  21. torchaudio/compliance/__init__.py +5 -0
  22. torchaudio/compliance/kaldi.py +813 -0
  23. torchaudio/datasets/__init__.py +47 -0
  24. torchaudio/datasets/cmuarctic.py +157 -0
  25. torchaudio/datasets/cmudict.py +186 -0
  26. torchaudio/datasets/commonvoice.py +86 -0
  27. torchaudio/datasets/dr_vctk.py +121 -0
  28. torchaudio/datasets/fluentcommands.py +108 -0
  29. torchaudio/datasets/gtzan.py +1118 -0
  30. torchaudio/datasets/iemocap.py +147 -0
  31. torchaudio/datasets/librilight_limited.py +111 -0
  32. torchaudio/datasets/librimix.py +133 -0
  33. torchaudio/datasets/librispeech.py +174 -0
  34. torchaudio/datasets/librispeech_biasing.py +189 -0
  35. torchaudio/datasets/libritts.py +168 -0
  36. torchaudio/datasets/ljspeech.py +107 -0
  37. torchaudio/datasets/musdb_hq.py +139 -0
  38. torchaudio/datasets/quesst14.py +136 -0
  39. torchaudio/datasets/snips.py +157 -0
  40. torchaudio/datasets/speechcommands.py +183 -0
  41. torchaudio/datasets/tedlium.py +218 -0
  42. torchaudio/datasets/utils.py +54 -0
  43. torchaudio/datasets/vctk.py +143 -0
  44. torchaudio/datasets/voxceleb1.py +309 -0
  45. torchaudio/datasets/yesno.py +89 -0
  46. torchaudio/functional/__init__.py +127 -0
  47. torchaudio/functional/_alignment.py +128 -0
  48. torchaudio/functional/filtering.py +1670 -0
  49. torchaudio/functional/functional.py +2535 -0
  50. torchaudio/io/__init__.py +13 -0
  51. torchaudio/io/_effector.py +347 -0
  52. torchaudio/io/_playback.py +72 -0
  53. torchaudio/kaldi_io.py +144 -0
  54. torchaudio/lib/__init__.py +0 -0
  55. torchaudio/lib/_torchaudio.so +0 -0
  56. torchaudio/lib/_torchaudio_sox.so +0 -0
  57. torchaudio/lib/libctc_prefix_decoder.so +0 -0
  58. torchaudio/lib/libtorchaudio.so +0 -0
  59. torchaudio/lib/libtorchaudio_sox.so +0 -0
  60. torchaudio/lib/pybind11_prefixctc.so +0 -0
  61. torchaudio/models/__init__.py +85 -0
  62. torchaudio/models/_hdemucs.py +1008 -0
  63. torchaudio/models/conformer.py +293 -0
  64. torchaudio/models/conv_tasnet.py +330 -0
  65. torchaudio/models/decoder/__init__.py +46 -0
  66. torchaudio/models/decoder/_ctc_decoder.py +568 -0
  67. torchaudio/models/decoder/_cuda_ctc_decoder.py +187 -0
  68. torchaudio/models/deepspeech.py +84 -0
  69. torchaudio/models/emformer.py +884 -0
  70. torchaudio/models/rnnt.py +816 -0
  71. torchaudio/models/rnnt_decoder.py +339 -0
  72. torchaudio/models/squim/__init__.py +11 -0
  73. torchaudio/models/squim/objective.py +326 -0
  74. torchaudio/models/squim/subjective.py +150 -0
  75. torchaudio/models/tacotron2.py +1046 -0
  76. torchaudio/models/wav2letter.py +72 -0
  77. torchaudio/models/wav2vec2/__init__.py +45 -0
  78. torchaudio/models/wav2vec2/components.py +1167 -0
  79. torchaudio/models/wav2vec2/model.py +1579 -0
  80. torchaudio/models/wav2vec2/utils/__init__.py +7 -0
  81. torchaudio/models/wav2vec2/utils/import_fairseq.py +213 -0
  82. torchaudio/models/wav2vec2/utils/import_huggingface.py +134 -0
  83. torchaudio/models/wav2vec2/wavlm_attention.py +214 -0
  84. torchaudio/models/wavernn.py +409 -0
  85. torchaudio/pipelines/__init__.py +102 -0
  86. torchaudio/pipelines/_source_separation_pipeline.py +109 -0
  87. torchaudio/pipelines/_squim_pipeline.py +156 -0
  88. torchaudio/pipelines/_tts/__init__.py +16 -0
  89. torchaudio/pipelines/_tts/impl.py +385 -0
  90. torchaudio/pipelines/_tts/interface.py +255 -0
  91. torchaudio/pipelines/_tts/utils.py +228 -0
  92. torchaudio/pipelines/_wav2vec2/__init__.py +0 -0
  93. torchaudio/pipelines/_wav2vec2/aligner.py +87 -0
  94. torchaudio/pipelines/_wav2vec2/impl.py +1699 -0
  95. torchaudio/pipelines/_wav2vec2/utils.py +346 -0
  96. torchaudio/pipelines/rnnt_pipeline.py +380 -0
  97. torchaudio/prototype/__init__.py +0 -0
  98. torchaudio/prototype/datasets/__init__.py +4 -0
  99. torchaudio/prototype/datasets/musan.py +67 -0
  100. torchaudio/prototype/functional/__init__.py +26 -0
  101. torchaudio/prototype/functional/_dsp.py +433 -0
  102. torchaudio/prototype/functional/_rir.py +379 -0
  103. torchaudio/prototype/functional/functional.py +190 -0
  104. torchaudio/prototype/models/__init__.py +36 -0
  105. torchaudio/prototype/models/_conformer_wav2vec2.py +794 -0
  106. torchaudio/prototype/models/_emformer_hubert.py +333 -0
  107. torchaudio/prototype/models/conv_emformer.py +525 -0
  108. torchaudio/prototype/models/hifi_gan.py +336 -0
  109. torchaudio/prototype/models/rnnt.py +711 -0
  110. torchaudio/prototype/models/rnnt_decoder.py +399 -0
  111. torchaudio/prototype/pipelines/__init__.py +12 -0
  112. torchaudio/prototype/pipelines/_vggish/__init__.py +3 -0
  113. torchaudio/prototype/pipelines/_vggish/_vggish_impl.py +233 -0
  114. torchaudio/prototype/pipelines/_vggish/_vggish_pipeline.py +82 -0
  115. torchaudio/prototype/pipelines/hifigan_pipeline.py +228 -0
  116. torchaudio/prototype/pipelines/rnnt_pipeline.py +58 -0
  117. torchaudio/prototype/transforms/__init__.py +9 -0
  118. torchaudio/prototype/transforms/_transforms.py +456 -0
  119. torchaudio/sox_effects/__init__.py +10 -0
  120. torchaudio/sox_effects/sox_effects.py +272 -0
  121. torchaudio/transforms/__init__.py +75 -0
  122. torchaudio/transforms/_multi_channel.py +467 -0
  123. torchaudio/transforms/_transforms.py +2137 -0
  124. torchaudio/utils/__init__.py +11 -0
  125. torchaudio/utils/download.py +89 -0
  126. torchaudio/utils/ffmpeg_utils.py +11 -0
  127. torchaudio/utils/sox_utils.py +99 -0
  128. torchaudio/version.py +2 -0
  129. torchaudio-2.7.0.dist-info/LICENSE +25 -0
  130. torchaudio-2.7.0.dist-info/METADATA +114 -0
  131. torchaudio-2.7.0.dist-info/RECORD +148 -0
  132. torchaudio-2.7.0.dist-info/WHEEL +5 -0
  133. torchaudio-2.7.0.dist-info/top_level.txt +2 -0
  134. torio/__init__.py +8 -0
  135. torio/_extension/__init__.py +13 -0
  136. torio/_extension/utils.py +147 -0
  137. torio/io/__init__.py +9 -0
  138. torio/io/_streaming_media_decoder.py +978 -0
  139. torio/io/_streaming_media_encoder.py +502 -0
  140. torio/lib/__init__.py +0 -0
  141. torio/lib/_torio_ffmpeg4.so +0 -0
  142. torio/lib/_torio_ffmpeg5.so +0 -0
  143. torio/lib/_torio_ffmpeg6.so +0 -0
  144. torio/lib/libtorio_ffmpeg4.so +0 -0
  145. torio/lib/libtorio_ffmpeg5.so +0 -0
  146. torio/lib/libtorio_ffmpeg6.so +0 -0
  147. torio/utils/__init__.py +4 -0
  148. torio/utils/ffmpeg_utils.py +247 -0
@@ -0,0 +1,2535 @@
1
+ # -*- coding: utf-8 -*-
2
+
3
+ import math
4
+ import tempfile
5
+ import warnings
6
+ from collections.abc import Sequence
7
+ from typing import List, Optional, Tuple, Union
8
+
9
+ import torch
10
+ import torchaudio
11
+ from torch import Tensor
12
+ from torchaudio._internal.module_utils import deprecated
13
+
14
+ from .filtering import highpass_biquad, treble_biquad
15
+
16
+ __all__ = [
17
+ "spectrogram",
18
+ "inverse_spectrogram",
19
+ "griffinlim",
20
+ "amplitude_to_DB",
21
+ "DB_to_amplitude",
22
+ "compute_deltas",
23
+ "melscale_fbanks",
24
+ "linear_fbanks",
25
+ "create_dct",
26
+ "compute_deltas",
27
+ "detect_pitch_frequency",
28
+ "DB_to_amplitude",
29
+ "mu_law_encoding",
30
+ "mu_law_decoding",
31
+ "phase_vocoder",
32
+ "mask_along_axis",
33
+ "mask_along_axis_iid",
34
+ "sliding_window_cmn",
35
+ "spectral_centroid",
36
+ "apply_codec",
37
+ "resample",
38
+ "edit_distance",
39
+ "loudness",
40
+ "pitch_shift",
41
+ "rnnt_loss",
42
+ "psd",
43
+ "mvdr_weights_souden",
44
+ "mvdr_weights_rtf",
45
+ "rtf_evd",
46
+ "rtf_power",
47
+ "apply_beamforming",
48
+ "fftconvolve",
49
+ "convolve",
50
+ "add_noise",
51
+ "speed",
52
+ "preemphasis",
53
+ "deemphasis",
54
+ ]
55
+
56
+
57
+ def spectrogram(
58
+ waveform: Tensor,
59
+ pad: int,
60
+ window: Tensor,
61
+ n_fft: int,
62
+ hop_length: int,
63
+ win_length: int,
64
+ power: Optional[float],
65
+ normalized: Union[bool, str],
66
+ center: bool = True,
67
+ pad_mode: str = "reflect",
68
+ onesided: bool = True,
69
+ return_complex: Optional[bool] = None,
70
+ ) -> Tensor:
71
+ r"""Create a spectrogram or a batch of spectrograms from a raw audio signal.
72
+ The spectrogram can be either magnitude-only or complex.
73
+
74
+ .. devices:: CPU CUDA
75
+
76
+ .. properties:: Autograd TorchScript
77
+
78
+ Args:
79
+ waveform (Tensor): Tensor of audio of dimension `(..., time)`
80
+ pad (int): Two sided padding of signal
81
+ window (Tensor): Window tensor that is applied/multiplied to each frame/window
82
+ n_fft (int): Size of FFT
83
+ hop_length (int): Length of hop between STFT windows
84
+ win_length (int): Window size
85
+ power (float or None): Exponent for the magnitude spectrogram,
86
+ (must be > 0) e.g., 1 for magnitude, 2 for power, etc.
87
+ If None, then the complex spectrum is returned instead.
88
+ normalized (bool or str): Whether to normalize by magnitude after stft. If input is str, choices are
89
+ ``"window"`` and ``"frame_length"``, if specific normalization type is desirable. ``True`` maps to
90
+ ``"window"``. When normalized on ``"window"``, waveform is normalized upon the window's L2 energy. If
91
+ normalized on ``"frame_length"``, waveform is normalized by dividing by
92
+ :math:`(\text{frame\_length})^{0.5}`.
93
+ center (bool, optional): whether to pad :attr:`waveform` on both sides so
94
+ that the :math:`t`-th frame is centered at time :math:`t \times \text{hop\_length}`.
95
+ Default: ``True``
96
+ pad_mode (string, optional): controls the padding method used when
97
+ :attr:`center` is ``True``. Default: ``"reflect"``
98
+ onesided (bool, optional): controls whether to return half of results to
99
+ avoid redundancy. Default: ``True``
100
+ return_complex (bool, optional):
101
+ Deprecated and not used.
102
+
103
+ Returns:
104
+ Tensor: Dimension `(..., freq, time)`, freq is
105
+ ``n_fft // 2 + 1`` and ``n_fft`` is the number of
106
+ Fourier bins, and time is the number of window hops (n_frame).
107
+ """
108
+ if return_complex is not None:
109
+ warnings.warn(
110
+ "`return_complex` argument is now deprecated and is not effective."
111
+ "`torchaudio.functional.spectrogram(power=None)` always returns a tensor with "
112
+ "complex dtype. Please remove the argument in the function call."
113
+ )
114
+
115
+ if pad > 0:
116
+ # TODO add "with torch.no_grad():" back when JIT supports it
117
+ waveform = torch.nn.functional.pad(waveform, (pad, pad), "constant")
118
+
119
+ frame_length_norm, window_norm = _get_spec_norms(normalized)
120
+
121
+ # pack batch
122
+ shape = waveform.size()
123
+ waveform = waveform.reshape(-1, shape[-1])
124
+
125
+ # default values are consistent with librosa.core.spectrum._spectrogram
126
+ spec_f = torch.stft(
127
+ input=waveform,
128
+ n_fft=n_fft,
129
+ hop_length=hop_length,
130
+ win_length=win_length,
131
+ window=window,
132
+ center=center,
133
+ pad_mode=pad_mode,
134
+ normalized=frame_length_norm,
135
+ onesided=onesided,
136
+ return_complex=True,
137
+ )
138
+
139
+ # unpack batch
140
+ spec_f = spec_f.reshape(shape[:-1] + spec_f.shape[-2:])
141
+
142
+ if window_norm:
143
+ spec_f /= window.pow(2.0).sum().sqrt()
144
+ if power is not None:
145
+ if power == 1.0:
146
+ return spec_f.abs()
147
+ return spec_f.abs().pow(power)
148
+ return spec_f
149
+
150
+
151
+ def inverse_spectrogram(
152
+ spectrogram: Tensor,
153
+ length: Optional[int],
154
+ pad: int,
155
+ window: Tensor,
156
+ n_fft: int,
157
+ hop_length: int,
158
+ win_length: int,
159
+ normalized: Union[bool, str],
160
+ center: bool = True,
161
+ pad_mode: str = "reflect",
162
+ onesided: bool = True,
163
+ ) -> Tensor:
164
+ r"""Create an inverse spectrogram or a batch of inverse spectrograms from the provided
165
+ complex-valued spectrogram.
166
+
167
+ .. devices:: CPU CUDA
168
+
169
+ .. properties:: Autograd TorchScript
170
+
171
+ Args:
172
+ spectrogram (Tensor): Complex tensor of audio of dimension (..., freq, time).
173
+ length (int or None): The output length of the waveform.
174
+ pad (int): Two sided padding of signal. It is only effective when ``length`` is provided.
175
+ window (Tensor): Window tensor that is applied/multiplied to each frame/window
176
+ n_fft (int): Size of FFT
177
+ hop_length (int): Length of hop between STFT windows
178
+ win_length (int): Window size
179
+ normalized (bool or str): Whether the stft output was normalized by magnitude. If input is str, choices are
180
+ ``"window"`` and ``"frame_length"``, dependent on normalization mode. ``True`` maps to
181
+ ``"window"``.
182
+ center (bool, optional): whether the waveform was padded on both sides so
183
+ that the :math:`t`-th frame is centered at time :math:`t \times \text{hop\_length}`.
184
+ Default: ``True``
185
+ pad_mode (string, optional): controls the padding method used when
186
+ :attr:`center` is ``True``. This parameter is provided for compatibility with the
187
+ spectrogram function and is not used. Default: ``"reflect"``
188
+ onesided (bool, optional): controls whether spectrogram was done in onesided mode.
189
+ Default: ``True``
190
+
191
+ Returns:
192
+ Tensor: Dimension `(..., time)`. Least squares estimation of the original signal.
193
+ """
194
+
195
+ frame_length_norm, window_norm = _get_spec_norms(normalized)
196
+
197
+ if not spectrogram.is_complex():
198
+ raise ValueError("Expected `spectrogram` to be complex dtype.")
199
+
200
+ if window_norm:
201
+ spectrogram = spectrogram * window.pow(2.0).sum().sqrt()
202
+
203
+ # pack batch
204
+ shape = spectrogram.size()
205
+ spectrogram = spectrogram.reshape(-1, shape[-2], shape[-1])
206
+
207
+ # default values are consistent with librosa.core.spectrum._spectrogram
208
+ waveform = torch.istft(
209
+ input=spectrogram,
210
+ n_fft=n_fft,
211
+ hop_length=hop_length,
212
+ win_length=win_length,
213
+ window=window,
214
+ center=center,
215
+ normalized=frame_length_norm,
216
+ onesided=onesided,
217
+ length=length + 2 * pad if length is not None else None,
218
+ return_complex=False,
219
+ )
220
+
221
+ if length is not None and pad > 0:
222
+ # remove padding from front and back
223
+ waveform = waveform[:, pad:-pad]
224
+
225
+ # unpack batch
226
+ waveform = waveform.reshape(shape[:-2] + waveform.shape[-1:])
227
+
228
+ return waveform
229
+
230
+
231
+ def _get_spec_norms(normalized: Union[str, bool]):
232
+ frame_length_norm, window_norm = False, False
233
+ if torch.jit.isinstance(normalized, str):
234
+ if normalized not in ["frame_length", "window"]:
235
+ raise ValueError("Invalid normalized parameter: {}".format(normalized))
236
+ if normalized == "frame_length":
237
+ frame_length_norm = True
238
+ elif normalized == "window":
239
+ window_norm = True
240
+ elif torch.jit.isinstance(normalized, bool):
241
+ if normalized:
242
+ window_norm = True
243
+ else:
244
+ raise TypeError("Input type not supported")
245
+ return frame_length_norm, window_norm
246
+
247
+
248
+ def _get_complex_dtype(real_dtype: torch.dtype):
249
+ if real_dtype == torch.double:
250
+ return torch.cdouble
251
+ if real_dtype == torch.float:
252
+ return torch.cfloat
253
+ if real_dtype == torch.half:
254
+ return torch.complex32
255
+ raise ValueError(f"Unexpected dtype {real_dtype}")
256
+
257
+
258
+ def griffinlim(
259
+ specgram: Tensor,
260
+ window: Tensor,
261
+ n_fft: int,
262
+ hop_length: int,
263
+ win_length: int,
264
+ power: float,
265
+ n_iter: int,
266
+ momentum: float,
267
+ length: Optional[int],
268
+ rand_init: bool,
269
+ ) -> Tensor:
270
+ r"""Compute waveform from a linear scale magnitude spectrogram using the Griffin-Lim transformation.
271
+
272
+ .. devices:: CPU CUDA
273
+
274
+ .. properties:: Autograd TorchScript
275
+
276
+ Implementation ported from
277
+ *librosa* :cite:`brian_mcfee-proc-scipy-2015`, *A fast Griffin-Lim algorithm* :cite:`6701851`
278
+ and *Signal estimation from modified short-time Fourier transform* :cite:`1172092`.
279
+
280
+ Args:
281
+ specgram (Tensor): A magnitude-only STFT spectrogram of dimension `(..., freq, frames)`
282
+ where freq is ``n_fft // 2 + 1``.
283
+ window (Tensor): Window tensor that is applied/multiplied to each frame/window
284
+ n_fft (int): Size of FFT, creates ``n_fft // 2 + 1`` bins
285
+ hop_length (int): Length of hop between STFT windows. (
286
+ Default: ``win_length // 2``)
287
+ win_length (int): Window size. (Default: ``n_fft``)
288
+ power (float): Exponent for the magnitude spectrogram,
289
+ (must be > 0) e.g., 1 for magnitude, 2 for power, etc.
290
+ n_iter (int): Number of iteration for phase recovery process.
291
+ momentum (float): The momentum parameter for fast Griffin-Lim.
292
+ Setting this to 0 recovers the original Griffin-Lim method.
293
+ Values near 1 can lead to faster convergence, but above 1 may not converge.
294
+ length (int or None): Array length of the expected output.
295
+ rand_init (bool): Initializes phase randomly if True, to zero otherwise.
296
+
297
+ Returns:
298
+ Tensor: waveform of `(..., time)`, where time equals the ``length`` parameter if given.
299
+ """
300
+ if not 0 <= momentum < 1:
301
+ raise ValueError("momentum must be in range [0, 1). Found: {}".format(momentum))
302
+
303
+ momentum = momentum / (1 + momentum)
304
+
305
+ # pack batch
306
+ shape = specgram.size()
307
+ specgram = specgram.reshape([-1] + list(shape[-2:]))
308
+
309
+ specgram = specgram.pow(1 / power)
310
+
311
+ # initialize the phase
312
+ if rand_init:
313
+ angles = torch.rand(specgram.size(), dtype=_get_complex_dtype(specgram.dtype), device=specgram.device)
314
+ else:
315
+ angles = torch.full(specgram.size(), 1, dtype=_get_complex_dtype(specgram.dtype), device=specgram.device)
316
+
317
+ # And initialize the previous iterate to 0
318
+ tprev = torch.tensor(0.0, dtype=specgram.dtype, device=specgram.device)
319
+ for _ in range(n_iter):
320
+ # Invert with our current estimate of the phases
321
+ inverse = torch.istft(
322
+ specgram * angles, n_fft=n_fft, hop_length=hop_length, win_length=win_length, window=window, length=length
323
+ )
324
+
325
+ # Rebuild the spectrogram
326
+ rebuilt = torch.stft(
327
+ input=inverse,
328
+ n_fft=n_fft,
329
+ hop_length=hop_length,
330
+ win_length=win_length,
331
+ window=window,
332
+ center=True,
333
+ pad_mode="reflect",
334
+ normalized=False,
335
+ onesided=True,
336
+ return_complex=True,
337
+ )
338
+
339
+ # Update our phase estimates
340
+ angles = rebuilt
341
+ if momentum:
342
+ angles = angles - tprev.mul_(momentum)
343
+ angles = angles.div(angles.abs().add(1e-16))
344
+
345
+ # Store the previous iterate
346
+ tprev = rebuilt
347
+
348
+ # Return the final phase estimates
349
+ waveform = torch.istft(
350
+ specgram * angles, n_fft=n_fft, hop_length=hop_length, win_length=win_length, window=window, length=length
351
+ )
352
+
353
+ # unpack batch
354
+ waveform = waveform.reshape(shape[:-2] + waveform.shape[-1:])
355
+
356
+ return waveform
357
+
358
+
359
+ def amplitude_to_DB(
360
+ x: Tensor, multiplier: float, amin: float, db_multiplier: float, top_db: Optional[float] = None
361
+ ) -> Tensor:
362
+ r"""Turn a spectrogram from the power/amplitude scale to the decibel scale.
363
+
364
+ .. devices:: CPU CUDA
365
+
366
+ .. properties:: Autograd TorchScript
367
+
368
+ The output of each tensor in a batch depends on the maximum value of that tensor,
369
+ and so may return different values for an audio clip split into snippets vs. a full clip.
370
+
371
+ Args:
372
+
373
+ x (Tensor): Input spectrogram(s) before being converted to decibel scale.
374
+ The expected shapes are ``(freq, time)``, ``(channel, freq, time)`` or
375
+ ``(..., batch, channel, freq, time)``.
376
+
377
+ .. note::
378
+
379
+ When ``top_db`` is specified, cut-off values are computed for each audio
380
+ in the batch. Therefore if the input shape is 4D (or larger), different
381
+ cut-off values are used for audio data in the batch.
382
+ If the input shape is 2D or 3D, a single cutoff value is used.
383
+
384
+ multiplier (float): Use 10. for power and 20. for amplitude
385
+ amin (float): Number to clamp ``x``
386
+ db_multiplier (float): Log10(max(reference value and amin))
387
+ top_db (float or None, optional): Minimum negative cut-off in decibels. A reasonable number
388
+ is 80. (Default: ``None``)
389
+
390
+ Returns:
391
+ Tensor: Output tensor in decibel scale
392
+ """
393
+ x_db = multiplier * torch.log10(torch.clamp(x, min=amin))
394
+ x_db -= multiplier * db_multiplier
395
+
396
+ if top_db is not None:
397
+ # Expand batch
398
+ shape = x_db.size()
399
+ packed_channels = shape[-3] if x_db.dim() > 2 else 1
400
+ x_db = x_db.reshape(-1, packed_channels, shape[-2], shape[-1])
401
+
402
+ x_db = torch.max(x_db, (x_db.amax(dim=(-3, -2, -1)) - top_db).view(-1, 1, 1, 1))
403
+
404
+ # Repack batch
405
+ x_db = x_db.reshape(shape)
406
+
407
+ return x_db
408
+
409
+
410
+ def DB_to_amplitude(x: Tensor, ref: float, power: float) -> Tensor:
411
+ r"""Turn a tensor from the decibel scale to the power/amplitude scale.
412
+
413
+ .. devices:: CPU CUDA
414
+
415
+ .. properties:: TorchScript
416
+
417
+ Args:
418
+ x (Tensor): Input tensor before being converted to power/amplitude scale.
419
+ ref (float): Reference which the output will be scaled by.
420
+ power (float): If power equals 1, will compute DB to power. If 0.5, will compute DB to amplitude.
421
+
422
+ Returns:
423
+ Tensor: Output tensor in power/amplitude scale.
424
+ """
425
+ return ref * torch.pow(torch.pow(10.0, 0.1 * x), power)
426
+
427
+
428
+ def _hz_to_mel(freq: float, mel_scale: str = "htk") -> float:
429
+ r"""Convert Hz to Mels.
430
+
431
+ Args:
432
+ freqs (float): Frequencies in Hz
433
+ mel_scale (str, optional): Scale to use: ``htk`` or ``slaney``. (Default: ``htk``)
434
+
435
+ Returns:
436
+ mels (float): Frequency in Mels
437
+ """
438
+
439
+ if mel_scale not in ["slaney", "htk"]:
440
+ raise ValueError('mel_scale should be one of "htk" or "slaney".')
441
+
442
+ if mel_scale == "htk":
443
+ return 2595.0 * math.log10(1.0 + (freq / 700.0))
444
+
445
+ # Fill in the linear part
446
+ f_min = 0.0
447
+ f_sp = 200.0 / 3
448
+
449
+ mels = (freq - f_min) / f_sp
450
+
451
+ # Fill in the log-scale part
452
+ min_log_hz = 1000.0
453
+ min_log_mel = (min_log_hz - f_min) / f_sp
454
+ logstep = math.log(6.4) / 27.0
455
+
456
+ if freq >= min_log_hz:
457
+ mels = min_log_mel + math.log(freq / min_log_hz) / logstep
458
+
459
+ return mels
460
+
461
+
462
+ def _mel_to_hz(mels: Tensor, mel_scale: str = "htk") -> Tensor:
463
+ """Convert mel bin numbers to frequencies.
464
+
465
+ Args:
466
+ mels (Tensor): Mel frequencies
467
+ mel_scale (str, optional): Scale to use: ``htk`` or ``slaney``. (Default: ``htk``)
468
+
469
+ Returns:
470
+ freqs (Tensor): Mels converted in Hz
471
+ """
472
+
473
+ if mel_scale not in ["slaney", "htk"]:
474
+ raise ValueError('mel_scale should be one of "htk" or "slaney".')
475
+
476
+ if mel_scale == "htk":
477
+ return 700.0 * (10.0 ** (mels / 2595.0) - 1.0)
478
+
479
+ # Fill in the linear scale
480
+ f_min = 0.0
481
+ f_sp = 200.0 / 3
482
+ freqs = f_min + f_sp * mels
483
+
484
+ # And now the nonlinear scale
485
+ min_log_hz = 1000.0
486
+ min_log_mel = (min_log_hz - f_min) / f_sp
487
+ logstep = math.log(6.4) / 27.0
488
+
489
+ log_t = mels >= min_log_mel
490
+ freqs[log_t] = min_log_hz * torch.exp(logstep * (mels[log_t] - min_log_mel))
491
+
492
+ return freqs
493
+
494
+
495
+ def _create_triangular_filterbank(
496
+ all_freqs: Tensor,
497
+ f_pts: Tensor,
498
+ ) -> Tensor:
499
+ """Create a triangular filter bank.
500
+
501
+ Args:
502
+ all_freqs (Tensor): STFT freq points of size (`n_freqs`).
503
+ f_pts (Tensor): Filter mid points of size (`n_filter`).
504
+
505
+ Returns:
506
+ fb (Tensor): The filter bank of size (`n_freqs`, `n_filter`).
507
+ """
508
+ # Adopted from Librosa
509
+ # calculate the difference between each filter mid point and each stft freq point in hertz
510
+ f_diff = f_pts[1:] - f_pts[:-1] # (n_filter + 1)
511
+ slopes = f_pts.unsqueeze(0) - all_freqs.unsqueeze(1) # (n_freqs, n_filter + 2)
512
+ # create overlapping triangles
513
+ zero = torch.zeros(1)
514
+ down_slopes = (-1.0 * slopes[:, :-2]) / f_diff[:-1] # (n_freqs, n_filter)
515
+ up_slopes = slopes[:, 2:] / f_diff[1:] # (n_freqs, n_filter)
516
+ fb = torch.max(zero, torch.min(down_slopes, up_slopes))
517
+
518
+ return fb
519
+
520
+
521
+ def melscale_fbanks(
522
+ n_freqs: int,
523
+ f_min: float,
524
+ f_max: float,
525
+ n_mels: int,
526
+ sample_rate: int,
527
+ norm: Optional[str] = None,
528
+ mel_scale: str = "htk",
529
+ ) -> Tensor:
530
+ r"""Create a frequency bin conversion matrix.
531
+
532
+ .. devices:: CPU
533
+
534
+ .. properties:: TorchScript
535
+
536
+ Note:
537
+ For the sake of the numerical compatibility with librosa, not all the coefficients
538
+ in the resulting filter bank has magnitude of 1.
539
+
540
+ .. image:: https://download.pytorch.org/torchaudio/doc-assets/mel_fbanks.png
541
+ :alt: Visualization of generated filter bank
542
+
543
+ Args:
544
+ n_freqs (int): Number of frequencies to highlight/apply
545
+ f_min (float): Minimum frequency (Hz)
546
+ f_max (float): Maximum frequency (Hz)
547
+ n_mels (int): Number of mel filterbanks
548
+ sample_rate (int): Sample rate of the audio waveform
549
+ norm (str or None, optional): If "slaney", divide the triangular mel weights by the width of the mel band
550
+ (area normalization). (Default: ``None``)
551
+ mel_scale (str, optional): Scale to use: ``htk`` or ``slaney``. (Default: ``htk``)
552
+
553
+ Returns:
554
+ Tensor: Triangular filter banks (fb matrix) of size (``n_freqs``, ``n_mels``)
555
+ meaning number of frequencies to highlight/apply to x the number of filterbanks.
556
+ Each column is a filterbank so that assuming there is a matrix A of
557
+ size (..., ``n_freqs``), the applied result would be
558
+ ``A @ melscale_fbanks(A.size(-1), ...)``.
559
+
560
+ """
561
+
562
+ if norm is not None and norm != "slaney":
563
+ raise ValueError('norm must be one of None or "slaney"')
564
+
565
+ # freq bins
566
+ all_freqs = torch.linspace(0, sample_rate // 2, n_freqs)
567
+
568
+ # calculate mel freq bins
569
+ m_min = _hz_to_mel(f_min, mel_scale=mel_scale)
570
+ m_max = _hz_to_mel(f_max, mel_scale=mel_scale)
571
+
572
+ m_pts = torch.linspace(m_min, m_max, n_mels + 2)
573
+ f_pts = _mel_to_hz(m_pts, mel_scale=mel_scale)
574
+
575
+ # create filterbank
576
+ fb = _create_triangular_filterbank(all_freqs, f_pts)
577
+
578
+ if norm is not None and norm == "slaney":
579
+ # Slaney-style mel is scaled to be approx constant energy per channel
580
+ enorm = 2.0 / (f_pts[2 : n_mels + 2] - f_pts[:n_mels])
581
+ fb *= enorm.unsqueeze(0)
582
+
583
+ if (fb.max(dim=0).values == 0.0).any():
584
+ warnings.warn(
585
+ "At least one mel filterbank has all zero values. "
586
+ f"The value for `n_mels` ({n_mels}) may be set too high. "
587
+ f"Or, the value for `n_freqs` ({n_freqs}) may be set too low."
588
+ )
589
+
590
+ return fb
591
+
592
+
593
+ def linear_fbanks(
594
+ n_freqs: int,
595
+ f_min: float,
596
+ f_max: float,
597
+ n_filter: int,
598
+ sample_rate: int,
599
+ ) -> Tensor:
600
+ r"""Creates a linear triangular filterbank.
601
+
602
+ .. devices:: CPU
603
+
604
+ .. properties:: TorchScript
605
+
606
+ Note:
607
+ For the sake of the numerical compatibility with librosa, not all the coefficients
608
+ in the resulting filter bank has magnitude of 1.
609
+
610
+ .. image:: https://download.pytorch.org/torchaudio/doc-assets/lin_fbanks.png
611
+ :alt: Visualization of generated filter bank
612
+
613
+ Args:
614
+ n_freqs (int): Number of frequencies to highlight/apply
615
+ f_min (float): Minimum frequency (Hz)
616
+ f_max (float): Maximum frequency (Hz)
617
+ n_filter (int): Number of (linear) triangular filter
618
+ sample_rate (int): Sample rate of the audio waveform
619
+
620
+ Returns:
621
+ Tensor: Triangular filter banks (fb matrix) of size (``n_freqs``, ``n_filter``)
622
+ meaning number of frequencies to highlight/apply to x the number of filterbanks.
623
+ Each column is a filterbank so that assuming there is a matrix A of
624
+ size (..., ``n_freqs``), the applied result would be
625
+ ``A * linear_fbanks(A.size(-1), ...)``.
626
+ """
627
+ # freq bins
628
+ all_freqs = torch.linspace(0, sample_rate // 2, n_freqs)
629
+
630
+ # filter mid-points
631
+ f_pts = torch.linspace(f_min, f_max, n_filter + 2)
632
+
633
+ # create filterbank
634
+ fb = _create_triangular_filterbank(all_freqs, f_pts)
635
+
636
+ return fb
637
+
638
+
639
+ def create_dct(n_mfcc: int, n_mels: int, norm: Optional[str]) -> Tensor:
640
+ r"""Create a DCT transformation matrix with shape (``n_mels``, ``n_mfcc``),
641
+ normalized depending on norm.
642
+
643
+ .. devices:: CPU
644
+
645
+ .. properties:: TorchScript
646
+
647
+ Args:
648
+ n_mfcc (int): Number of mfc coefficients to retain
649
+ n_mels (int): Number of mel filterbanks
650
+ norm (str or None): Norm to use (either "ortho" or None)
651
+
652
+ Returns:
653
+ Tensor: The transformation matrix, to be right-multiplied to
654
+ row-wise data of size (``n_mels``, ``n_mfcc``).
655
+ """
656
+
657
+ if norm is not None and norm != "ortho":
658
+ raise ValueError('norm must be either "ortho" or None')
659
+
660
+ # http://en.wikipedia.org/wiki/Discrete_cosine_transform#DCT-II
661
+ n = torch.arange(float(n_mels))
662
+ k = torch.arange(float(n_mfcc)).unsqueeze(1)
663
+ dct = torch.cos(math.pi / float(n_mels) * (n + 0.5) * k) # size (n_mfcc, n_mels)
664
+
665
+ if norm is None:
666
+ dct *= 2.0
667
+ else:
668
+ dct[0] *= 1.0 / math.sqrt(2.0)
669
+ dct *= math.sqrt(2.0 / float(n_mels))
670
+ return dct.t()
671
+
672
+
673
+ def mu_law_encoding(x: Tensor, quantization_channels: int) -> Tensor:
674
+ r"""Encode signal based on mu-law companding.
675
+
676
+ .. devices:: CPU CUDA
677
+
678
+ .. properties:: TorchScript
679
+
680
+ For more info see the
681
+ `Wikipedia Entry <https://en.wikipedia.org/wiki/%CE%9C-law_algorithm>`_
682
+
683
+ This algorithm expects the signal has been scaled to between -1 and 1 and
684
+ returns a signal encoded with values from 0 to quantization_channels - 1.
685
+
686
+ Args:
687
+ x (Tensor): Input tensor
688
+ quantization_channels (int): Number of channels
689
+
690
+ Returns:
691
+ Tensor: Input after mu-law encoding
692
+ """
693
+ mu = quantization_channels - 1.0
694
+ if not x.is_floating_point():
695
+ warnings.warn(
696
+ "The input Tensor must be of floating type. \
697
+ This will be an error in the v0.12 release."
698
+ )
699
+ x = x.to(torch.float)
700
+ mu = torch.tensor(mu, dtype=x.dtype)
701
+ x_mu = torch.sign(x) * torch.log1p(mu * torch.abs(x)) / torch.log1p(mu)
702
+ x_mu = ((x_mu + 1) / 2 * mu + 0.5).to(torch.int64)
703
+ return x_mu
704
+
705
+
706
+ def mu_law_decoding(x_mu: Tensor, quantization_channels: int) -> Tensor:
707
+ r"""Decode mu-law encoded signal.
708
+
709
+ .. devices:: CPU CUDA
710
+
711
+ .. properties:: TorchScript
712
+
713
+ For more info see the
714
+ `Wikipedia Entry <https://en.wikipedia.org/wiki/%CE%9C-law_algorithm>`_
715
+
716
+ This expects an input with values between 0 and quantization_channels - 1
717
+ and returns a signal scaled between -1 and 1.
718
+
719
+ Args:
720
+ x_mu (Tensor): Input tensor
721
+ quantization_channels (int): Number of channels
722
+
723
+ Returns:
724
+ Tensor: Input after mu-law decoding
725
+ """
726
+ mu = quantization_channels - 1.0
727
+ if not x_mu.is_floating_point():
728
+ x_mu = x_mu.to(torch.float)
729
+ mu = torch.tensor(mu, dtype=x_mu.dtype)
730
+ x = ((x_mu) / mu) * 2 - 1.0
731
+ x = torch.sign(x) * (torch.exp(torch.abs(x) * torch.log1p(mu)) - 1.0) / mu
732
+ return x
733
+
734
+
735
+ def phase_vocoder(complex_specgrams: Tensor, rate: float, phase_advance: Tensor) -> Tensor:
736
+ r"""Given a STFT tensor, speed up in time without modifying pitch by a factor of ``rate``.
737
+
738
+ .. devices:: CPU CUDA
739
+
740
+ .. properties:: Autograd TorchScript
741
+
742
+ Args:
743
+ complex_specgrams (Tensor):
744
+ A tensor of dimension `(..., freq, num_frame)` with complex dtype.
745
+ rate (float): Speed-up factor
746
+ phase_advance (Tensor): Expected phase advance in each bin. Dimension of `(freq, 1)`
747
+
748
+ Returns:
749
+ Tensor:
750
+ Stretched spectrogram. The resulting tensor is of the same dtype as the input
751
+ spectrogram, but the number of frames is changed to ``ceil(num_frame / rate)``.
752
+
753
+ Example
754
+ >>> freq, hop_length = 1025, 512
755
+ >>> # (channel, freq, time)
756
+ >>> complex_specgrams = torch.randn(2, freq, 300, dtype=torch.cfloat)
757
+ >>> rate = 1.3 # Speed up by 30%
758
+ >>> phase_advance = torch.linspace(
759
+ >>> 0, math.pi * hop_length, freq)[..., None]
760
+ >>> x = phase_vocoder(complex_specgrams, rate, phase_advance)
761
+ >>> x.shape # with 231 == ceil(300 / 1.3)
762
+ torch.Size([2, 1025, 231])
763
+ """
764
+ if rate == 1.0:
765
+ return complex_specgrams
766
+
767
+ # pack batch
768
+ shape = complex_specgrams.size()
769
+ complex_specgrams = complex_specgrams.reshape([-1] + list(shape[-2:]))
770
+
771
+ # Figures out the corresponding real dtype, i.e. complex128 -> float64, complex64 -> float32
772
+ # Note torch.real is a view so it does not incur any memory copy.
773
+ real_dtype = torch.real(complex_specgrams).dtype
774
+ time_steps = torch.arange(0, complex_specgrams.size(-1), rate, device=complex_specgrams.device, dtype=real_dtype)
775
+
776
+ alphas = time_steps % 1.0
777
+ phase_0 = complex_specgrams[..., :1].angle()
778
+
779
+ # Time Padding
780
+ complex_specgrams = torch.nn.functional.pad(complex_specgrams, [0, 2])
781
+
782
+ # (new_bins, freq, 2)
783
+ complex_specgrams_0 = complex_specgrams.index_select(-1, time_steps.long())
784
+ complex_specgrams_1 = complex_specgrams.index_select(-1, (time_steps + 1).long())
785
+
786
+ angle_0 = complex_specgrams_0.angle()
787
+ angle_1 = complex_specgrams_1.angle()
788
+
789
+ norm_0 = complex_specgrams_0.abs()
790
+ norm_1 = complex_specgrams_1.abs()
791
+
792
+ phase = angle_1 - angle_0 - phase_advance
793
+ phase = phase - 2 * math.pi * torch.round(phase / (2 * math.pi))
794
+
795
+ # Compute Phase Accum
796
+ phase = phase + phase_advance
797
+ phase = torch.cat([phase_0, phase[..., :-1]], dim=-1)
798
+ phase_acc = torch.cumsum(phase, -1)
799
+
800
+ mag = alphas * norm_1 + (1 - alphas) * norm_0
801
+
802
+ complex_specgrams_stretch = torch.polar(mag, phase_acc)
803
+
804
+ # unpack batch
805
+ complex_specgrams_stretch = complex_specgrams_stretch.reshape(shape[:-2] + complex_specgrams_stretch.shape[1:])
806
+ return complex_specgrams_stretch
807
+
808
+
809
+ def _get_mask_param(mask_param: int, p: float, axis_length: int) -> int:
810
+ if p == 1.0:
811
+ return mask_param
812
+ else:
813
+ return min(mask_param, int(axis_length * p))
814
+
815
+
816
+ def mask_along_axis_iid(
817
+ specgrams: Tensor,
818
+ mask_param: int,
819
+ mask_value: float,
820
+ axis: int,
821
+ p: float = 1.0,
822
+ ) -> Tensor:
823
+ r"""Apply a mask along ``axis``.
824
+
825
+ .. devices:: CPU CUDA
826
+
827
+ .. properties:: Autograd TorchScript
828
+
829
+ Mask will be applied from indices ``[v_0, v_0 + v)``,
830
+ where ``v`` is sampled from ``uniform(0, max_v)`` and
831
+ ``v_0`` from ``uniform(0, specgrams.size(axis) - v)``,
832
+ with ``max_v = mask_param`` when ``p = 1.0`` and
833
+ ``max_v = min(mask_param, floor(specgrams.size(axis) * p))`` otherwise.
834
+
835
+ Args:
836
+ specgrams (Tensor): Real spectrograms `(..., freq, time)`, with at least 3 dimensions.
837
+ mask_param (int): Number of columns to be masked will be uniformly sampled from [0, mask_param]
838
+ mask_value (float): Value to assign to the masked columns
839
+ axis (int): Axis to apply masking on, which should be the one of the last two dimensions.
840
+ p (float, optional): maximum proportion of columns that can be masked. (Default: 1.0)
841
+
842
+ Returns:
843
+ Tensor: Masked spectrograms with the same dimensions as input specgrams Tensor`
844
+ """
845
+
846
+ dim = specgrams.dim()
847
+
848
+ if dim < 3:
849
+ raise ValueError(f"Spectrogram must have at least three dimensions ({dim} given).")
850
+
851
+ if axis not in [dim - 2, dim - 1]:
852
+ raise ValueError(
853
+ f"Only Frequency and Time masking are supported (axis {dim-2} and axis {dim-1} supported; {axis} given)."
854
+ )
855
+
856
+ if not 0.0 <= p <= 1.0:
857
+ raise ValueError(f"The value of p must be between 0.0 and 1.0 ({p} given).")
858
+
859
+ mask_param = _get_mask_param(mask_param, p, specgrams.shape[axis])
860
+ if mask_param < 1:
861
+ return specgrams
862
+
863
+ device = specgrams.device
864
+ dtype = specgrams.dtype
865
+
866
+ value = torch.rand(specgrams.shape[: (dim - 2)], device=device, dtype=dtype) * mask_param
867
+ min_value = torch.rand(specgrams.shape[: (dim - 2)], device=device, dtype=dtype) * (specgrams.size(axis) - value)
868
+
869
+ # Create broadcastable mask
870
+ mask_start = min_value.long()[..., None, None]
871
+ mask_end = (min_value.long() + value.long())[..., None, None]
872
+ mask = torch.arange(0, specgrams.size(axis), device=device, dtype=dtype)
873
+
874
+ # Per batch example masking
875
+ specgrams = specgrams.transpose(axis, -1)
876
+ specgrams = specgrams.masked_fill((mask >= mask_start) & (mask < mask_end), mask_value)
877
+ specgrams = specgrams.transpose(axis, -1)
878
+
879
+ return specgrams
880
+
881
+
882
+ def mask_along_axis(
883
+ specgram: Tensor,
884
+ mask_param: int,
885
+ mask_value: float,
886
+ axis: int,
887
+ p: float = 1.0,
888
+ ) -> Tensor:
889
+ r"""Apply a mask along ``axis``.
890
+
891
+ .. devices:: CPU CUDA
892
+
893
+ .. properties:: Autograd TorchScript
894
+
895
+ Mask will be applied from indices ``[v_0, v_0 + v)``,
896
+ where ``v`` is sampled from ``uniform(0, max_v)`` and
897
+ ``v_0`` from ``uniform(0, specgram.size(axis) - v)``, with
898
+ ``max_v = mask_param`` when ``p = 1.0`` and
899
+ ``max_v = min(mask_param, floor(specgram.size(axis) * p))``
900
+ otherwise.
901
+ All examples will have the same mask interval.
902
+
903
+ Args:
904
+ specgram (Tensor): Real spectrograms `(..., freq, time)`, with at least 2 dimensions.
905
+ mask_param (int): Number of columns to be masked will be uniformly sampled from [0, mask_param]
906
+ mask_value (float): Value to assign to the masked columns
907
+ axis (int): Axis to apply masking on, which should be the one of the last two dimensions.
908
+ p (float, optional): maximum proportion of columns that can be masked. (Default: 1.0)
909
+
910
+ Returns:
911
+ Tensor: Masked spectrograms with the same dimensions as input specgram Tensor
912
+ """
913
+ dim = specgram.dim()
914
+
915
+ if dim < 2:
916
+ raise ValueError(f"Spectrogram must have at least two dimensions (time and frequency) ({dim} given).")
917
+
918
+ if axis not in [dim - 2, dim - 1]:
919
+ raise ValueError(
920
+ f"Only Frequency and Time masking are supported (axis {dim-2} and axis {dim-1} supported; {axis} given)."
921
+ )
922
+
923
+ if not 0.0 <= p <= 1.0:
924
+ raise ValueError(f"The value of p must be between 0.0 and 1.0 ({p} given).")
925
+
926
+ mask_param = _get_mask_param(mask_param, p, specgram.shape[axis])
927
+ if mask_param < 1:
928
+ return specgram
929
+
930
+ # pack batch
931
+ shape = specgram.size()
932
+ specgram = specgram.reshape([-1] + list(shape[-2:]))
933
+ # After packing, specgram is a 3D tensor, and the axis corresponding to the to-be-masked dimension
934
+ # is now (axis - dim + 3), e.g. a tensor of shape (10, 2, 50, 10, 2) becomes a tensor of shape (1000, 10, 2).
935
+ value = torch.rand(1) * mask_param
936
+ min_value = torch.rand(1) * (specgram.size(axis - dim + 3) - value)
937
+
938
+ mask_start = (min_value.long()).squeeze()
939
+ mask_end = (min_value.long() + value.long()).squeeze()
940
+ mask = torch.arange(0, specgram.shape[axis - dim + 3], device=specgram.device, dtype=specgram.dtype)
941
+ mask = (mask >= mask_start) & (mask < mask_end)
942
+ # unsqueeze the mask if the axis is frequency
943
+ if axis == dim - 2:
944
+ mask = mask.unsqueeze(-1)
945
+
946
+ if mask_end - mask_start >= mask_param:
947
+ raise ValueError("Number of columns to be masked should be less than mask_param")
948
+
949
+ specgram = specgram.masked_fill(mask, mask_value)
950
+
951
+ # unpack batch
952
+ specgram = specgram.reshape(shape[:-2] + specgram.shape[-2:])
953
+
954
+ return specgram
955
+
956
+
957
+ def compute_deltas(specgram: Tensor, win_length: int = 5, mode: str = "replicate") -> Tensor:
958
+ r"""Compute delta coefficients of a tensor, usually a spectrogram:
959
+
960
+ .. devices:: CPU CUDA
961
+
962
+ .. properties:: TorchScript
963
+
964
+ .. math::
965
+ d_t = \frac{\sum_{n=1}^{\text{N}} n (c_{t+n} - c_{t-n})}{2 \sum_{n=1}^{\text{N}} n^2}
966
+
967
+ where :math:`d_t` is the deltas at time :math:`t`,
968
+ :math:`c_t` is the spectrogram coeffcients at time :math:`t`,
969
+ :math:`N` is ``(win_length-1)//2``.
970
+
971
+ Args:
972
+ specgram (Tensor): Tensor of audio of dimension `(..., freq, time)`
973
+ win_length (int, optional): The window length used for computing delta (Default: ``5``)
974
+ mode (str, optional): Mode parameter passed to padding (Default: ``"replicate"``)
975
+
976
+ Returns:
977
+ Tensor: Tensor of deltas of dimension `(..., freq, time)`
978
+
979
+ Example
980
+ >>> specgram = torch.randn(1, 40, 1000)
981
+ >>> delta = compute_deltas(specgram)
982
+ >>> delta2 = compute_deltas(delta)
983
+ """
984
+ device = specgram.device
985
+ dtype = specgram.dtype
986
+
987
+ # pack batch
988
+ shape = specgram.size()
989
+ specgram = specgram.reshape(1, -1, shape[-1])
990
+
991
+ if win_length < 3:
992
+ raise ValueError(f"Window length should be greater than or equal to 3. Found win_length {win_length}")
993
+
994
+ n = (win_length - 1) // 2
995
+
996
+ # twice sum of integer squared
997
+ denom = n * (n + 1) * (2 * n + 1) / 3
998
+
999
+ specgram = torch.nn.functional.pad(specgram, (n, n), mode=mode)
1000
+
1001
+ kernel = torch.arange(-n, n + 1, 1, device=device, dtype=dtype).repeat(specgram.shape[1], 1, 1)
1002
+
1003
+ output = torch.nn.functional.conv1d(specgram, kernel, groups=specgram.shape[1]) / denom
1004
+
1005
+ # unpack batch
1006
+ output = output.reshape(shape)
1007
+
1008
+ return output
1009
+
1010
+
1011
+ def _compute_nccf(waveform: Tensor, sample_rate: int, frame_time: float, freq_low: int) -> Tensor:
1012
+ r"""
1013
+ Compute Normalized Cross-Correlation Function (NCCF).
1014
+
1015
+ .. math::
1016
+ \phi_i(m) = \frac{\sum_{n=b_i}^{b_i + N-1} w(n) w(m+n)}{\sqrt{E(b_i) E(m+b_i)}},
1017
+
1018
+ where
1019
+ :math:`\phi_i(m)` is the NCCF at frame :math:`i` with lag :math:`m`,
1020
+ :math:`w` is the waveform,
1021
+ :math:`N` is the length of a frame,
1022
+ :math:`b_i` is the beginning of frame :math:`i`,
1023
+ :math:`E(j)` is the energy :math:`\sum_{n=j}^{j+N-1} w^2(n)`.
1024
+ """
1025
+
1026
+ EPSILON = 10 ** (-9)
1027
+
1028
+ # Number of lags to check
1029
+ lags = int(math.ceil(sample_rate / freq_low))
1030
+
1031
+ frame_size = int(math.ceil(sample_rate * frame_time))
1032
+
1033
+ waveform_length = waveform.size()[-1]
1034
+ num_of_frames = int(math.ceil(waveform_length / frame_size))
1035
+
1036
+ p = lags + num_of_frames * frame_size - waveform_length
1037
+ waveform = torch.nn.functional.pad(waveform, (0, p))
1038
+
1039
+ # Compute lags
1040
+ output_lag = []
1041
+ for lag in range(1, lags + 1):
1042
+ s1 = waveform[..., :-lag].unfold(-1, frame_size, frame_size)[..., :num_of_frames, :]
1043
+ s2 = waveform[..., lag:].unfold(-1, frame_size, frame_size)[..., :num_of_frames, :]
1044
+
1045
+ output_frames = (
1046
+ (s1 * s2).sum(-1)
1047
+ / (EPSILON + torch.linalg.vector_norm(s1, ord=2, dim=-1)).pow(2)
1048
+ / (EPSILON + torch.linalg.vector_norm(s2, ord=2, dim=-1)).pow(2)
1049
+ )
1050
+
1051
+ output_lag.append(output_frames.unsqueeze(-1))
1052
+
1053
+ nccf = torch.cat(output_lag, -1)
1054
+
1055
+ return nccf
1056
+
1057
+
1058
+ def _combine_max(a: Tuple[Tensor, Tensor], b: Tuple[Tensor, Tensor], thresh: float = 0.99) -> Tuple[Tensor, Tensor]:
1059
+ """
1060
+ Take value from first if bigger than a multiplicative factor of the second, elementwise.
1061
+ """
1062
+ mask = a[0] > thresh * b[0]
1063
+ values = mask * a[0] + ~mask * b[0]
1064
+ indices = mask * a[1] + ~mask * b[1]
1065
+ return values, indices
1066
+
1067
+
1068
+ def _find_max_per_frame(nccf: Tensor, sample_rate: int, freq_high: int) -> Tensor:
1069
+ r"""
1070
+ For each frame, take the highest value of NCCF,
1071
+ apply centered median smoothing, and convert to frequency.
1072
+
1073
+ Note: If the max among all the lags is very close
1074
+ to the first half of lags, then the latter is taken.
1075
+ """
1076
+
1077
+ lag_min = int(math.ceil(sample_rate / freq_high))
1078
+
1079
+ # Find near enough max that is smallest
1080
+
1081
+ best = torch.max(nccf[..., lag_min:], -1)
1082
+
1083
+ half_size = nccf.shape[-1] // 2
1084
+ half = torch.max(nccf[..., lag_min:half_size], -1)
1085
+
1086
+ best = _combine_max(half, best)
1087
+ indices = best[1]
1088
+
1089
+ # Add back minimal lag
1090
+ indices += lag_min
1091
+ # Add 1 empirical calibration offset
1092
+ indices += 1
1093
+
1094
+ return indices
1095
+
1096
+
1097
+ def _median_smoothing(indices: Tensor, win_length: int) -> Tensor:
1098
+ r"""
1099
+ Apply median smoothing to the 1D tensor over the given window.
1100
+ """
1101
+
1102
+ # Centered windowed
1103
+ pad_length = (win_length - 1) // 2
1104
+
1105
+ # "replicate" padding in any dimension
1106
+ indices = torch.nn.functional.pad(indices, (pad_length, 0), mode="constant", value=0.0)
1107
+
1108
+ indices[..., :pad_length] = torch.cat(pad_length * [indices[..., pad_length].unsqueeze(-1)], dim=-1)
1109
+ roll = indices.unfold(-1, win_length, 1)
1110
+
1111
+ values, _ = torch.median(roll, -1)
1112
+ return values
1113
+
1114
+
1115
+ def detect_pitch_frequency(
1116
+ waveform: Tensor,
1117
+ sample_rate: int,
1118
+ frame_time: float = 10 ** (-2),
1119
+ win_length: int = 30,
1120
+ freq_low: int = 85,
1121
+ freq_high: int = 3400,
1122
+ ) -> Tensor:
1123
+ r"""Detect pitch frequency.
1124
+
1125
+ .. devices:: CPU CUDA
1126
+
1127
+ .. properties:: TorchScript
1128
+
1129
+ It is implemented using normalized cross-correlation function and median smoothing.
1130
+
1131
+ Args:
1132
+ waveform (Tensor): Tensor of audio of dimension `(..., freq, time)`
1133
+ sample_rate (int): The sample rate of the waveform (Hz)
1134
+ frame_time (float, optional): Duration of a frame (Default: ``10 ** (-2)``).
1135
+ win_length (int, optional): The window length for median smoothing (in number of frames) (Default: ``30``).
1136
+ freq_low (int, optional): Lowest frequency that can be detected (Hz) (Default: ``85``).
1137
+ freq_high (int, optional): Highest frequency that can be detected (Hz) (Default: ``3400``).
1138
+
1139
+ Returns:
1140
+ Tensor: Tensor of freq of dimension `(..., frame)`
1141
+ """
1142
+ # pack batch
1143
+ shape = list(waveform.size())
1144
+ waveform = waveform.reshape([-1] + shape[-1:])
1145
+
1146
+ nccf = _compute_nccf(waveform, sample_rate, frame_time, freq_low)
1147
+ indices = _find_max_per_frame(nccf, sample_rate, freq_high)
1148
+ indices = _median_smoothing(indices, win_length)
1149
+
1150
+ # Convert indices to frequency
1151
+ EPSILON = 10 ** (-9)
1152
+ freq = sample_rate / (EPSILON + indices.to(torch.float))
1153
+
1154
+ # unpack batch
1155
+ freq = freq.reshape(shape[:-1] + list(freq.shape[-1:]))
1156
+
1157
+ return freq
1158
+
1159
+
1160
+ def sliding_window_cmn(
1161
+ specgram: Tensor,
1162
+ cmn_window: int = 600,
1163
+ min_cmn_window: int = 100,
1164
+ center: bool = False,
1165
+ norm_vars: bool = False,
1166
+ ) -> Tensor:
1167
+ r"""
1168
+ Apply sliding-window cepstral mean (and optionally variance) normalization per utterance.
1169
+
1170
+ .. devices:: CPU CUDA
1171
+
1172
+ .. properties:: TorchScript
1173
+
1174
+ Args:
1175
+ specgram (Tensor): Tensor of spectrogram of dimension `(..., time, freq)`
1176
+ cmn_window (int, optional): Window in frames for running average CMN computation (int, default = 600)
1177
+ min_cmn_window (int, optional): Minimum CMN window used at start of decoding (adds latency only at start).
1178
+ Only applicable if center == false, ignored if center==true (int, default = 100)
1179
+ center (bool, optional): If true, use a window centered on the current frame
1180
+ (to the extent possible, modulo end effects). If false, window is to the left. (bool, default = false)
1181
+ norm_vars (bool, optional): If true, normalize variance to one. (bool, default = false)
1182
+
1183
+ Returns:
1184
+ Tensor: Tensor matching input shape `(..., freq, time)`
1185
+ """
1186
+ input_shape = specgram.shape
1187
+ num_frames, num_feats = input_shape[-2:]
1188
+ specgram = specgram.view(-1, num_frames, num_feats)
1189
+ num_channels = specgram.shape[0]
1190
+
1191
+ dtype = specgram.dtype
1192
+ device = specgram.device
1193
+ last_window_start = last_window_end = -1
1194
+ cur_sum = torch.zeros(num_channels, num_feats, dtype=dtype, device=device)
1195
+ cur_sumsq = torch.zeros(num_channels, num_feats, dtype=dtype, device=device)
1196
+ cmn_specgram = torch.zeros(num_channels, num_frames, num_feats, dtype=dtype, device=device)
1197
+ for t in range(num_frames):
1198
+ window_start = 0
1199
+ window_end = 0
1200
+ if center:
1201
+ window_start = t - cmn_window // 2
1202
+ window_end = window_start + cmn_window
1203
+ else:
1204
+ window_start = t - cmn_window
1205
+ window_end = t + 1
1206
+ if window_start < 0:
1207
+ window_end -= window_start
1208
+ window_start = 0
1209
+ if not center:
1210
+ if window_end > t:
1211
+ window_end = max(t + 1, min_cmn_window)
1212
+ if window_end > num_frames:
1213
+ window_start -= window_end - num_frames
1214
+ window_end = num_frames
1215
+ if window_start < 0:
1216
+ window_start = 0
1217
+ if last_window_start == -1:
1218
+ input_part = specgram[:, window_start : window_end - window_start, :]
1219
+ cur_sum += torch.sum(input_part, 1)
1220
+ if norm_vars:
1221
+ cur_sumsq += torch.cumsum(input_part**2, 1)[:, -1, :]
1222
+ else:
1223
+ if window_start > last_window_start:
1224
+ frame_to_remove = specgram[:, last_window_start, :]
1225
+ cur_sum -= frame_to_remove
1226
+ if norm_vars:
1227
+ cur_sumsq -= frame_to_remove**2
1228
+ if window_end > last_window_end:
1229
+ frame_to_add = specgram[:, last_window_end, :]
1230
+ cur_sum += frame_to_add
1231
+ if norm_vars:
1232
+ cur_sumsq += frame_to_add**2
1233
+ window_frames = window_end - window_start
1234
+ last_window_start = window_start
1235
+ last_window_end = window_end
1236
+ cmn_specgram[:, t, :] = specgram[:, t, :] - cur_sum / window_frames
1237
+ if norm_vars:
1238
+ if window_frames == 1:
1239
+ cmn_specgram[:, t, :] = torch.zeros(num_channels, num_feats, dtype=dtype, device=device)
1240
+ else:
1241
+ variance = cur_sumsq
1242
+ variance = variance / window_frames
1243
+ variance -= (cur_sum**2) / (window_frames**2)
1244
+ variance = torch.pow(variance, -0.5)
1245
+ cmn_specgram[:, t, :] *= variance
1246
+
1247
+ cmn_specgram = cmn_specgram.view(input_shape[:-2] + (num_frames, num_feats))
1248
+ if len(input_shape) == 2:
1249
+ cmn_specgram = cmn_specgram.squeeze(0)
1250
+ return cmn_specgram
1251
+
1252
+
1253
+ def spectral_centroid(
1254
+ waveform: Tensor,
1255
+ sample_rate: int,
1256
+ pad: int,
1257
+ window: Tensor,
1258
+ n_fft: int,
1259
+ hop_length: int,
1260
+ win_length: int,
1261
+ ) -> Tensor:
1262
+ r"""Compute the spectral centroid for each channel along the time axis.
1263
+
1264
+ .. devices:: CPU CUDA
1265
+
1266
+ .. properties:: Autograd TorchScript
1267
+
1268
+ The spectral centroid is defined as the weighted average of the
1269
+ frequency values, weighted by their magnitude.
1270
+
1271
+ Args:
1272
+ waveform (Tensor): Tensor of audio of dimension `(..., time)`
1273
+ sample_rate (int): Sample rate of the audio waveform
1274
+ pad (int): Two sided padding of signal
1275
+ window (Tensor): Window tensor that is applied/multiplied to each frame/window
1276
+ n_fft (int): Size of FFT
1277
+ hop_length (int): Length of hop between STFT windows
1278
+ win_length (int): Window size
1279
+
1280
+ Returns:
1281
+ Tensor: Dimension `(..., time)`
1282
+ """
1283
+ specgram = spectrogram(
1284
+ waveform,
1285
+ pad=pad,
1286
+ window=window,
1287
+ n_fft=n_fft,
1288
+ hop_length=hop_length,
1289
+ win_length=win_length,
1290
+ power=1.0,
1291
+ normalized=False,
1292
+ )
1293
+ freqs = torch.linspace(0, sample_rate // 2, steps=1 + n_fft // 2, device=specgram.device).reshape((-1, 1))
1294
+ freq_dim = -2
1295
+ return (freqs * specgram).sum(dim=freq_dim) / specgram.sum(dim=freq_dim)
1296
+
1297
+
1298
+ @deprecated("Please migrate to :py:class:`torchaudio.io.AudioEffector`.", remove=False)
1299
+ def apply_codec(
1300
+ waveform: Tensor,
1301
+ sample_rate: int,
1302
+ format: str,
1303
+ channels_first: bool = True,
1304
+ compression: Optional[float] = None,
1305
+ encoding: Optional[str] = None,
1306
+ bits_per_sample: Optional[int] = None,
1307
+ ) -> Tensor:
1308
+ r"""
1309
+ Apply codecs as a form of augmentation.
1310
+
1311
+ .. devices:: CPU
1312
+
1313
+ Args:
1314
+ waveform (Tensor): Audio data. Must be 2 dimensional. See also ```channels_first```.
1315
+ sample_rate (int): Sample rate of the audio waveform.
1316
+ format (str): File format.
1317
+ channels_first (bool, optional):
1318
+ When True, both the input and output Tensor have dimension `(channel, time)`.
1319
+ Otherwise, they have dimension `(time, channel)`.
1320
+ compression (float or None, optional): Used for formats other than WAV.
1321
+ For more details see :py:func:`torchaudio.backend.sox_io_backend.save`.
1322
+ encoding (str or None, optional): Changes the encoding for the supported formats.
1323
+ For more details see :py:func:`torchaudio.backend.sox_io_backend.save`.
1324
+ bits_per_sample (int or None, optional): Changes the bit depth for the supported formats.
1325
+ For more details see :py:func:`torchaudio.backend.sox_io_backend.save`.
1326
+
1327
+ Returns:
1328
+ Tensor: Resulting Tensor.
1329
+ If ``channels_first=True``, it has `(channel, time)` else `(time, channel)`.
1330
+ """
1331
+ from torchaudio.backend import _sox_io_backend
1332
+
1333
+ with tempfile.NamedTemporaryFile() as f:
1334
+ torchaudio.backend._sox_io_backend.save(
1335
+ f.name, waveform, sample_rate, channels_first, compression, format, encoding, bits_per_sample
1336
+ )
1337
+ augmented, sr = _sox_io_backend.load(f.name, channels_first=channels_first, format=format)
1338
+ if sr != sample_rate:
1339
+ augmented = resample(augmented, sr, sample_rate)
1340
+ return augmented
1341
+
1342
+
1343
+ _CPU = torch.device("cpu")
1344
+
1345
+
1346
+ def _get_sinc_resample_kernel(
1347
+ orig_freq: int,
1348
+ new_freq: int,
1349
+ gcd: int,
1350
+ lowpass_filter_width: int = 6,
1351
+ rolloff: float = 0.99,
1352
+ resampling_method: str = "sinc_interp_hann",
1353
+ beta: Optional[float] = None,
1354
+ device: torch.device = _CPU,
1355
+ dtype: Optional[torch.dtype] = None,
1356
+ ):
1357
+ if not (int(orig_freq) == orig_freq and int(new_freq) == new_freq):
1358
+ raise Exception(
1359
+ "Frequencies must be of integer type to ensure quality resampling computation. "
1360
+ "To work around this, manually convert both frequencies to integer values "
1361
+ "that maintain their resampling rate ratio before passing them into the function. "
1362
+ "Example: To downsample a 44100 hz waveform by a factor of 8, use "
1363
+ "`orig_freq=8` and `new_freq=1` instead of `orig_freq=44100` and `new_freq=5512.5`. "
1364
+ "For more information, please refer to https://github.com/pytorch/audio/issues/1487."
1365
+ )
1366
+
1367
+ if resampling_method in ["sinc_interpolation", "kaiser_window"]:
1368
+ method_map = {
1369
+ "sinc_interpolation": "sinc_interp_hann",
1370
+ "kaiser_window": "sinc_interp_kaiser",
1371
+ }
1372
+ warnings.warn(
1373
+ f'"{resampling_method}" resampling method name is being deprecated and replaced by '
1374
+ f'"{method_map[resampling_method]}" in the next release. '
1375
+ "The default behavior remains unchanged.",
1376
+ stacklevel=3,
1377
+ )
1378
+ elif resampling_method not in ["sinc_interp_hann", "sinc_interp_kaiser"]:
1379
+ raise ValueError("Invalid resampling method: {}".format(resampling_method))
1380
+
1381
+ orig_freq = int(orig_freq) // gcd
1382
+ new_freq = int(new_freq) // gcd
1383
+
1384
+ if lowpass_filter_width <= 0:
1385
+ raise ValueError("Low pass filter width should be positive.")
1386
+ base_freq = min(orig_freq, new_freq)
1387
+ # This will perform antialiasing filtering by removing the highest frequencies.
1388
+ # At first I thought I only needed this when downsampling, but when upsampling
1389
+ # you will get edge artifacts without this, as the edge is equivalent to zero padding,
1390
+ # which will add high freq artifacts.
1391
+ base_freq *= rolloff
1392
+
1393
+ # The key idea of the algorithm is that x(t) can be exactly reconstructed from x[i] (tensor)
1394
+ # using the sinc interpolation formula:
1395
+ # x(t) = sum_i x[i] sinc(pi * orig_freq * (i / orig_freq - t))
1396
+ # We can then sample the function x(t) with a different sample rate:
1397
+ # y[j] = x(j / new_freq)
1398
+ # or,
1399
+ # y[j] = sum_i x[i] sinc(pi * orig_freq * (i / orig_freq - j / new_freq))
1400
+
1401
+ # We see here that y[j] is the convolution of x[i] with a specific filter, for which
1402
+ # we take an FIR approximation, stopping when we see at least `lowpass_filter_width` zeros crossing.
1403
+ # But y[j+1] is going to have a different set of weights and so on, until y[j + new_freq].
1404
+ # Indeed:
1405
+ # y[j + new_freq] = sum_i x[i] sinc(pi * orig_freq * ((i / orig_freq - (j + new_freq) / new_freq))
1406
+ # = sum_i x[i] sinc(pi * orig_freq * ((i - orig_freq) / orig_freq - j / new_freq))
1407
+ # = sum_i x[i + orig_freq] sinc(pi * orig_freq * (i / orig_freq - j / new_freq))
1408
+ # so y[j+new_freq] uses the same filter as y[j], but on a shifted version of x by `orig_freq`.
1409
+ # This will explain the F.conv1d after, with a stride of orig_freq.
1410
+ width = math.ceil(lowpass_filter_width * orig_freq / base_freq)
1411
+ # If orig_freq is still big after GCD reduction, most filters will be very unbalanced, i.e.,
1412
+ # they will have a lot of almost zero values to the left or to the right...
1413
+ # There is probably a way to evaluate those filters more efficiently, but this is kept for
1414
+ # future work.
1415
+ idx_dtype = dtype if dtype is not None else torch.float64
1416
+
1417
+ idx = torch.arange(-width, width + orig_freq, dtype=idx_dtype, device=device)[None, None] / orig_freq
1418
+
1419
+ t = torch.arange(0, -new_freq, -1, dtype=dtype, device=device)[:, None, None] / new_freq + idx
1420
+ t *= base_freq
1421
+ t = t.clamp_(-lowpass_filter_width, lowpass_filter_width)
1422
+
1423
+ # we do not use built in torch windows here as we need to evaluate the window
1424
+ # at specific positions, not over a regular grid.
1425
+ if resampling_method == "sinc_interp_hann":
1426
+ window = torch.cos(t * math.pi / lowpass_filter_width / 2) ** 2
1427
+ else:
1428
+ # sinc_interp_kaiser
1429
+ if beta is None:
1430
+ beta = 14.769656459379492
1431
+ beta_tensor = torch.tensor(float(beta))
1432
+ window = torch.i0(beta_tensor * torch.sqrt(1 - (t / lowpass_filter_width) ** 2)) / torch.i0(beta_tensor)
1433
+
1434
+ t *= math.pi
1435
+
1436
+ scale = base_freq / orig_freq
1437
+ kernels = torch.where(t == 0, torch.tensor(1.0).to(t), t.sin() / t)
1438
+ kernels *= window * scale
1439
+
1440
+ if dtype is None:
1441
+ kernels = kernels.to(dtype=torch.float32)
1442
+
1443
+ return kernels, width
1444
+
1445
+
1446
+ def _apply_sinc_resample_kernel(
1447
+ waveform: Tensor,
1448
+ orig_freq: int,
1449
+ new_freq: int,
1450
+ gcd: int,
1451
+ kernel: Tensor,
1452
+ width: int,
1453
+ ):
1454
+ if not waveform.is_floating_point():
1455
+ raise TypeError(f"Expected floating point type for waveform tensor, but received {waveform.dtype}.")
1456
+
1457
+ orig_freq = int(orig_freq) // gcd
1458
+ new_freq = int(new_freq) // gcd
1459
+
1460
+ # pack batch
1461
+ shape = waveform.size()
1462
+ waveform = waveform.view(-1, shape[-1])
1463
+
1464
+ num_wavs, length = waveform.shape
1465
+ waveform = torch.nn.functional.pad(waveform, (width, width + orig_freq))
1466
+ resampled = torch.nn.functional.conv1d(waveform[:, None], kernel, stride=orig_freq)
1467
+ resampled = resampled.transpose(1, 2).reshape(num_wavs, -1)
1468
+ target_length = torch.ceil(torch.as_tensor(new_freq * length / orig_freq)).long()
1469
+ resampled = resampled[..., :target_length]
1470
+
1471
+ # unpack batch
1472
+ resampled = resampled.view(shape[:-1] + resampled.shape[-1:])
1473
+ return resampled
1474
+
1475
+
1476
+ def resample(
1477
+ waveform: Tensor,
1478
+ orig_freq: int,
1479
+ new_freq: int,
1480
+ lowpass_filter_width: int = 6,
1481
+ rolloff: float = 0.99,
1482
+ resampling_method: str = "sinc_interp_hann",
1483
+ beta: Optional[float] = None,
1484
+ ) -> Tensor:
1485
+ r"""Resamples the waveform at the new frequency using bandlimited interpolation. :cite:`RESAMPLE`.
1486
+
1487
+ .. devices:: CPU CUDA
1488
+
1489
+ .. properties:: Autograd TorchScript
1490
+
1491
+ Note:
1492
+ ``transforms.Resample`` precomputes and reuses the resampling kernel, so using it will result in
1493
+ more efficient computation if resampling multiple waveforms with the same resampling parameters.
1494
+
1495
+ Args:
1496
+ waveform (Tensor): The input signal of dimension `(..., time)`
1497
+ orig_freq (int): The original frequency of the signal
1498
+ new_freq (int): The desired frequency
1499
+ lowpass_filter_width (int, optional): Controls the sharpness of the filter, more == sharper
1500
+ but less efficient. (Default: ``6``)
1501
+ rolloff (float, optional): The roll-off frequency of the filter, as a fraction of the Nyquist.
1502
+ Lower values reduce anti-aliasing, but also reduce some of the highest frequencies. (Default: ``0.99``)
1503
+ resampling_method (str, optional): The resampling method to use.
1504
+ Options: [``"sinc_interp_hann"``, ``"sinc_interp_kaiser"``] (Default: ``"sinc_interp_hann"``)
1505
+ beta (float or None, optional): The shape parameter used for kaiser window.
1506
+
1507
+ Returns:
1508
+ Tensor: The waveform at the new frequency of dimension `(..., time).`
1509
+ """
1510
+
1511
+ if orig_freq <= 0.0 or new_freq <= 0.0:
1512
+ raise ValueError("Original frequency and desired frequecy should be positive")
1513
+
1514
+ if orig_freq == new_freq:
1515
+ return waveform
1516
+
1517
+ gcd = math.gcd(int(orig_freq), int(new_freq))
1518
+
1519
+ kernel, width = _get_sinc_resample_kernel(
1520
+ orig_freq,
1521
+ new_freq,
1522
+ gcd,
1523
+ lowpass_filter_width,
1524
+ rolloff,
1525
+ resampling_method,
1526
+ beta,
1527
+ waveform.device,
1528
+ waveform.dtype,
1529
+ )
1530
+ resampled = _apply_sinc_resample_kernel(waveform, orig_freq, new_freq, gcd, kernel, width)
1531
+ return resampled
1532
+
1533
+
1534
+ @torch.jit.unused
1535
+ def edit_distance(seq1: Sequence, seq2: Sequence) -> int:
1536
+ """
1537
+ Calculate the word level edit (Levenshtein) distance between two sequences.
1538
+
1539
+ .. devices:: CPU
1540
+
1541
+ The function computes an edit distance allowing deletion, insertion and
1542
+ substitution. The result is an integer.
1543
+
1544
+ For most applications, the two input sequences should be the same type. If
1545
+ two strings are given, the output is the edit distance between the two
1546
+ strings (character edit distance). If two lists of strings are given, the
1547
+ output is the edit distance between sentences (word edit distance). Users
1548
+ may want to normalize the output by the length of the reference sequence.
1549
+
1550
+ Args:
1551
+ seq1 (Sequence): the first sequence to compare.
1552
+ seq2 (Sequence): the second sequence to compare.
1553
+ Returns:
1554
+ int: The distance between the first and second sequences.
1555
+ """
1556
+ len_sent2 = len(seq2)
1557
+ dold = list(range(len_sent2 + 1))
1558
+ dnew = [0 for _ in range(len_sent2 + 1)]
1559
+
1560
+ for i in range(1, len(seq1) + 1):
1561
+ dnew[0] = i
1562
+ for j in range(1, len_sent2 + 1):
1563
+ if seq1[i - 1] == seq2[j - 1]:
1564
+ dnew[j] = dold[j - 1]
1565
+ else:
1566
+ substitution = dold[j - 1] + 1
1567
+ insertion = dnew[j - 1] + 1
1568
+ deletion = dold[j] + 1
1569
+ dnew[j] = min(substitution, insertion, deletion)
1570
+
1571
+ dnew, dold = dold, dnew
1572
+
1573
+ return int(dold[-1])
1574
+
1575
+
1576
+ def loudness(waveform: Tensor, sample_rate: int):
1577
+ r"""Measure audio loudness according to the ITU-R BS.1770-4 recommendation.
1578
+
1579
+ .. devices:: CPU CUDA
1580
+
1581
+ .. properties:: TorchScript
1582
+
1583
+ Args:
1584
+ waveform(torch.Tensor): audio waveform of dimension `(..., channels, time)`
1585
+ sample_rate (int): sampling rate of the waveform
1586
+
1587
+ Returns:
1588
+ Tensor: loudness estimates (LKFS)
1589
+
1590
+ Reference:
1591
+ - https://www.itu.int/rec/R-REC-BS.1770-4-201510-I/en
1592
+ """
1593
+
1594
+ if waveform.size(-2) > 5:
1595
+ raise ValueError("Only up to 5 channels are supported.")
1596
+
1597
+ gate_duration = 0.4
1598
+ overlap = 0.75
1599
+ gamma_abs = -70.0
1600
+ kweight_bias = -0.691
1601
+ gate_samples = int(round(gate_duration * sample_rate))
1602
+ step = int(round(gate_samples * (1 - overlap)))
1603
+
1604
+ # Apply K-weighting
1605
+ waveform = treble_biquad(waveform, sample_rate, 4.0, 1500.0, 1 / math.sqrt(2))
1606
+ waveform = highpass_biquad(waveform, sample_rate, 38.0, 0.5)
1607
+
1608
+ # Compute the energy for each block
1609
+ energy = torch.square(waveform).unfold(-1, gate_samples, step)
1610
+ energy = torch.mean(energy, dim=-1)
1611
+
1612
+ # Compute channel-weighted summation
1613
+ g = torch.tensor([1.0, 1.0, 1.0, 1.41, 1.41], dtype=waveform.dtype, device=waveform.device)
1614
+ g = g[: energy.size(-2)]
1615
+
1616
+ energy_weighted = torch.sum(g.unsqueeze(-1) * energy, dim=-2)
1617
+ loudness = -0.691 + 10 * torch.log10(energy_weighted)
1618
+
1619
+ # Apply absolute gating of the blocks
1620
+ gated_blocks = loudness > gamma_abs
1621
+ gated_blocks = gated_blocks.unsqueeze(-2)
1622
+
1623
+ energy_filtered = torch.sum(gated_blocks * energy, dim=-1) / torch.count_nonzero(gated_blocks, dim=-1)
1624
+ energy_weighted = torch.sum(g * energy_filtered, dim=-1)
1625
+ gamma_rel = kweight_bias + 10 * torch.log10(energy_weighted) - 10
1626
+
1627
+ # Apply relative gating of the blocks
1628
+ gated_blocks = torch.logical_and(gated_blocks.squeeze(-2), loudness > gamma_rel.unsqueeze(-1))
1629
+ gated_blocks = gated_blocks.unsqueeze(-2)
1630
+
1631
+ energy_filtered = torch.sum(gated_blocks * energy, dim=-1) / torch.count_nonzero(gated_blocks, dim=-1)
1632
+ energy_weighted = torch.sum(g * energy_filtered, dim=-1)
1633
+ LKFS = kweight_bias + 10 * torch.log10(energy_weighted)
1634
+ return LKFS
1635
+
1636
+
1637
+ def pitch_shift(
1638
+ waveform: Tensor,
1639
+ sample_rate: int,
1640
+ n_steps: int,
1641
+ bins_per_octave: int = 12,
1642
+ n_fft: int = 512,
1643
+ win_length: Optional[int] = None,
1644
+ hop_length: Optional[int] = None,
1645
+ window: Optional[Tensor] = None,
1646
+ ) -> Tensor:
1647
+ """
1648
+ Shift the pitch of a waveform by ``n_steps`` steps.
1649
+
1650
+ .. devices:: CPU CUDA
1651
+
1652
+ .. properties:: TorchScript
1653
+
1654
+ Args:
1655
+ waveform (Tensor): The input waveform of shape `(..., time)`.
1656
+ sample_rate (int): Sample rate of `waveform`.
1657
+ n_steps (int): The (fractional) steps to shift `waveform`.
1658
+ bins_per_octave (int, optional): The number of steps per octave (Default: ``12``).
1659
+ n_fft (int, optional): Size of FFT, creates ``n_fft // 2 + 1`` bins (Default: ``512``).
1660
+ win_length (int or None, optional): Window size. If None, then ``n_fft`` is used. (Default: ``None``).
1661
+ hop_length (int or None, optional): Length of hop between STFT windows. If None, then
1662
+ ``win_length // 4`` is used (Default: ``None``).
1663
+ window (Tensor or None, optional): Window tensor that is applied/multiplied to each frame/window.
1664
+ If None, then ``torch.hann_window(win_length)`` is used (Default: ``None``).
1665
+
1666
+
1667
+ Returns:
1668
+ Tensor: The pitch-shifted audio waveform of shape `(..., time)`.
1669
+ """
1670
+ waveform_stretch = _stretch_waveform(
1671
+ waveform,
1672
+ n_steps,
1673
+ bins_per_octave,
1674
+ n_fft,
1675
+ win_length,
1676
+ hop_length,
1677
+ window,
1678
+ )
1679
+ rate = 2.0 ** (-float(n_steps) / bins_per_octave)
1680
+ waveform_shift = resample(waveform_stretch, int(sample_rate / rate), sample_rate)
1681
+
1682
+ return _fix_waveform_shape(waveform_shift, waveform.size())
1683
+
1684
+
1685
+ def _stretch_waveform(
1686
+ waveform: Tensor,
1687
+ n_steps: int,
1688
+ bins_per_octave: int = 12,
1689
+ n_fft: int = 512,
1690
+ win_length: Optional[int] = None,
1691
+ hop_length: Optional[int] = None,
1692
+ window: Optional[Tensor] = None,
1693
+ ) -> Tensor:
1694
+ """
1695
+ Pitch shift helper function to preprocess and stretch waveform before resampling step.
1696
+
1697
+ Args:
1698
+ See pitch_shift arg descriptions.
1699
+
1700
+ Returns:
1701
+ Tensor: The preprocessed waveform stretched prior to resampling.
1702
+ """
1703
+ if hop_length is None:
1704
+ hop_length = n_fft // 4
1705
+ if win_length is None:
1706
+ win_length = n_fft
1707
+ if window is None:
1708
+ window = torch.hann_window(window_length=win_length, device=waveform.device)
1709
+
1710
+ # pack batch
1711
+ shape = waveform.size()
1712
+ waveform = waveform.reshape(-1, shape[-1])
1713
+
1714
+ ori_len = shape[-1]
1715
+ rate = 2.0 ** (-float(n_steps) / bins_per_octave)
1716
+ spec_f = torch.stft(
1717
+ input=waveform,
1718
+ n_fft=n_fft,
1719
+ hop_length=hop_length,
1720
+ win_length=win_length,
1721
+ window=window,
1722
+ center=True,
1723
+ pad_mode="reflect",
1724
+ normalized=False,
1725
+ onesided=True,
1726
+ return_complex=True,
1727
+ )
1728
+ phase_advance = torch.linspace(0, math.pi * hop_length, spec_f.shape[-2], device=spec_f.device)[..., None]
1729
+ spec_stretch = phase_vocoder(spec_f, rate, phase_advance)
1730
+ len_stretch = int(round(ori_len / rate))
1731
+ waveform_stretch = torch.istft(
1732
+ spec_stretch, n_fft=n_fft, hop_length=hop_length, win_length=win_length, window=window, length=len_stretch
1733
+ )
1734
+ return waveform_stretch
1735
+
1736
+
1737
+ def _fix_waveform_shape(
1738
+ waveform_shift: Tensor,
1739
+ shape: List[int],
1740
+ ) -> Tensor:
1741
+ """
1742
+ PitchShift helper function to process after resampling step to fix the shape back.
1743
+
1744
+ Args:
1745
+ waveform_shift(Tensor): The waveform after stretch and resample
1746
+ shape (List[int]): The shape of initial waveform
1747
+
1748
+ Returns:
1749
+ Tensor: The pitch-shifted audio waveform of shape `(..., time)`.
1750
+ """
1751
+ ori_len = shape[-1]
1752
+ shift_len = waveform_shift.size()[-1]
1753
+ if shift_len > ori_len:
1754
+ waveform_shift = waveform_shift[..., :ori_len]
1755
+ else:
1756
+ waveform_shift = torch.nn.functional.pad(waveform_shift, [0, ori_len - shift_len])
1757
+
1758
+ # unpack batch
1759
+ waveform_shift = waveform_shift.view(shape[:-1] + waveform_shift.shape[-1:])
1760
+ return waveform_shift
1761
+
1762
+
1763
+ def rnnt_loss(
1764
+ logits: Tensor,
1765
+ targets: Tensor,
1766
+ logit_lengths: Tensor,
1767
+ target_lengths: Tensor,
1768
+ blank: int = -1,
1769
+ clamp: float = -1,
1770
+ reduction: str = "mean",
1771
+ fused_log_softmax: bool = True,
1772
+ ):
1773
+ """Compute the RNN Transducer loss from *Sequence Transduction with Recurrent Neural Networks*
1774
+ :cite:`graves2012sequence`.
1775
+
1776
+ .. devices:: CPU CUDA
1777
+
1778
+ .. properties:: Autograd TorchScript
1779
+
1780
+ The RNN Transducer loss extends the CTC loss by defining a distribution over output
1781
+ sequences of all lengths, and by jointly modelling both input-output and output-output
1782
+ dependencies.
1783
+
1784
+ Args:
1785
+ logits (Tensor): Tensor of dimension `(batch, max seq length, max target length + 1, class)`
1786
+ containing output from joiner
1787
+ targets (Tensor): Tensor of dimension `(batch, max target length)` containing targets with zero padded
1788
+ logit_lengths (Tensor): Tensor of dimension `(batch)` containing lengths of each sequence from encoder
1789
+ target_lengths (Tensor): Tensor of dimension `(batch)` containing lengths of targets for each sequence
1790
+ blank (int, optional): blank label (Default: ``-1``)
1791
+ clamp (float, optional): clamp for gradients (Default: ``-1``)
1792
+ reduction (string, optional): Specifies the reduction to apply to the output:
1793
+ ``"none"`` | ``"mean"`` | ``"sum"``. (Default: ``"mean"``)
1794
+ fused_log_softmax (bool): set to False if calling log_softmax outside of loss (Default: ``True``)
1795
+ Returns:
1796
+ Tensor: Loss with the reduction option applied. If ``reduction`` is ``"none"``, then size `(batch)`,
1797
+ otherwise scalar.
1798
+ """
1799
+ if reduction not in ["none", "mean", "sum"]:
1800
+ raise ValueError('reduction should be one of "none", "mean", or "sum"')
1801
+
1802
+ if blank < 0: # reinterpret blank index if blank < 0.
1803
+ blank = logits.shape[-1] + blank
1804
+
1805
+ costs, _ = torch.ops.torchaudio.rnnt_loss(
1806
+ logits=logits,
1807
+ targets=targets,
1808
+ logit_lengths=logit_lengths,
1809
+ target_lengths=target_lengths,
1810
+ blank=blank,
1811
+ clamp=clamp,
1812
+ fused_log_softmax=fused_log_softmax,
1813
+ )
1814
+
1815
+ if reduction == "mean":
1816
+ return costs.mean()
1817
+ elif reduction == "sum":
1818
+ return costs.sum()
1819
+
1820
+ return costs
1821
+
1822
+
1823
+ def psd(
1824
+ specgram: Tensor,
1825
+ mask: Optional[Tensor] = None,
1826
+ normalize: bool = True,
1827
+ eps: float = 1e-10,
1828
+ ) -> Tensor:
1829
+ """Compute cross-channel power spectral density (PSD) matrix.
1830
+
1831
+ .. devices:: CPU CUDA
1832
+
1833
+ .. properties:: Autograd TorchScript
1834
+
1835
+ Args:
1836
+ specgram (torch.Tensor): Multi-channel complex-valued spectrum.
1837
+ Tensor with dimensions `(..., channel, freq, time)`.
1838
+ mask (torch.Tensor or None, optional): Time-Frequency mask for normalization.
1839
+ Tensor with dimensions `(..., freq, time)`. (Default: ``None``)
1840
+ normalize (bool, optional): If ``True``, normalize the mask along the time dimension. (Default: ``True``)
1841
+ eps (float, optional): Value to add to the denominator in mask normalization. (Default: ``1e-15``)
1842
+
1843
+ Returns:
1844
+ torch.Tensor: The complex-valued PSD matrix of the input spectrum.
1845
+ Tensor with dimensions `(..., freq, channel, channel)`
1846
+ """
1847
+ specgram = specgram.transpose(-3, -2) # shape (freq, channel, time)
1848
+ # outer product:
1849
+ # (..., ch_1, time) x (..., ch_2, time) -> (..., time, ch_1, ch_2)
1850
+ psd = torch.einsum("...ct,...et->...tce", [specgram, specgram.conj()])
1851
+
1852
+ if mask is not None:
1853
+ if mask.shape[:-1] != specgram.shape[:-2] or mask.shape[-1] != specgram.shape[-1]:
1854
+ raise ValueError(
1855
+ "The dimensions of mask except the channel dimension should be the same as specgram."
1856
+ f"Found {mask.shape} for mask and {specgram.shape} for specgram."
1857
+ )
1858
+ # Normalized mask along time dimension:
1859
+ if normalize:
1860
+ mask = mask / (mask.sum(dim=-1, keepdim=True) + eps)
1861
+
1862
+ psd = psd * mask[..., None, None]
1863
+
1864
+ psd = psd.sum(dim=-3)
1865
+ return psd
1866
+
1867
+
1868
+ def _compute_mat_trace(input: torch.Tensor, dim1: int = -1, dim2: int = -2) -> torch.Tensor:
1869
+ r"""Compute the trace of a Tensor along ``dim1`` and ``dim2`` dimensions.
1870
+
1871
+ Args:
1872
+ input (torch.Tensor): Tensor with dimensions `(..., channel, channel)`.
1873
+ dim1 (int, optional): The first dimension of the diagonal matrix.
1874
+ (Default: ``-1``)
1875
+ dim2 (int, optional): The second dimension of the diagonal matrix.
1876
+ (Default: ``-2``)
1877
+
1878
+ Returns:
1879
+ Tensor: The trace of the input Tensor.
1880
+ """
1881
+ if input.ndim < 2:
1882
+ raise ValueError("The dimension of the tensor must be at least 2.")
1883
+ if input.shape[dim1] != input.shape[dim2]:
1884
+ raise ValueError("The size of ``dim1`` and ``dim2`` must be the same.")
1885
+ input = torch.diagonal(input, 0, dim1=dim1, dim2=dim2)
1886
+ return input.sum(dim=-1)
1887
+
1888
+
1889
+ def _tik_reg(mat: torch.Tensor, reg: float = 1e-7, eps: float = 1e-8) -> torch.Tensor:
1890
+ """Perform Tikhonov regularization (only modifying real part).
1891
+
1892
+ Args:
1893
+ mat (torch.Tensor): Input matrix with dimensions `(..., channel, channel)`.
1894
+ reg (float, optional): Regularization factor. (Default: 1e-8)
1895
+ eps (float, optional): Value to avoid the correlation matrix is all-zero. (Default: ``1e-8``)
1896
+
1897
+ Returns:
1898
+ Tensor: Regularized matrix with dimensions `(..., channel, channel)`.
1899
+ """
1900
+ # Add eps
1901
+ C = mat.size(-1)
1902
+ eye = torch.eye(C, dtype=mat.dtype, device=mat.device)
1903
+ epsilon = _compute_mat_trace(mat).real[..., None, None] * reg
1904
+ # in case that correlation_matrix is all-zero
1905
+ epsilon = epsilon + eps
1906
+ mat = mat + epsilon * eye[..., :, :]
1907
+ return mat
1908
+
1909
+
1910
+ def _assert_psd_matrices(psd_s: torch.Tensor, psd_n: torch.Tensor) -> None:
1911
+ """Assertion checks of the PSD matrices of target speech and noise.
1912
+
1913
+ Args:
1914
+ psd_s (torch.Tensor): The complex-valued power spectral density (PSD) matrix of target speech.
1915
+ Tensor with dimensions `(..., freq, channel, channel)`.
1916
+ psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise.
1917
+ Tensor with dimensions `(..., freq, channel, channel)`.
1918
+ """
1919
+ if psd_s.ndim < 3 or psd_n.ndim < 3:
1920
+ raise ValueError(
1921
+ "Expected at least 3D Tensor (..., freq, channel, channel) for psd_s and psd_n. "
1922
+ f"Found {psd_s.shape} for psd_s and {psd_n.shape} for psd_n."
1923
+ )
1924
+ if not (psd_s.is_complex() and psd_n.is_complex()):
1925
+ raise TypeError(
1926
+ "The type of psd_s and psd_n must be ``torch.cfloat`` or ``torch.cdouble``. "
1927
+ f"Found {psd_s.dtype} for psd_s and {psd_n.dtype} for psd_n."
1928
+ )
1929
+ if psd_s.shape != psd_n.shape:
1930
+ raise ValueError(
1931
+ f"The dimensions of psd_s and psd_n should be the same. Found {psd_s.shape} and {psd_n.shape}."
1932
+ )
1933
+ if psd_s.shape[-1] != psd_s.shape[-2]:
1934
+ raise ValueError(f"The last two dimensions of psd_s should be the same. Found {psd_s.shape}.")
1935
+
1936
+
1937
+ def mvdr_weights_souden(
1938
+ psd_s: Tensor,
1939
+ psd_n: Tensor,
1940
+ reference_channel: Union[int, Tensor],
1941
+ diagonal_loading: bool = True,
1942
+ diag_eps: float = 1e-7,
1943
+ eps: float = 1e-8,
1944
+ ) -> Tensor:
1945
+ r"""Compute the Minimum Variance Distortionless Response (*MVDR* :cite:`capon1969high`) beamforming weights
1946
+ by the method proposed by *Souden et, al.* :cite:`souden2009optimal`.
1947
+
1948
+ .. devices:: CPU CUDA
1949
+
1950
+ .. properties:: Autograd TorchScript
1951
+
1952
+ Given the power spectral density (PSD) matrix of target speech :math:`\bf{\Phi}_{\textbf{SS}}`,
1953
+ the PSD matrix of noise :math:`\bf{\Phi}_{\textbf{NN}}`, and a one-hot vector that represents the
1954
+ reference channel :math:`\bf{u}`, the method computes the MVDR beamforming weight martrix
1955
+ :math:`\textbf{w}_{\text{MVDR}}`. The formula is defined as:
1956
+
1957
+ .. math::
1958
+ \textbf{w}_{\text{MVDR}}(f) =
1959
+ \frac{{{\bf{\Phi}_{\textbf{NN}}^{-1}}(f){\bf{\Phi}_{\textbf{SS}}}}(f)}
1960
+ {\text{Trace}({{{\bf{\Phi}_{\textbf{NN}}^{-1}}(f) \bf{\Phi}_{\textbf{SS}}}(f))}}\bm{u}
1961
+
1962
+ Args:
1963
+ psd_s (torch.Tensor): The complex-valued power spectral density (PSD) matrix of target speech.
1964
+ Tensor with dimensions `(..., freq, channel, channel)`.
1965
+ psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise.
1966
+ Tensor with dimensions `(..., freq, channel, channel)`.
1967
+ reference_channel (int or torch.Tensor): Specifies the reference channel.
1968
+ If the dtype is ``int``, it represents the reference channel index.
1969
+ If the dtype is ``torch.Tensor``, its shape is `(..., channel)`, where the ``channel`` dimension
1970
+ is one-hot.
1971
+ diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to ``psd_n``.
1972
+ (Default: ``True``)
1973
+ diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading.
1974
+ It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``)
1975
+ eps (float, optional): Value to add to the denominator in the beamforming weight formula.
1976
+ (Default: ``1e-8``)
1977
+
1978
+ Returns:
1979
+ torch.Tensor: The complex-valued MVDR beamforming weight matrix with dimensions `(..., freq, channel)`.
1980
+ """
1981
+ _assert_psd_matrices(psd_s, psd_n)
1982
+
1983
+ if diagonal_loading:
1984
+ psd_n = _tik_reg(psd_n, reg=diag_eps)
1985
+ numerator = torch.linalg.solve(psd_n, psd_s) # psd_n.inv() @ psd_s
1986
+ # ws: (..., C, C) / (...,) -> (..., C, C)
1987
+ ws = numerator / (_compute_mat_trace(numerator)[..., None, None] + eps)
1988
+ if torch.jit.isinstance(reference_channel, int):
1989
+ beamform_weights = ws[..., :, reference_channel]
1990
+ elif torch.jit.isinstance(reference_channel, Tensor):
1991
+ reference_channel = reference_channel.to(psd_n.dtype)
1992
+ # h: (..., F, C_1, C_2) x (..., C_2) -> (..., F, C_1)
1993
+ beamform_weights = torch.einsum("...c,...c->...", [ws, reference_channel[..., None, None, :]])
1994
+ else:
1995
+ raise TypeError(f'Expected "int" or "Tensor" for reference_channel. Found: {type(reference_channel)}.')
1996
+
1997
+ return beamform_weights
1998
+
1999
+
2000
+ def mvdr_weights_rtf(
2001
+ rtf: Tensor,
2002
+ psd_n: Tensor,
2003
+ reference_channel: Optional[Union[int, Tensor]] = None,
2004
+ diagonal_loading: bool = True,
2005
+ diag_eps: float = 1e-7,
2006
+ eps: float = 1e-8,
2007
+ ) -> Tensor:
2008
+ r"""Compute the Minimum Variance Distortionless Response (*MVDR* :cite:`capon1969high`) beamforming weights
2009
+ based on the relative transfer function (RTF) and power spectral density (PSD) matrix of noise.
2010
+
2011
+ .. devices:: CPU CUDA
2012
+
2013
+ .. properties:: Autograd TorchScript
2014
+
2015
+ Given the relative transfer function (RTF) matrix or the steering vector of target speech :math:`\bm{v}`,
2016
+ the PSD matrix of noise :math:`\bf{\Phi}_{\textbf{NN}}`, and a one-hot vector that represents the
2017
+ reference channel :math:`\bf{u}`, the method computes the MVDR beamforming weight martrix
2018
+ :math:`\textbf{w}_{\text{MVDR}}`. The formula is defined as:
2019
+
2020
+ .. math::
2021
+ \textbf{w}_{\text{MVDR}}(f) =
2022
+ \frac{{{\bf{\Phi}_{\textbf{NN}}^{-1}}(f){\bm{v}}(f)}}
2023
+ {{\bm{v}^{\mathsf{H}}}(f){\bf{\Phi}_{\textbf{NN}}^{-1}}(f){\bm{v}}(f)}
2024
+
2025
+ where :math:`(.)^{\mathsf{H}}` denotes the Hermitian Conjugate operation.
2026
+
2027
+ Args:
2028
+ rtf (torch.Tensor): The complex-valued RTF vector of target speech.
2029
+ Tensor with dimensions `(..., freq, channel)`.
2030
+ psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise.
2031
+ Tensor with dimensions `(..., freq, channel, channel)`.
2032
+ reference_channel (int or torch.Tensor): Specifies the reference channel.
2033
+ If the dtype is ``int``, it represents the reference channel index.
2034
+ If the dtype is ``torch.Tensor``, its shape is `(..., channel)`, where the ``channel`` dimension
2035
+ is one-hot.
2036
+ diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to ``psd_n``.
2037
+ (Default: ``True``)
2038
+ diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading.
2039
+ It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``)
2040
+ eps (float, optional): Value to add to the denominator in the beamforming weight formula.
2041
+ (Default: ``1e-8``)
2042
+
2043
+ Returns:
2044
+ torch.Tensor: The complex-valued MVDR beamforming weight matrix with dimensions `(..., freq, channel)`.
2045
+ """
2046
+ if rtf.ndim < 2:
2047
+ raise ValueError(f"Expected at least 2D Tensor (..., freq, channel) for rtf. Found {rtf.shape}.")
2048
+ if psd_n.ndim < 3:
2049
+ raise ValueError(f"Expected at least 3D Tensor (..., freq, channel, channel) for psd_n. Found {psd_n.shape}.")
2050
+ if not (rtf.is_complex() and psd_n.is_complex()):
2051
+ raise TypeError(
2052
+ "The type of rtf and psd_n must be ``torch.cfloat`` or ``torch.cdouble``. "
2053
+ f"Found {rtf.dtype} for rtf and {psd_n.dtype} for psd_n."
2054
+ )
2055
+ if rtf.shape != psd_n.shape[:-1]:
2056
+ raise ValueError(
2057
+ "The dimensions of rtf and the dimensions withou the last dimension of psd_n should be the same. "
2058
+ f"Found {rtf.shape} for rtf and {psd_n.shape} for psd_n."
2059
+ )
2060
+ if psd_n.shape[-1] != psd_n.shape[-2]:
2061
+ raise ValueError(f"The last two dimensions of psd_n should be the same. Found {psd_n.shape}.")
2062
+
2063
+ if diagonal_loading:
2064
+ psd_n = _tik_reg(psd_n, reg=diag_eps)
2065
+ # numerator = psd_n.inv() @ stv
2066
+ numerator = torch.linalg.solve(psd_n, rtf.unsqueeze(-1)).squeeze(-1) # (..., freq, channel)
2067
+ # denominator = stv^H @ psd_n.inv() @ stv
2068
+ denominator = torch.einsum("...d,...d->...", [rtf.conj(), numerator])
2069
+ beamform_weights = numerator / (denominator.real.unsqueeze(-1) + eps)
2070
+ # normalize the numerator
2071
+ if reference_channel is not None:
2072
+ if torch.jit.isinstance(reference_channel, int):
2073
+ scale = rtf[..., reference_channel].conj()
2074
+ elif torch.jit.isinstance(reference_channel, Tensor):
2075
+ reference_channel = reference_channel.to(psd_n.dtype)
2076
+ scale = torch.einsum("...c,...c->...", [rtf.conj(), reference_channel[..., None, :]])
2077
+ else:
2078
+ raise TypeError(f'Expected "int" or "Tensor" for reference_channel. Found: {type(reference_channel)}.')
2079
+
2080
+ beamform_weights = beamform_weights * scale[..., None]
2081
+
2082
+ return beamform_weights
2083
+
2084
+
2085
+ def rtf_evd(psd_s: Tensor) -> Tensor:
2086
+ r"""Estimate the relative transfer function (RTF) or the steering vector by eigenvalue decomposition.
2087
+
2088
+ .. devices:: CPU CUDA
2089
+
2090
+ .. properties:: TorchScript
2091
+
2092
+ Args:
2093
+ psd_s (Tensor): The complex-valued power spectral density (PSD) matrix of target speech.
2094
+ Tensor of dimension `(..., freq, channel, channel)`
2095
+
2096
+ Returns:
2097
+ Tensor: The estimated complex-valued RTF of target speech.
2098
+ Tensor of dimension `(..., freq, channel)`
2099
+ """
2100
+ if not psd_s.is_complex():
2101
+ raise TypeError(f"The type of psd_s must be ``torch.cfloat`` or ``torch.cdouble``. Found {psd_s.dtype}.")
2102
+ if psd_s.shape[-1] != psd_s.shape[-2]:
2103
+ raise ValueError(f"The last two dimensions of psd_s should be the same. Found {psd_s.shape}.")
2104
+ _, v = torch.linalg.eigh(psd_s) # v is sorted along with eigenvalues in ascending order
2105
+ rtf = v[..., -1] # choose the eigenvector with max eigenvalue
2106
+ return rtf
2107
+
2108
+
2109
+ def rtf_power(
2110
+ psd_s: Tensor,
2111
+ psd_n: Tensor,
2112
+ reference_channel: Union[int, Tensor],
2113
+ n_iter: int = 3,
2114
+ diagonal_loading: bool = True,
2115
+ diag_eps: float = 1e-7,
2116
+ ) -> Tensor:
2117
+ r"""Estimate the relative transfer function (RTF) or the steering vector by the power method.
2118
+
2119
+ .. devices:: CPU CUDA
2120
+
2121
+ .. properties:: Autograd TorchScript
2122
+
2123
+ Args:
2124
+ psd_s (torch.Tensor): The complex-valued power spectral density (PSD) matrix of target speech.
2125
+ Tensor with dimensions `(..., freq, channel, channel)`.
2126
+ psd_n (torch.Tensor): The complex-valued power spectral density (PSD) matrix of noise.
2127
+ Tensor with dimensions `(..., freq, channel, channel)`.
2128
+ reference_channel (int or torch.Tensor): Specifies the reference channel.
2129
+ If the dtype is ``int``, it represents the reference channel index.
2130
+ If the dtype is ``torch.Tensor``, its shape is `(..., channel)`, where the ``channel`` dimension
2131
+ is one-hot.
2132
+ diagonal_loading (bool, optional): If ``True``, enables applying diagonal loading to ``psd_n``.
2133
+ (Default: ``True``)
2134
+ diag_eps (float, optional): The coefficient multiplied to the identity matrix for diagonal loading.
2135
+ It is only effective when ``diagonal_loading`` is set to ``True``. (Default: ``1e-7``)
2136
+
2137
+ Returns:
2138
+ torch.Tensor: The estimated complex-valued RTF of target speech.
2139
+ Tensor of dimension `(..., freq, channel)`.
2140
+ """
2141
+ _assert_psd_matrices(psd_s, psd_n)
2142
+ if n_iter <= 0:
2143
+ raise ValueError("The number of iteration must be greater than 0.")
2144
+
2145
+ # Apply diagonal loading to psd_n to improve robustness.
2146
+ if diagonal_loading:
2147
+ psd_n = _tik_reg(psd_n, reg=diag_eps)
2148
+ # phi is regarded as the first iteration
2149
+ phi = torch.linalg.solve(psd_n, psd_s) # psd_n.inv() @ psd_s
2150
+ if torch.jit.isinstance(reference_channel, int):
2151
+ rtf = phi[..., reference_channel]
2152
+ elif torch.jit.isinstance(reference_channel, Tensor):
2153
+ reference_channel = reference_channel.to(psd_n.dtype)
2154
+ rtf = torch.einsum("...c,...c->...", [phi, reference_channel[..., None, None, :]])
2155
+ else:
2156
+ raise TypeError(f'Expected "int" or "Tensor" for reference_channel. Found: {type(reference_channel)}.')
2157
+ rtf = rtf.unsqueeze(-1) # (..., freq, channel, 1)
2158
+ if n_iter >= 2:
2159
+ # The number of iterations in the for loop is `n_iter - 2`
2160
+ # because the `phi` above and `torch.matmul(psd_s, rtf)` are regarded as
2161
+ # two iterations.
2162
+ for _ in range(n_iter - 2):
2163
+ rtf = torch.matmul(phi, rtf)
2164
+ rtf = torch.matmul(psd_s, rtf)
2165
+ else:
2166
+ # if there is only one iteration, the rtf is the psd_s[..., referenc_channel]
2167
+ # which is psd_n @ phi @ ref_channel
2168
+ rtf = torch.matmul(psd_n, rtf)
2169
+ return rtf.squeeze(-1)
2170
+
2171
+
2172
+ def apply_beamforming(beamform_weights: Tensor, specgram: Tensor) -> Tensor:
2173
+ r"""Apply the beamforming weight to the multi-channel noisy spectrum to obtain the single-channel enhanced spectrum.
2174
+
2175
+ .. devices:: CPU CUDA
2176
+
2177
+ .. properties:: Autograd TorchScript
2178
+
2179
+ .. math::
2180
+ \hat{\textbf{S}}(f) = \textbf{w}_{\text{bf}}(f)^{\mathsf{H}} \textbf{Y}(f)
2181
+
2182
+ where :math:`\textbf{w}_{\text{bf}}(f)` is the beamforming weight for the :math:`f`-th frequency bin,
2183
+ :math:`\textbf{Y}` is the multi-channel spectrum for the :math:`f`-th frequency bin.
2184
+
2185
+ Args:
2186
+ beamform_weights (Tensor): The complex-valued beamforming weight matrix.
2187
+ Tensor of dimension `(..., freq, channel)`
2188
+ specgram (Tensor): The multi-channel complex-valued noisy spectrum.
2189
+ Tensor of dimension `(..., channel, freq, time)`
2190
+
2191
+ Returns:
2192
+ Tensor: The single-channel complex-valued enhanced spectrum.
2193
+ Tensor of dimension `(..., freq, time)`
2194
+ """
2195
+ if beamform_weights.shape[:-2] != specgram.shape[:-3]:
2196
+ raise ValueError(
2197
+ "The dimensions except the last two dimensions of beamform_weights should be the same "
2198
+ "as the dimensions except the last three dimensions of specgram. "
2199
+ f"Found {beamform_weights.shape} for beamform_weights and {specgram.shape} for specgram."
2200
+ )
2201
+
2202
+ if not (beamform_weights.is_complex() and specgram.is_complex()):
2203
+ raise TypeError(
2204
+ "The type of beamform_weights and specgram must be ``torch.cfloat`` or ``torch.cdouble``. "
2205
+ f"Found {beamform_weights.dtype} for beamform_weights and {specgram.dtype} for specgram."
2206
+ )
2207
+
2208
+ # (..., freq, channel) x (..., channel, freq, time) -> (..., freq, time)
2209
+ specgram_enhanced = torch.einsum("...fc,...cft->...ft", [beamform_weights.conj(), specgram])
2210
+ return specgram_enhanced
2211
+
2212
+
2213
+ def _check_shape_compatible(x: torch.Tensor, y: torch.Tensor) -> None:
2214
+ if x.ndim != y.ndim:
2215
+ raise ValueError(f"The operands must be the same dimension (got {x.ndim} and {y.ndim}).")
2216
+
2217
+ for i in range(x.ndim - 1):
2218
+ xi = x.size(i)
2219
+ yi = y.size(i)
2220
+ if xi == yi or xi == 1 or yi == 1:
2221
+ continue
2222
+ raise ValueError(f"Leading dimensions of x and y are not broadcastable (got {x.shape} and {y.shape}).")
2223
+
2224
+
2225
+ def _check_convolve_mode(mode: str) -> None:
2226
+ valid_convolve_modes = ["full", "valid", "same"]
2227
+ if mode not in valid_convolve_modes:
2228
+ raise ValueError(f"Unrecognized mode value '{mode}'. Please specify one of {valid_convolve_modes}.")
2229
+
2230
+
2231
+ def _apply_convolve_mode(conv_result: torch.Tensor, x_length: int, y_length: int, mode: str) -> torch.Tensor:
2232
+ valid_convolve_modes = ["full", "valid", "same"]
2233
+ if mode == "full":
2234
+ return conv_result
2235
+ elif mode == "valid":
2236
+ target_length = max(x_length, y_length) - min(x_length, y_length) + 1
2237
+ start_idx = (conv_result.size(-1) - target_length) // 2
2238
+ return conv_result[..., start_idx : start_idx + target_length]
2239
+ elif mode == "same":
2240
+ start_idx = (conv_result.size(-1) - x_length) // 2
2241
+ return conv_result[..., start_idx : start_idx + x_length]
2242
+ else:
2243
+ raise ValueError(f"Unrecognized mode value '{mode}'. Please specify one of {valid_convolve_modes}.")
2244
+
2245
+
2246
+ def fftconvolve(x: torch.Tensor, y: torch.Tensor, mode: str = "full") -> torch.Tensor:
2247
+ r"""
2248
+ Convolves inputs along their last dimension using FFT. For inputs with large last dimensions, this function
2249
+ is generally much faster than :meth:`convolve`.
2250
+ Note that, in contrast to :meth:`torch.nn.functional.conv1d`, which actually applies the valid cross-correlation
2251
+ operator, this function applies the true `convolution`_ operator.
2252
+ Also note that this function can only output float tensors (int tensor inputs will be cast to float).
2253
+
2254
+ .. devices:: CPU CUDA
2255
+
2256
+ .. properties:: Autograd TorchScript
2257
+
2258
+ Args:
2259
+ x (torch.Tensor): First convolution operand, with shape `(..., N)`.
2260
+ y (torch.Tensor): Second convolution operand, with shape `(..., M)`
2261
+ (leading dimensions must be broadcast-able with those of ``x``).
2262
+ mode (str, optional): Must be one of ("full", "valid", "same").
2263
+
2264
+ * "full": Returns the full convolution result, with shape `(..., N + M - 1)`. (Default)
2265
+ * "valid": Returns the segment of the full convolution result corresponding to where
2266
+ the two inputs overlap completely, with shape `(..., max(N, M) - min(N, M) + 1)`.
2267
+ * "same": Returns the center segment of the full convolution result, with shape `(..., N)`.
2268
+
2269
+ Returns:
2270
+ torch.Tensor: Result of convolving ``x`` and ``y``, with shape `(..., L)`, where
2271
+ the leading dimensions match those of ``x`` and `L` is dictated by ``mode``.
2272
+
2273
+ .. _convolution:
2274
+ https://en.wikipedia.org/wiki/Convolution
2275
+ """
2276
+ _check_shape_compatible(x, y)
2277
+ _check_convolve_mode(mode)
2278
+
2279
+ n = x.size(-1) + y.size(-1) - 1
2280
+ fresult = torch.fft.rfft(x, n=n) * torch.fft.rfft(y, n=n)
2281
+ result = torch.fft.irfft(fresult, n=n)
2282
+ return _apply_convolve_mode(result, x.size(-1), y.size(-1), mode)
2283
+
2284
+
2285
+ def convolve(x: torch.Tensor, y: torch.Tensor, mode: str = "full") -> torch.Tensor:
2286
+ r"""
2287
+ Convolves inputs along their last dimension using the direct method.
2288
+ Note that, in contrast to :meth:`torch.nn.functional.conv1d`, which actually applies the valid cross-correlation
2289
+ operator, this function applies the true `convolution`_ operator.
2290
+
2291
+ .. devices:: CPU CUDA
2292
+
2293
+ .. properties:: Autograd TorchScript
2294
+
2295
+ Args:
2296
+ x (torch.Tensor): First convolution operand, with shape `(..., N)`.
2297
+ y (torch.Tensor): Second convolution operand, with shape `(..., M)`
2298
+ (leading dimensions must be broadcast-able with those of ``x``).
2299
+ mode (str, optional): Must be one of ("full", "valid", "same").
2300
+
2301
+ * "full": Returns the full convolution result, with shape `(..., N + M - 1)`. (Default)
2302
+ * "valid": Returns the segment of the full convolution result corresponding to where
2303
+ the two inputs overlap completely, with shape `(..., max(N, M) - min(N, M) + 1)`.
2304
+ * "same": Returns the center segment of the full convolution result, with shape `(..., N)`.
2305
+
2306
+ Returns:
2307
+ torch.Tensor: Result of convolving ``x`` and ``y``, with shape `(..., L)`, where
2308
+ the leading dimensions match those of ``x`` and `L` is dictated by ``mode``.
2309
+
2310
+ .. _convolution:
2311
+ https://en.wikipedia.org/wiki/Convolution
2312
+ """
2313
+ _check_shape_compatible(x, y)
2314
+ _check_convolve_mode(mode)
2315
+
2316
+ x_size, y_size = x.size(-1), y.size(-1)
2317
+
2318
+ if x.size(-1) < y.size(-1):
2319
+ x, y = y, x
2320
+
2321
+ if x.shape[:-1] != y.shape[:-1]:
2322
+ new_shape = [max(i, j) for i, j in zip(x.shape[:-1], y.shape[:-1])]
2323
+ x = x.broadcast_to(new_shape + [x.shape[-1]])
2324
+ y = y.broadcast_to(new_shape + [y.shape[-1]])
2325
+
2326
+ num_signals = torch.tensor(x.shape[:-1]).prod()
2327
+ reshaped_x = x.reshape((int(num_signals), x.size(-1)))
2328
+ reshaped_y = y.reshape((int(num_signals), y.size(-1)))
2329
+ output = torch.nn.functional.conv1d(
2330
+ input=reshaped_x,
2331
+ weight=reshaped_y.flip(-1).unsqueeze(1),
2332
+ stride=1,
2333
+ groups=reshaped_x.size(0),
2334
+ padding=reshaped_y.size(-1) - 1,
2335
+ )
2336
+ output_shape = x.shape[:-1] + (-1,)
2337
+ result = output.reshape(output_shape)
2338
+ return _apply_convolve_mode(result, x_size, y_size, mode)
2339
+
2340
+
2341
+ def add_noise(
2342
+ waveform: torch.Tensor, noise: torch.Tensor, snr: torch.Tensor, lengths: Optional[torch.Tensor] = None
2343
+ ) -> torch.Tensor:
2344
+ r"""Scales and adds noise to waveform per signal-to-noise ratio.
2345
+
2346
+ Specifically, for each pair of waveform vector :math:`x \in \mathbb{R}^L` and noise vector
2347
+ :math:`n \in \mathbb{R}^L`, the function computes output :math:`y` as
2348
+
2349
+ .. math::
2350
+ y = x + a n \, \text{,}
2351
+
2352
+ where
2353
+
2354
+ .. math::
2355
+ a = \sqrt{ \frac{ ||x||_{2}^{2} }{ ||n||_{2}^{2} } \cdot 10^{-\frac{\text{SNR}}{10}} } \, \text{,}
2356
+
2357
+ with :math:`\text{SNR}` being the desired signal-to-noise ratio between :math:`x` and :math:`n`, in dB.
2358
+
2359
+ Note that this function broadcasts singleton leading dimensions in its inputs in a manner that is
2360
+ consistent with the above formulae and PyTorch's broadcasting semantics.
2361
+
2362
+ .. devices:: CPU CUDA
2363
+
2364
+ .. properties:: Autograd TorchScript
2365
+
2366
+ Args:
2367
+ waveform (torch.Tensor): Input waveform, with shape `(..., L)`.
2368
+ noise (torch.Tensor): Noise, with shape `(..., L)` (same shape as ``waveform``).
2369
+ snr (torch.Tensor): Signal-to-noise ratios in dB, with shape `(...,)`.
2370
+ lengths (torch.Tensor or None, optional): Valid lengths of signals in ``waveform`` and ``noise``, with shape
2371
+ `(...,)` (leading dimensions must match those of ``waveform``). If ``None``, all elements in ``waveform``
2372
+ and ``noise`` are treated as valid. (Default: ``None``)
2373
+
2374
+ Returns:
2375
+ torch.Tensor: Result of scaling and adding ``noise`` to ``waveform``, with shape `(..., L)`
2376
+ (same shape as ``waveform``).
2377
+ """
2378
+
2379
+ if not (waveform.ndim - 1 == noise.ndim - 1 == snr.ndim and (lengths is None or lengths.ndim == snr.ndim)):
2380
+ raise ValueError("Input leading dimensions don't match.")
2381
+
2382
+ L = waveform.size(-1)
2383
+
2384
+ if L != noise.size(-1):
2385
+ raise ValueError(f"Length dimensions of waveform and noise don't match (got {L} and {noise.size(-1)}).")
2386
+
2387
+ # compute scale
2388
+ if lengths is not None:
2389
+ mask = torch.arange(0, L, device=lengths.device).expand(waveform.shape) < lengths.unsqueeze(
2390
+ -1
2391
+ ) # (*, L) < (*, 1) = (*, L)
2392
+ masked_waveform = waveform * mask
2393
+ masked_noise = noise * mask
2394
+ else:
2395
+ masked_waveform = waveform
2396
+ masked_noise = noise
2397
+
2398
+ energy_signal = torch.linalg.vector_norm(masked_waveform, ord=2, dim=-1) ** 2 # (*,)
2399
+ energy_noise = torch.linalg.vector_norm(masked_noise, ord=2, dim=-1) ** 2 # (*,)
2400
+ original_snr_db = 10 * (torch.log10(energy_signal) - torch.log10(energy_noise))
2401
+ scale = 10 ** ((original_snr_db - snr) / 20.0) # (*,)
2402
+
2403
+ # scale noise
2404
+ scaled_noise = scale.unsqueeze(-1) * noise # (*, 1) * (*, L) = (*, L)
2405
+
2406
+ return waveform + scaled_noise # (*, L)
2407
+
2408
+
2409
+ def speed(
2410
+ waveform: torch.Tensor, orig_freq: int, factor: float, lengths: Optional[torch.Tensor] = None
2411
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
2412
+ r"""Adjusts waveform speed.
2413
+
2414
+ .. devices:: CPU CUDA
2415
+
2416
+ .. properties:: Autograd TorchScript
2417
+
2418
+ Args:
2419
+ waveform (torch.Tensor): Input signals, with shape `(..., time)`.
2420
+ orig_freq (int): Original frequency of the signals in ``waveform``.
2421
+ factor (float): Factor by which to adjust speed of input. Values greater than 1.0
2422
+ compress ``waveform`` in time, whereas values less than 1.0 stretch ``waveform`` in time.
2423
+ lengths (torch.Tensor or None, optional): Valid lengths of signals in ``waveform``, with shape `(...)`.
2424
+ If ``None``, all elements in ``waveform`` are treated as valid. (Default: ``None``)
2425
+
2426
+ Returns:
2427
+ (torch.Tensor, torch.Tensor or None):
2428
+ torch.Tensor
2429
+ Speed-adjusted waveform, with shape `(..., new_time).`
2430
+ torch.Tensor or None
2431
+ If ``lengths`` is not ``None``, valid lengths of signals in speed-adjusted waveform,
2432
+ with shape `(...)`; otherwise, ``None``.
2433
+ """
2434
+
2435
+ source_sample_rate = int(factor * orig_freq)
2436
+ target_sample_rate = int(orig_freq)
2437
+
2438
+ gcd = math.gcd(source_sample_rate, target_sample_rate)
2439
+ source_sample_rate = source_sample_rate // gcd
2440
+ target_sample_rate = target_sample_rate // gcd
2441
+
2442
+ if lengths is None:
2443
+ out_lengths = None
2444
+ else:
2445
+ out_lengths = torch.ceil(lengths * target_sample_rate / source_sample_rate).to(lengths.dtype)
2446
+
2447
+ return resample(waveform, source_sample_rate, target_sample_rate), out_lengths
2448
+
2449
+
2450
+ def preemphasis(waveform, coeff: float = 0.97) -> torch.Tensor:
2451
+ r"""Pre-emphasizes a waveform along its last dimension, i.e.
2452
+ for each signal :math:`x` in ``waveform``, computes
2453
+ output :math:`y` as
2454
+
2455
+ .. math::
2456
+ y[i] = x[i] - \text{coeff} \cdot x[i - 1]
2457
+
2458
+ .. devices:: CPU CUDA
2459
+
2460
+ .. properties:: Autograd TorchScript
2461
+
2462
+ Args:
2463
+ waveform (torch.Tensor): Waveform, with shape `(..., N)`.
2464
+ coeff (float, optional): Pre-emphasis coefficient. Typically between 0.0 and 1.0.
2465
+ (Default: 0.97)
2466
+
2467
+ Returns:
2468
+ torch.Tensor: Pre-emphasized waveform, with shape `(..., N)`.
2469
+ """
2470
+ waveform = waveform.clone()
2471
+ waveform[..., 1:] -= coeff * waveform[..., :-1]
2472
+ return waveform
2473
+
2474
+
2475
+ def deemphasis(waveform, coeff: float = 0.97) -> torch.Tensor:
2476
+ r"""De-emphasizes a waveform along its last dimension.
2477
+ Inverse of :meth:`preemphasis`. Concretely, for each signal
2478
+ :math:`x` in ``waveform``, computes output :math:`y` as
2479
+
2480
+ .. math::
2481
+ y[i] = x[i] + \text{coeff} \cdot y[i - 1]
2482
+
2483
+ .. devices:: CPU CUDA
2484
+
2485
+ .. properties:: Autograd TorchScript
2486
+
2487
+ Args:
2488
+ waveform (torch.Tensor): Waveform, with shape `(..., N)`.
2489
+ coeff (float, optional): De-emphasis coefficient. Typically between 0.0 and 1.0.
2490
+ (Default: 0.97)
2491
+
2492
+ Returns:
2493
+ torch.Tensor: De-emphasized waveform, with shape `(..., N)`.
2494
+ """
2495
+ a_coeffs = torch.tensor([1.0, -coeff], dtype=waveform.dtype, device=waveform.device)
2496
+ b_coeffs = torch.tensor([1.0, 0.0], dtype=waveform.dtype, device=waveform.device)
2497
+ return torchaudio.functional.lfilter(waveform, a_coeffs=a_coeffs, b_coeffs=b_coeffs)
2498
+
2499
+
2500
+ def frechet_distance(mu_x, sigma_x, mu_y, sigma_y):
2501
+ r"""Computes the Fréchet distance between two multivariate normal distributions :cite:`dowson1982frechet`.
2502
+
2503
+ Concretely, for multivariate Gaussians :math:`X(\mu_X, \Sigma_X)`
2504
+ and :math:`Y(\mu_Y, \Sigma_Y)`, the function computes and returns :math:`F` as
2505
+
2506
+ .. math::
2507
+ F(X, Y) = || \mu_X - \mu_Y ||_2^2
2508
+ + \text{Tr}\left( \Sigma_X + \Sigma_Y - 2 \sqrt{\Sigma_X \Sigma_Y} \right)
2509
+
2510
+ Args:
2511
+ mu_x (torch.Tensor): mean :math:`\mu_X` of multivariate Gaussian :math:`X`, with shape `(N,)`.
2512
+ sigma_x (torch.Tensor): covariance matrix :math:`\Sigma_X` of :math:`X`, with shape `(N, N)`.
2513
+ mu_y (torch.Tensor): mean :math:`\mu_Y` of multivariate Gaussian :math:`Y`, with shape `(N,)`.
2514
+ sigma_y (torch.Tensor): covariance matrix :math:`\Sigma_Y` of :math:`Y`, with shape `(N, N)`.
2515
+
2516
+ Returns:
2517
+ torch.Tensor: the Fréchet distance between :math:`X` and :math:`Y`.
2518
+ """
2519
+ if len(mu_x.size()) != 1:
2520
+ raise ValueError(f"Input mu_x must be one-dimensional; got dimension {len(mu_x.size())}.")
2521
+ if len(sigma_x.size()) != 2:
2522
+ raise ValueError(f"Input sigma_x must be two-dimensional; got dimension {len(sigma_x.size())}.")
2523
+ if sigma_x.size(0) != sigma_x.size(1) != mu_x.size(0):
2524
+ raise ValueError("Each of sigma_x's dimensions must match mu_x's size.")
2525
+ if mu_x.size() != mu_y.size():
2526
+ raise ValueError(f"Inputs mu_x and mu_y must have the same shape; got {mu_x.size()} and {mu_y.size()}.")
2527
+ if sigma_x.size() != sigma_y.size():
2528
+ raise ValueError(
2529
+ f"Inputs sigma_x and sigma_y must have the same shape; got {sigma_x.size()} and {sigma_y.size()}."
2530
+ )
2531
+
2532
+ a = (mu_x - mu_y).square().sum()
2533
+ b = sigma_x.trace() + sigma_y.trace()
2534
+ c = torch.linalg.eigvals(sigma_x @ sigma_y).sqrt().real.sum()
2535
+ return a + b - 2 * c