RFC: The future of Kaldi compliance module

[mthrok]

Request For Comment: The future of Kaldi-compatible features

Problems

torchaudio.compliance.kaldi implements functionalities that tries to reproduce Kaldi's feature extractions, and this module has many issues, and causing headache for maintainers.

1. Inconsistent design
   While the rest of the torchaudio library is standardized to work with floating-point Tensor with value range [-1.0, 1.0], the ported Kaldi implementations are not necessarily following this (example)
2. Does not support batch dimension well. support batching for kaldi compliant feature extraction functions #675 torchaudio.compliance.kaldi.fbank #1245
3. compliance is not the right name Compliance should be named compatibility (or similar) #281
4. While it's called compliance, it does not match with the result of Kaldi's CLI
   I think the natural expectation that users get from the compliance.kaldi is that you can get a result that matches Kaldi, (and possibly in an easy manner)
   The spectrogram computed by "torchaudio.compliance.kaldi.spectrogram" and "compute-spectrogram-feats" are different #332, Problems with Kaldi MFCCs #328, Fbank features are different from Kaldi Fbank #400,
5. It's slow more efficient resample module #908 Torchaudio resampling could be faster and simpler #1057
   The code translated verbatim from the original Kaldi's C++ code typically is slower than the original implementation. (100x is not uncommon) To overcome this the code has to be changed to adopt PyTorch-style operation. (Kaldi does par element access very efficiently, which incurs a lot of overhead for PyTorch framework) This incurs huge maintenance overhead, 1 for implementing our custom code and 2 for catching up with upstream Kaldi.

Possible solutions

Before going into the detail of possible solutions, I note that removing Kaldi compatibility features from torchaudio was also mentioned as a possibility. It was not included in the original plan of torchaudio. I am not familiar on this matter and I am not advocating this but keeping it as a possibility.

The following describes some partial solutions and considerations I have put so far.

Module location.

For problems like 1 (inconsistent design) and 3 (naming), we can add a new interface to torchaudio.functional.

We have some Kaldi-compatible functions in torchaudio.functional module, such as sliding_window_cmn and compute_kaldi_pitch. They are placed under torchaudio.functional because they are confirmed to work on floating-point Tensors and their behaviors are consistent with other feature implementations.

We can do the similar thing so that we can provide the same set of Kaldi features abut that works on float-point Tensor like the other feature extractions too. We can start by adding interface to torchaudio.functional that does value normalization and call the corresponding function in compliance.kaldi module. Then deprecate the compliance.kaldi module and eventually remove it.

Implementation

For problems like 4 (numerical parity) and 5(speed). We have other approaches to make Kaldi features available and compatible with PyTorch.

1. Keep porting Kaldi in Python
   The current state.
2. Build and bind libkaldi-feat
   This will resolve the most of the headaches. The usual maintenance cost will be drastically reduced, though we need to figure out a robust way to build Kaldi with MKL that PyTorch is using. Users will get what they would naturally expect. The same result as Kaldi.
3. Re-building Kaldi's vector/matrix classes with PyTorch's Tensor class.
   This is somewhat like a hybrid approach of 1 and 2. Detail can be found in https://mthrok.github.io/tkaldi/. This is how I added Pitch feature in Add Kaldi Pitch feature #1243

The following table summarizes the pros and cons.

For problems 2 (batch support), one simple approach that can be added is to use at::parallel_for to parallelize the batch computation. This can be applied if the core implementation is in C++.

		Bind libkaldi-feat	Reimpl libkaldi-matrix	Reimplement in Python
Numerical Compatibility		✅ Baseline	✅ Easy	🚫 Extremely Difficult (None of the existing features meet this criteria)
Execution Speed		✅ Baseline	🚫 Slow (60x~)	🚫 Slow on CPU Even slower on GPU
Dev Scalability (Easy to add another feature?)		✅ Easy (Add new wrapper function)	🍰 Small Effort (Extend Matrix interface as needed)	🚫 Very Time Consuming (Understand, translate the code, verify the result)
Upstream Adaptation (Easy to follow the changes Kaldi makes?)		✅ Easy (Change the upstream commit, and wrapper function)	✅ Easy (Pull the upstream code, Change the wrapper function)	🚫 Practically impossible (I do not know where the existing code comes from)
Maint Cost	Initial build setup cost	🤨 Moderate Effort (Custom Build + MKL setup)	🍰 Small Effort (Custom Build)	✅ None
	Long term maintenance	🍰 Small Effort (Mostly about wrapper func)	🍰 Small effort (Mostly about wrapper func)	🚫 High (All the related codes, 1K LoC)

[mthrok]

cc Kaldi experts @danpovey @csukuangfj @pzelasko @mravanelli @sw005320

Since I have very limited experience in Kaldi and do not know how Kaldi project thinks of PyTorch/torchaudio, it is nice if you can your feedback. Thanks!

[pzelasko]

The binding of libkaldi-feat looks like a good solution. We are already following this approach for some components in K2/Snowfall (e.g. kaldialign for WER/alignment computation and kaldilm for LM compilation into FSA).
However my concerns are:

• will it run on GPU?
• will it be compatible with autograd?
• will you be able to efficiently batch it?

If the answer to any of these is no, then we should consider other options.

[danpovey]

Some comments:

• The existing Kaldi project is kind of in maintenance mode right now, as I am working on new projects (k2/lhotse/snowfall) that are intended to replace it.
• I suspect many of the uses of the kaldi compliance module right now are some kind of legacy support, e.g. people who have converted Kaldi models to PyTorch for inference.
• There is probably a quite limited set of configurations that people are actually using in practice.

I suspect the users of torchaudio's Kaldi compliance fall into very distinct categories: (1) people training new models who don't care about compatibility, (2), and this category is decreasing over time: people porting models from Kaldi who are probably using one of Kaldi's standard feature configurations.

For feature extractors that need to support autograd, I think that can be done separately from the Kaldi compliance module if you end up doing the Kaldi compliance in a way that's not convenient for that.
In principle Kaldi's matrix library can be configured to use various BLAS implementations; IMO it would be more ideal to figure out how to use whatever BLAS is being used by NumPy or PyTorch and try to use that, but the build issues might be complex.

It might be possible to start from the kaldi10feat project (https://github.com/danpovey/kaldi10feat) and modify it to be compatible with current Kaldi. (Note: kaldi10 is a project that I was planning but which does not exist and will not exist). It has a NumPy dependency. It was intended to be a simple feature configuration without all the bells and whistles and options of Kaldi's, but with reasonable defaults. Some of the differences include:

• I think it gets the frame count in a way that corresponds to --snip-edges=false in Kaldi, which is a simpler formula and tends to be less hassle but is not Kaldi's default (for compatibility reasons we were never able to change the default in Kaldi).
• It doesn't use pre-emphasis. (--preemph-coeff=0.0)
• It puts the signal in the range [-1,1] instead of [-32k, 32k]
• It doesn't use dithering (--dither=0.0)
• It sets the energy floor to a different value than Kaldi's default, see the code for details.
• Currently it only goes as far as mel filterbank values, without supporting MFCC.
  Those might be all the differences, if I am not forgetting anything.

[csukuangfj]

I would propose another alternative: reimplement, instead of binding, libkaldi-feat using torch::Tensor in C++
and provide a Python interface for it.

I've got a preliminary working version that produces the same output as kaldi's compute-fbank-feats using
default parameters with --dither=0.

The following sample wave abc.wav

sox -n -r 16000 -b 16 abc.wav synth 320 sine 300

is used to benchmark the implementation.

I've modified compute-fbank-feats.cc to include only the feature extraction time. The change is shown below

diff --git a/src/featbin/compute-fbank-feats.cc b/src/featbin/compute-fbank-feats.cc
index e52b30baf..63735c985 100644
--- a/src/featbin/compute-fbank-feats.cc
+++ b/src/featbin/compute-fbank-feats.cc
@@ -18,6 +18,8 @@
 // See the Apache 2 License for the specific language governing permissions and
 // limitations under the License.

+#include <chrono>
+
 #include "base/kaldi-common.h"
 #include "feat/feature-fbank.h"
 #include "feat/wave-reader.h"
@@ -145,8 +147,16 @@ int main(int argc, char *argv[]) {
       SubVector<BaseFloat> waveform(wave_data.Data(), this_chan);
       Matrix<BaseFloat> features;
       try {
+        std::chrono::steady_clock::time_point begin = std::chrono::steady_clock::now();
         fbank.ComputeFeatures(waveform, wave_data.SampFreq(),
                               vtln_warp_local, &features);
+        std::chrono::steady_clock::time_point end = std::chrono::steady_clock::now();
+        std::cout << "Time difference = "
+                  << std::chrono::duration_cast<std::chrono::microseconds>(
+                         end - begin)
+                             .count() /
+                         1000000.
+                  << "[s]" << std::endl;
       } catch (...) {
         KALDI_WARN << "Failed to compute features for utterance " << utt;
         continue;

And the benchmark command is

(base) fangjun:~/open-source/kaldi/src/featbin$ ./compute-fbank-feats --dither=0 scp:1.scp ark:1.ark
./compute-fbank-feats --dither=0 scp:1.scp ark:1.ark
Time difference = 0.313595[s]
LOG (compute-fbank-feats[5.5.880~4-3e446]:main():compute-fbank-feats.cc:195)  Done 1 out of 1 utterances.

1.scp contains

1 abc.wav

The execution time is about 0.3 seconds and does not vary too much across multiple runs.

The reimplementation, which is still in its early draft (https://github.com/csukuangfj/kaldifeat), needs only about 0.1 seconds.
The benchmark command for it is

git clone https://github.com/csukuangfj/kaldifeat
cd kaldifeat
mkdir build
cd build
make -j
./kaldifeat/python/tests/test_kaldifeat.py

The outputs are

(py38) fangjun:~/open-source/kaldifeat/build$ ../kaldifeat/python/tests/test_kaldifeat.py
elapsed seconds: 0.096343
(py38) fangjun:~/open-source/kaldifeat/build$ ../kaldifeat/python/tests/test_kaldifeat.py
elapsed seconds: 0.083852
(py38) fangjun:~/open-source/kaldifeat/build$ ../kaldifeat/python/tests/test_kaldifeat.py
elapsed seconds: 0.141683

[csukuangfj]

will it be compatible with autograd?

@pzelasko

As there are no learnable parameters in the feature extraction, could you explain what's the use of autograd
in the context of feature extraction?

[pzelasko]

@csukuangfj Adversarial attacks and defense research

[csukuangfj]

It's even faster when CUDA is available:

(py38) fangjun:~/open-source/kaldifeat/build$ ../kaldifeat/python/tests/test_kaldifeat.py
elapsed seconds cpu: 0.153896
elapsed seconds cuda:0: 0.005204
(py38) fangjun:~/open-source/kaldifeat/build$ ../kaldifeat/python/tests/test_kaldifeat.py
elapsed seconds cpu: 0.080523
elapsed seconds cuda:0: 0.00535
(py38) fangjun:~/open-source/kaldifeat/build$ ../kaldifeat/python/tests/test_kaldifeat.py
elapsed seconds cpu: 0.085813
elapsed seconds cuda:0: 0.005682

[danpovey]

@csukuangfj let's wait to see whether that's the direction they want to go in before spending lots of time on this.
Since your repo is public, if they want to use that as a starting point they can easily do that, but right now I think getting our new model working is more urgent.

[csukuangfj]

@csukuangfj let's wait to see whether that's the direction they want to go in before spending lots of time on this.

Since your repo is public, if they want to use that as a starting point they can easily do that, but right now I think getting our new model working is more urgent.

Will focus on our snowfall training.

[mthrok]

Hi @pzelasko @danpovey @csukuangfj

Thanks for sharing your thoughts.

So torchaudio team wants to resolve the issue with the current module. That is inconsistent design / bad performance / high maintenance cost. But to do that we need a direction and that's where we would like to hear from the users.

I was pushing for binding libkaldi only because this easily solves performance and maintenance cost. This is very attractive option to me since torchaudio's engineering resource is very limited. Of course, with this approach, we cannot add anew feature like differentiability, which blocks the research for the defense against adversarial attack, differential DSP or E2E training with downstream task.

Meanwhile @cpuhrsch shares the similar view with @csukuangfj. We should be pushing the boundary taking advantage of PyTorch provides, which I do agree, if we have infinite amount of engineering resource.

My understanding of @danpovey's is that complete API compatibility with Kaldi is not necessary and we can reduce the options, which could simplify some implementations.

Does either of K2 or Lhotse provide feature extractions? My understanding of the scene of Speech Recognition is that K2 handles WFST(LM), PyTorch model handles AM, Lhotse handles data handling and training recipes. I am not sure if Lhotse is going to implement the feature extractions on its own, but if Lhotse is going to add an efficient feature extraction, torchaudio can just retire the current module and do nothing. In contrary, if adding efficient implementations in torchaudio is beneficial to the whole ecosystem, we can prioritize that too.

So I think the long term direction is to

• Provide features that are useful for future audience of the ecosystem (K2, Lhotse, PyTorch, torchaudio etc...)
• Efficient implementation with GPU support
• Full backward API compatibility with Kaldi's feature is not necessary, but having numerical parity with basic configuration would be nice.
• Autograd support is valuable

@csukuangfj
Your kaldifeat looks great. I assume that that is not your primal work. I have integrated it to torchaudio #1326, I will try to see if this approach can support Autograd.

[danpovey]

For the time being our plan was to rely on torchaudio for feature extraction, I'm afraid.
That's what we are currently doing.

[h-vetinari]

Hey all

I've been trying to package torchaudio for conda-forge, and one of the first things I wanted to do (based on how conda-forge does things) is reuse our own kaldi - which we have built such that we can change the BLAS implementation, in particular it should be compatible with MKL.

The problem I have is that I don't know how to patch this, because it's unclear how much the additional implementation is wrapper/extension/reimplementation.

For the moment I think I'll have to disable it, even though that's a pity because it seems we have all the necessary pieces...