Repository: pytorch/audio

Issue:
Add batch dimension inside the computation of lfilter
## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

Add batch dimension inside the computation of lfilter to make thing fast

## Motivation

It is recently possible to integrate digital filters inside the neural network because of the backprop support for lfilter (https://github.com/pytorch/audio/pull/1310). 

However, currently, the filter only supports the 1D waveform, 1D a_coefficients, 1D b_coeficient. It is not possible to include the batch dimension in the computation. The only way to do it is naive for-loop, but it is very time-consuming. 

Therefore, I modify some code to let the computation include batch dimension, but as the issues (https://github.com/pytorch/audio/issues/1408) say, compute lfilter without CPU optimization is slow. 

I am ask for help whether any people have any other idea to add dimension inside the computation of lfilter ? 


<!-- Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too -->

## Pitch
I would love to see that any other elegant way to include the batch dimension when running the lfilter or 
maybe someone know how to include batch dimension inside the CPU optimization loop which is used in lfilter. 

<!-- A clear and concise description of what you want to happen. -->

## Alternatives

<!-- A clear and concise description of any alternative solutions or features you've considered, if any. -->

I modify the original lfilter to include the batch dimension. 
The major modification is 

1.  calculate `windowed_input_signal` by `F.conv1d` with group\
2. replace `input_signal_windows.t()` to `input_signal_windows.permute(2, 0, 1)` (inside `_lfilter_core_generic_loop`)
3. replace `addmv_` to `einsum` (inside `_lfilter_core_generic_loop`)

```python
def lfilter_batch(waveform, a_coeffs, b_coeffs):
    
    batch    = waveform.size(0)
    n_sample = waveform.size(-1)
    n_order  = a_coeffs.size(-1)
        
    # Pad the input and create output
    padded_waveform        = F.pad(waveform, [n_order - 1, 0])
    padded_output_waveform = torch.zeros_like(padded_waveform)
    
    
    # Set up the coefficients matrix
    # Flip coefficients' order
    a_coeffs_flipped = a_coeffs.flip(-1)
    b_coeffs_flipped = b_coeffs.flip(-1)
    
    # calculate windowed_input_signal in parallel using convolution
    
    input_signal_windows = F.conv1d(padded_waveform.permute(1, 0, 2), b_coeffs_flipped.unsqueeze(1), groups=batch)
    input_signal_windows = input_signal_windows.permute(1, 0, 2)
    a_coeffs         = a_coeffs.unsqueeze(1)
    a_coeffs_flipped = a_coeffs_flipped.unsqueeze(1)
        
    input_signal_windows.div_(a_coeffs[:, :, :1])
    a_coeffs_flipped.div_(a_coeffs[:, :, :1])
    a_coeffs_flipped = a_coeffs_flipped.squeeze(1)

    
    padded_output_waveform = _lfilter_core_generic_loop_batch(input_signal_windows, a_coeffs_flipped, padded_output_waveform)
    output = padded_output_waveform[:, :, n_order - 1:]
    return output


def _lfilter_core_execute(i_sample, n_order, o0, a_coeffs_flipped, out):    
    windowed_output_signal = out[:, :, i_sample:i_sample + n_order]   
    out[:, :, i_sample + n_order - 1] = o0 - torch.einsum('bcd, bd->bc', windowed_output_signal, a_coeffs_flipped)


def _lfilter_core_generic_loop_batch(input_signal_windows, a_coeffs_flipped, padded_output_waveform):
    n_order = a_coeffs_flipped.size(-1)
    out     = padded_output_waveform
    input_signal_windows_transpose = input_signal_windows.permute(2, 0, 1)
    for (i_sample, o0) in enumerate(input_signal_windows_transpose):
        _lfilter_core_execute(i_sample, n_order, o0, a_coeffs_flipped, out)
    return out
```

## Additional context

<!-- Add any other context or screenshots about the feature request here. -->



Discussion:

Resolving patch:
diff --git a/torchaudio/functional/filtering.py b/torchaudio/functional/filtering.py
--- a/torchaudio/functional/filtering.py
+++ b/torchaudio/functional/filtering.py
@@ -930,6 +930,7 @@ def lfilter(
     a_coeffs: Tensor,
     b_coeffs: Tensor,
     clamp: bool = True,
+    batching: bool = True
 ) -> Tensor:
     r"""Perform an IIR filter by evaluating difference equation.
 
@@ -948,6 +949,10 @@ def lfilter(
                                 Lower delays coefficients are first, e.g. ``[b0, b1, b2, ...]``.
                                 Must be same size as a_coeffs (pad with 0's as necessary).
         clamp (bool, optional): If ``True``, clamp the output signal to be in the range [-1, 1] (Default: ``True``)
+        batching (bool, optional): Activate when coefficients are in 2D. If ``True``, then waveform should be at least
+                                    2D, and the size of second axis from last should equals to ``num_filters``.
+                                    The output can be expressed as ``output[..., i, :] = lfilter(waveform[..., i, :],
+                                    a_coeffs[i], b_coeffs[i], clamp=clamp, batching=False)``. (Default: ``True``)
 
     Returns:
         Tensor: Waveform with dimension of either ``(..., num_filters, time)`` if ``a_coeffs`` and ``b_coeffs``
@@ -957,7 +962,11 @@ def lfilter(
     assert a_coeffs.ndim <= 2
 
     if a_coeffs.ndim > 1:
-        waveform = torch.stack([waveform] * a_coeffs.shape[0], -2)
+        if batching:
+            assert waveform.ndim > 1
+            assert waveform.shape[-2] == a_coeffs.shape[0]
+        else:
+            waveform = torch.stack([waveform] * a_coeffs.shape[0], -2)
     else:
         a_coeffs = a_coeffs.unsqueeze(0)
         b_coeffs = b_coeffs.unsqueeze(0)


Test patch:
diff --git a/test/torchaudio_unittest/functional/autograd_impl.py b/test/torchaudio_unittest/functional/autograd_impl.py
--- a/test/torchaudio_unittest/functional/autograd_impl.py
+++ b/test/torchaudio_unittest/functional/autograd_impl.py
@@ -1,4 +1,5 @@
 from typing import Callable, Tuple
+from functools import partial
 import torch
 from parameterized import parameterized
 from torch import Tensor
@@ -62,6 +63,15 @@ def test_lfilter_all_inputs(self):
     def test_lfilter_filterbanks(self):
         torch.random.manual_seed(2434)
         x = get_whitenoise(sample_rate=22050, duration=0.01, n_channels=3)
+        a = torch.tensor([[0.7, 0.2, 0.6],
+                          [0.8, 0.2, 0.9]])
+        b = torch.tensor([[0.4, 0.2, 0.9],
+                          [0.7, 0.2, 0.6]])
+        self.assert_grad(partial(F.lfilter, batching=False), (x, a, b))
+
+    def test_lfilter_batching(self):
+        torch.random.manual_seed(2434)
+        x = get_whitenoise(sample_rate=22050, duration=0.01, n_channels=2)
         a = torch.tensor([[0.7, 0.2, 0.6],
                           [0.8, 0.2, 0.9]])
         b = torch.tensor([[0.4, 0.2, 0.9],
diff --git a/test/torchaudio_unittest/functional/batch_consistency_test.py b/test/torchaudio_unittest/functional/batch_consistency_test.py
--- a/test/torchaudio_unittest/functional/batch_consistency_test.py
+++ b/test/torchaudio_unittest/functional/batch_consistency_test.py
@@ -217,3 +217,18 @@ def test_compute_kaldi_pitch(self):
         batch = waveform.view(self.batch_size, n_channels, waveform.size(-1))
         self.assert_batch_consistency(
             F.compute_kaldi_pitch, batch, sample_rate=sample_rate)
+
+    def test_lfilter(self):
+        signal_length = 2048
+        torch.manual_seed(2434)
+        x = torch.randn(self.batch_size, signal_length)
+        a = torch.rand(self.batch_size, 3)
+        b = torch.rand(self.batch_size, 3)
+
+        batchwise_output = F.lfilter(x, a, b, batching=True)
+        itemwise_output = torch.stack([
+            F.lfilter(x[i], a[i], b[i])
+            for i in range(self.batch_size)
+        ])
+
+        self.assertEqual(batchwise_output, itemwise_output)
diff --git a/test/torchaudio_unittest/functional/functional_impl.py b/test/torchaudio_unittest/functional/functional_impl.py
--- a/test/torchaudio_unittest/functional/functional_impl.py
+++ b/test/torchaudio_unittest/functional/functional_impl.py
@@ -80,7 +80,7 @@ def test_lfilter_shape(self, input_shape, coeff_shape, target_shape):
         waveform = torch.rand(*input_shape, dtype=self.dtype, device=self.device)
         b_coeffs = torch.rand(*coeff_shape, dtype=self.dtype, device=self.device)
         a_coeffs = torch.rand(*coeff_shape, dtype=self.dtype, device=self.device)
-        output_waveform = F.lfilter(waveform, a_coeffs, b_coeffs)
+        output_waveform = F.lfilter(waveform, a_coeffs, b_coeffs, batching=False)
         assert input_shape == waveform.size()
         assert target_shape == output_waveform.size()