Add batch dimension inside the computation of lfilter

[ChenPaulYu]

🚀 Feature

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 (#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 (#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 ?

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.

Alternatives

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)

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

[cpuhrsch]

cc @mthrok @yoyololicon

[yoyolicoris]

@ChenPaulYu
Awesome, I would like to see this happen.

To include batch dimension inside the CPU optimization loop, we can add one extra for-loop or use multithreading along the batch dimension.

Also, could you help me benchmark this line

torch.einsum('bcd, bd->bc', windowed_output_signal, a_coeffs_flipped)

compare with the following alternative implementation:

(windowed_output_signal @ a_coeffs_flipped.unsqueeze(2))[:, :, 0]

cuz einsum may introduce some extra overhead.

[ChenPaulYu]

Hi @yoyololicon ,

That is sound nice for the c++ implementation idea, I will also take look at how to write c++ extension for pytorch.
Of course, I can help to benchmark the einsum.

[mthrok]

Hi @ChenPaulYu

Thanks for the suggestion. I think this is a good idea.
One simple approach in C++ to make it parallel along batch dimension is to use at::parallel_for, like in

audio/torchaudio/csrc/kaldi.cpp

Lines 76 to 84 in 3ab521f

	auto batch_size = wave_.size(0);
	std::vector<torch::Tensor> results(batch_size);
	at::parallel_for(0, batch_size, 1, [&](int64_t begin, int64_t end) {
	for (auto i = begin; i < end; ++i) {
	results[i] = compute_kaldi_pitch(wave_.index({i}), opts);
	}
	});
	return torch::stack(results, 0);
	}

but I am not sure how robust at::parallel_for is, so we will need to run a benchmark. (I suspect we might need to set a collect compilation flag to do this. At the moment, torchaudio's build step does not set any OpenMP flag.)

Is your Python-level modification faster than the implementation in master branch?

[mthrok]

Note: #1483 has been landed and now torchaudio's build process can handle CUDA code as well.

Adding the source code and listing it here should do the trick.

[yoyolicoris]

@ChenPaulYu We have landed the batching feature in the new commit 8094751.
Please give it a try and will close this issue.