Use the torchaudio resampler instead of the custom implementation (speechbrain#2410)

asumagic · web-flow · commit e602161f4d30 · 2024-02-14T10:02:53.000+01:00
* Use the torchaudio resampler for better performance

* Add intersphinx mapping for torchaudio

* Resampler fixes and better docstring

* Make the linter happy

* Make the linter happy again

* Keep compat by migrating the resampler to whatever input device
diff --git a/docs/conf.py b/docs/conf.py
@@ -61,6 +61,7 @@
     "python": ("https://docs.python.org/", None),
     "numpy": ("http://docs.scipy.org/doc/numpy/", None),
     "torch": ("https://pytorch.org/docs/master/", None),
+    "torchaudio": ("https://pytorch.org/audio/stable/", None),
 }
 
 # AUTODOC:
diff --git a/speechbrain/augment/time_domain.py b/speechbrain/augment/time_domain.py
@@ -11,10 +11,10 @@
 """
 
 # Importing libraries
-import math
 import random
 import torch
 import torch.nn.functional as F
+import torchaudio
 from speechbrain.dataio.legacy import ExtendedCSVDataset
 from speechbrain.dataio.dataloader import make_dataloader
 from speechbrain.processing.signal_processing import (
@@ -508,21 +508,22 @@ def forward(self, waveform):
 
 
 class Resample(torch.nn.Module):
-    """This class resamples an audio signal using sinc-based interpolation.
-
-    It is a modification of the `resample` function from torchaudio
-    (https://pytorch.org/audio/stable/tutorials/audio_resampling_tutorial.html)
+    """This class resamples audio using the
+    :class:`torchaudio resampler <torchaudio.transforms.Resample>` based on
+    sinc interpolation.
 
     Arguments
     ---------
     orig_freq : int
         the sampling frequency of the input signal.
     new_freq : int
         the new sampling frequency after this operation is performed.
-    lowpass_filter_width : int
-        Controls the sharpness of the filter, larger numbers result in a
-        sharper filter, but they are less efficient. Values from 4 to 10 are
-        allowed.
+    *args
+        additional arguments forwarded to the
+        :class:`torchaudio.transforms.Resample` constructor
+    **kwargs
+        additional keyword arguments forwarded to the
+        :class:`torchaudio.transforms.Resample` constructor
 
     Example
     -------
@@ -537,49 +538,28 @@ class Resample(torch.nn.Module):
     torch.Size([1, 26087])
     """
 
-    def __init__(self, orig_freq=16000, new_freq=16000, lowpass_filter_width=6):
+    def __init__(self, orig_freq=16000, new_freq=16000, *args, **kwargs):
         super().__init__()
+
         self.orig_freq = orig_freq
         self.new_freq = new_freq
-        self.lowpass_filter_width = lowpass_filter_width
-
-        # Compute rate for striding
-        self._compute_strides()
-        assert self.orig_freq % self.conv_stride == 0
-        assert self.new_freq % self.conv_transpose_stride == 0
 
-    def _compute_strides(self):
-        """Compute the phases in polyphase filter.
-
-        (almost directly from torchaudio.compliance.kaldi)
-        """
-
-        # Compute new unit based on ratio of in/out frequencies
-        base_freq = math.gcd(self.orig_freq, self.new_freq)
-        input_samples_in_unit = self.orig_freq // base_freq
-        self.output_samples = self.new_freq // base_freq
-
-        # Store the appropriate stride based on the new units
-        self.conv_stride = input_samples_in_unit
-        self.conv_transpose_stride = self.output_samples
+        self.resampler = torchaudio.transforms.Resample(
+            orig_freq=orig_freq, new_freq=new_freq, *args, **kwargs,
+        )
 
     def forward(self, waveforms):
         """
         Arguments
         ---------
         waveforms : tensor
             Shape should be `[batch, time]` or `[batch, time, channels]`.
-        lengths : tensor
-            Shape should be a single dimension, `[batch]`.
 
         Returns
         -------
         Tensor of shape `[batch, time]` or `[batch, time, channels]`.
         """
 
-        if not hasattr(self, "first_indices"):
-            self._indices_and_weights(waveforms)
-
         # Don't do anything if the frequencies are the same
         if self.orig_freq == self.new_freq:
             return waveforms
@@ -593,8 +573,15 @@ def forward(self, waveforms):
         else:
             raise ValueError("Input must be 2 or 3 dimensions")
 
+        # If necessary, migrate the resampler to the current device, for
+        # backwards compat with scripts that do not call `resampler.to()`
+        # themselves.
+        # Please do not reuse the sample resampler for tensors that live on
+        # different devices, though.
+        self.resampler.to(waveforms.device)  # in-place
+
         # Do resampling
-        resampled_waveform = self._perform_resample(waveforms)
+        resampled_waveform = self.resampler(waveforms)
 
         if unsqueezed:
             resampled_waveform = resampled_waveform.squeeze(1)
@@ -603,219 +590,6 @@ def forward(self, waveforms):
 
         return resampled_waveform
 
-    def _perform_resample(self, waveforms):
-        """Resamples the waveform at the new frequency.
-
-        This matches Kaldi's OfflineFeatureTpl ResampleWaveform which uses a
-        LinearResample (resample a signal at linearly spaced intervals to
-        up/downsample a signal). LinearResample (LR) means that the output
-        signal is at linearly spaced intervals (i.e the output signal has a
-        frequency of `new_freq`). It uses sinc/bandlimited interpolation to
-        upsample/downsample the signal.
-
-        (almost directly from torchaudio.compliance.kaldi)
-
-        https://ccrma.stanford.edu/~jos/resample/
-        Theory_Ideal_Bandlimited_Interpolation.html
-
-        https://github.com/kaldi-asr/kaldi/blob/master/src/feat/resample.h#L56
-
-        Arguments
-        ---------
-        waveforms : tensor
-            The batch of audio signals to resample.
-
-        Returns
-        -------
-        The waveforms at the new frequency.
-        """
-
-        # Compute output size and initialize
-        batch_size, num_channels, wave_len = waveforms.size()
-        window_size = self.weights.size(1)
-        tot_output_samp = self._output_samples(wave_len)
-        resampled_waveform = torch.zeros(
-            (batch_size, num_channels, tot_output_samp), device=waveforms.device
-        )
-        self.weights = self.weights.to(waveforms.device)
-
-        # Check weights are on correct device
-        if waveforms.device != self.weights.device:
-            self.weights = self.weights.to(waveforms.device)
-
-        # eye size: (num_channels, num_channels, 1)
-        eye = torch.eye(num_channels, device=waveforms.device).unsqueeze(2)
-
-        # Iterate over the phases in the polyphase filter
-        for i in range(self.first_indices.size(0)):
-            wave_to_conv = waveforms
-            first_index = int(self.first_indices[i].item())
-            if first_index >= 0:
-                # trim the signal as the filter will not be applied
-                # before the first_index
-                wave_to_conv = wave_to_conv[..., first_index:]
-
-            # pad the right of the signal to allow partial convolutions
-            # meaning compute values for partial windows (e.g. end of the
-            # window is outside the signal length)
-            max_index = (tot_output_samp - 1) // self.output_samples
-            end_index = max_index * self.conv_stride + window_size
-            current_wave_len = wave_len - first_index
-            right_padding = max(0, end_index + 1 - current_wave_len)
-            left_padding = max(0, -first_index)
-            wave_to_conv = torch.nn.functional.pad(
-                wave_to_conv, (left_padding, right_padding)
-            )
-            conv_wave = torch.nn.functional.conv1d(
-                input=wave_to_conv,
-                weight=self.weights[i].repeat(num_channels, 1, 1),
-                stride=self.conv_stride,
-                groups=num_channels,
-            )
-
-            # we want conv_wave[:, i] to be at
-            # output[:, i + n*conv_transpose_stride]
-            dilated_conv_wave = torch.nn.functional.conv_transpose1d(
-                conv_wave, eye, stride=self.conv_transpose_stride
-            )
-
-            # pad dilated_conv_wave so it reaches the output length if needed.
-            left_padding = i
-            previous_padding = left_padding + dilated_conv_wave.size(-1)
-            right_padding = max(0, tot_output_samp - previous_padding)
-            dilated_conv_wave = torch.nn.functional.pad(
-                dilated_conv_wave, (left_padding, right_padding)
-            )
-            dilated_conv_wave = dilated_conv_wave[..., :tot_output_samp]
-
-            resampled_waveform += dilated_conv_wave
-
-        return resampled_waveform
-
-    def _output_samples(self, input_num_samp):
-        """Based on LinearResample::GetNumOutputSamples.
-
-        LinearResample (LR) means that the output signal is at
-        linearly spaced intervals (i.e the output signal has a
-        frequency of ``new_freq``). It uses sinc/bandlimited
-        interpolation to upsample/downsample the signal.
-
-        (almost directly from torchaudio.compliance.kaldi)
-
-        Arguments
-        ---------
-        input_num_samp : int
-            The number of samples in each example in the batch.
-
-        Returns
-        -------
-        Number of samples in the output waveform.
-        """
-
-        # For exact computation, we measure time in "ticks" of 1.0 / tick_freq,
-        # where tick_freq is the least common multiple of samp_in and
-        # samp_out.
-        samp_in = int(self.orig_freq)
-        samp_out = int(self.new_freq)
-
-        tick_freq = abs(samp_in * samp_out) // math.gcd(samp_in, samp_out)
-        ticks_per_input_period = tick_freq // samp_in
-
-        # work out the number of ticks in the time interval
-        # [ 0, input_num_samp/samp_in ).
-        interval_length = input_num_samp * ticks_per_input_period
-        if interval_length <= 0:
-            return 0
-        ticks_per_output_period = tick_freq // samp_out
-
-        # Get the last output-sample in the closed interval,
-        # i.e. replacing [ ) with [ ]. Note: integer division rounds down.
-        # See http://en.wikipedia.org/wiki/Interval_(mathematics) for an
-        # explanation of the notation.
-        last_output_samp = interval_length // ticks_per_output_period
-
-        # We need the last output-sample in the open interval, so if it
-        # takes us to the end of the interval exactly, subtract one.
-        if last_output_samp * ticks_per_output_period == interval_length:
-            last_output_samp -= 1
-
-        # First output-sample index is zero, so the number of output samples
-        # is the last output-sample plus one.
-        num_output_samp = last_output_samp + 1
-
-        return num_output_samp
-
-    def _indices_and_weights(self, waveforms):
-        """Based on LinearResample::SetIndexesAndWeights
-
-        Retrieves the weights for resampling as well as the indices in which
-        they are valid. LinearResample (LR) means that the output signal is at
-        linearly spaced intervals (i.e the output signal has a frequency
-        of ``new_freq``). It uses sinc/bandlimited interpolation to
-        upsample/downsample the signal.
-
-        Returns
-        -------
-        - the place where each filter should start being applied
-        - the filters to be applied to the signal for resampling
-        """
-
-        # Lowpass filter frequency depends on smaller of two frequencies
-        min_freq = min(self.orig_freq, self.new_freq)
-        lowpass_cutoff = 0.99 * 0.5 * min_freq
-
-        assert lowpass_cutoff * 2 <= min_freq
-        window_width = self.lowpass_filter_width / (2.0 * lowpass_cutoff)
-
-        assert lowpass_cutoff < min(self.orig_freq, self.new_freq) / 2
-        output_t = torch.arange(
-            start=0.0, end=self.output_samples, device=waveforms.device
-        )
-        output_t /= self.new_freq
-        min_t = output_t - window_width
-        max_t = output_t + window_width
-
-        min_input_index = torch.ceil(min_t * self.orig_freq)
-        max_input_index = torch.floor(max_t * self.orig_freq)
-        num_indices = max_input_index - min_input_index + 1
-
-        max_weight_width = num_indices.max()
-        j = torch.arange(max_weight_width, device=waveforms.device)
-        input_index = min_input_index.unsqueeze(1) + j.unsqueeze(0)
-        delta_t = (input_index / self.orig_freq) - output_t.unsqueeze(1)
-
-        weights = torch.zeros_like(delta_t)
-        inside_window_indices = delta_t.abs().lt(window_width)
-
-        # raised-cosine (Hanning) window with width `window_width`
-        weights[inside_window_indices] = 0.5 * (
-            1
-            + torch.cos(
-                2
-                * math.pi
-                * lowpass_cutoff
-                / self.lowpass_filter_width
-                * delta_t[inside_window_indices]
-            )
-        )
-
-        t_eq_zero_indices = delta_t.eq(0.0)
-        t_not_eq_zero_indices = ~t_eq_zero_indices
-
-        # sinc filter function
-        weights[t_not_eq_zero_indices] *= torch.sin(
-            2 * math.pi * lowpass_cutoff * delta_t[t_not_eq_zero_indices]
-        ) / (math.pi * delta_t[t_not_eq_zero_indices])
-
-        # limit of the function at t = 0
-        weights[t_eq_zero_indices] *= 2 * lowpass_cutoff
-
-        # size (output_samples, max_weight_width)
-        weights /= self.orig_freq
-
-        self.first_indices = min_input_index
-        self.weights = weights
-
 
 class DropFreq(torch.nn.Module):
     """This class drops a random frequency from the signal.

Original file line number	Diff line number	Diff line change
`@@ -61,6 +61,7 @@`
`61`	`61`	`"python": ("https://docs.python.org/", None),`
`62`	`62`	`"numpy": ("http://docs.scipy.org/doc/numpy/", None),`
`63`	`63`	`"torch": ("https://pytorch.org/docs/master/", None),`
	`64`	`+ "torchaudio": ("https://pytorch.org/audio/stable/", None),`
`64`	`65`	`}`
`65`	`66`
`66`	`67`	`# AUTODOC:`