Add support for pickling. (#17)

* Add pickle support to FrameExtractionOptions.

* Add pickle support to MelBanksOptions.

* Add pickle support.

* Update README to mention other projects.

* Fix style issues.

README.md

@@ -213,6 +213,26 @@ for more examples.
 - ``kaldifeat`` supports batch processing as well as chunk processing
 - ``kaldifeat`` uses the same options as `Kaldi`'s `compute-fbank-feats` and `compute-mfcc-feats`
 
+# Usage in other projects
+
+## icefall
+
+[icefall](https://github.com/k2-fsa/icefall) uses kaldifeat to extract features for a pre-trained model.
+
+See <https://github.com/k2-fsa/icefall/blob/master/egs/librispeech/ASR/conformer_ctc/pretrained.py>.
+
+## k2
+
+[k2](https://github.com/k2-fsa/k2) uses kaldifeat's C++ API.
+
+See <https://github.com/k2-fsa/k2/blob/v2.0-pre/k2/torch/csrc/features.cu>.
+
+## lhotse
+
+[lhotse](https://github.com/lhotse-speech/lhotse) uses kaldifeat to extract features on GPU.
+
+See <https://github.com/lhotse-speech/lhotse/blob/master/lhotse/features/kaldifeat.py>.
+
 # Installation
 
 ## From conda (Only for Linux + CUDA)

kaldifeat/python/csrc/feature-fbank.cc

@@ -4,6 +4,7 @@
 
 #include "kaldifeat/python/csrc/feature-fbank.h"
 
+#include <memory>
 #include <string>
 
 #include "kaldifeat/csrc/feature-fbank.h"
@@ -37,9 +38,12 @@ static void PybindFbankOptions(py::module &m) {
            [](const PyClass &self) -> std::string { return self.ToString(); })
       .def("as_dict",
            [](const PyClass &self) -> py::dict { return AsDict(self); })
-      .def_static("from_dict", [](py::dict dict) -> PyClass {
+      .def_static(
-        return FbankOptionsFromDict(dict);
+          "from_dict",
-      });
+          [](py::dict dict) -> PyClass { return FbankOptionsFromDict(dict); })
+      .def(py::pickle(
+          [](const PyClass &self) -> py::dict { return AsDict(self); },
+          [](py::dict dict) -> PyClass { return FbankOptionsFromDict(dict); }));
 }
 
 static void PybindFbank(py::module &m) {
@@ -49,7 +53,14 @@ static void PybindFbank(py::module &m) {
       .def("dim", &PyClass::Dim)
       .def_property_readonly("options", &PyClass::GetOptions)
       .def("compute_features", &PyClass::ComputeFeatures, py::arg("wave"),
-           py::arg("vtln_warp"));
+           py::arg("vtln_warp"))
+      .def(py::pickle(
+          [](const PyClass &self) -> py::dict {
+            return AsDict(self.GetOptions());
+          },
+          [](py::dict dict) -> std::unique_ptr<PyClass> {
+            return std::make_unique<PyClass>(FbankOptionsFromDict(dict));
+          }));
 }
 
 void PybindFeatureFbank(py::module &m) {

kaldifeat/python/csrc/feature-mfcc.cc

@@ -4,6 +4,7 @@
 
 #include "kaldifeat/python/csrc/feature-mfcc.h"
 
+#include <memory>
 #include <string>
 
 #include "kaldifeat/csrc/feature-mfcc.h"
@@ -37,9 +38,12 @@ void PybindMfccOptions(py::module &m) {
            [](const PyClass &self) -> std::string { return self.ToString(); })
       .def("as_dict",
            [](const PyClass &self) -> py::dict { return AsDict(self); })
-      .def_static("from_dict", [](py::dict dict) -> PyClass {
+      .def_static(
-        return MfccOptionsFromDict(dict);
+          "from_dict",
-      });
+          [](py::dict dict) -> PyClass { return MfccOptionsFromDict(dict); })
+      .def(py::pickle(
+          [](const PyClass &self) -> py::dict { return AsDict(self); },
+          [](py::dict dict) -> PyClass { return MfccOptionsFromDict(dict); }));
 }
 
 static void PybindMfcc(py::module &m) {
@@ -49,7 +53,14 @@ static void PybindMfcc(py::module &m) {
       .def("dim", &PyClass::Dim)
       .def_property_readonly("options", &PyClass::GetOptions)
       .def("compute_features", &PyClass::ComputeFeatures, py::arg("wave"),
-           py::arg("vtln_warp"));
+           py::arg("vtln_warp"))
+      .def(py::pickle(
+          [](const PyClass &self) -> py::dict {
+            return AsDict(self.GetOptions());
+          },
+          [](py::dict dict) -> std::unique_ptr<PyClass> {
+            return std::make_unique<PyClass>(MfccOptionsFromDict(dict));
+          }));
 }
 
 void PybindFeatureMfcc(py::module &m) {

kaldifeat/python/csrc/feature-plp.cc

@@ -4,6 +4,7 @@
 
 #include "kaldifeat/python/csrc/feature-plp.h"
 
+#include <memory>
 #include <string>
 
 #include "kaldifeat/csrc/feature-plp.h"
@@ -40,9 +41,12 @@ void PybindPlpOptions(py::module &m) {
            [](const PyClass &self) -> std::string { return self.ToString(); })
       .def("as_dict",
            [](const PyClass &self) -> py::dict { return AsDict(self); })
-      .def_static("from_dict", [](py::dict dict) -> PyClass {
+      .def_static(
-        return PlpOptionsFromDict(dict);
+          "from_dict",
-      });
+          [](py::dict dict) -> PyClass { return PlpOptionsFromDict(dict); })
+      .def(py::pickle(
+          [](const PyClass &self) -> py::dict { return AsDict(self); },
+          [](py::dict dict) -> PyClass { return PlpOptionsFromDict(dict); }));
 }
 
 static void PybindPlp(py::module &m) {
@@ -52,7 +56,14 @@ static void PybindPlp(py::module &m) {
       .def("dim", &PyClass::Dim)
       .def_property_readonly("options", &PyClass::GetOptions)
       .def("compute_features", &PyClass::ComputeFeatures, py::arg("wave"),
-           py::arg("vtln_warp"));
+           py::arg("vtln_warp"))
+      .def(py::pickle(
+          [](const PyClass &self) -> py::dict {
+            return AsDict(self.GetOptions());
+          },
+          [](py::dict dict) -> std::unique_ptr<PyClass> {
+            return std::make_unique<PyClass>(PlpOptionsFromDict(dict));
+          }));
 }
 
 void PybindFeaturePlp(py::module &m) {

kaldifeat/python/csrc/feature-spectrogram.cc

@@ -4,6 +4,7 @@
 
 #include "kaldifeat/python/csrc/feature-spectrogram.h"
 
+#include <memory>
 #include <string>
 
 #include "kaldifeat/csrc/feature-spectrogram.h"
@@ -34,9 +35,15 @@ static void PybindSpectrogramOptions(py::module &m) {
            [](const PyClass &self) -> std::string { return self.ToString(); })
       .def("as_dict",
            [](const PyClass &self) -> py::dict { return AsDict(self); })
-      .def_static("from_dict", [](py::dict dict) -> PyClass {
+      .def_static("from_dict",
+                  [](py::dict dict) -> PyClass {
                     return SpectrogramOptionsFromDict(dict);
-      });
+                  })
+      .def(py::pickle(
+          [](const PyClass &self) -> py::dict { return AsDict(self); },
+          [](py::dict dict) -> PyClass {
+            return SpectrogramOptionsFromDict(dict);
+          }));
 }
 
 static void PybindSpectrogram(py::module &m) {
@@ -46,7 +53,14 @@ static void PybindSpectrogram(py::module &m) {
       .def("dim", &PyClass::Dim)
       .def_property_readonly("options", &PyClass::GetOptions)
       .def("compute_features", &PyClass::ComputeFeatures, py::arg("wave"),
-           py::arg("vtln_warp"));
+           py::arg("vtln_warp"))
+      .def(py::pickle(
+          [](const PyClass &self) -> py::dict {
+            return AsDict(self.GetOptions());
+          },
+          [](py::dict dict) -> std::unique_ptr<PyClass> {
+            return std::make_unique<PyClass>(SpectrogramOptionsFromDict(dict));
+          }));
 }
 
 void PybindFeatureSpectrogram(py::module &m) {

kaldifeat/python/csrc/feature-window.cc

@@ -12,39 +12,39 @@
 namespace kaldifeat {
 
 static void PybindFrameExtractionOptions(py::module &m) {
-  py::class_<FrameExtractionOptions>(m, "FrameExtractionOptions")
+  using PyClass = FrameExtractionOptions;
+  py::class_<PyClass>(m, "FrameExtractionOptions")
       .def(py::init<>())
-      .def_readwrite("samp_freq", &FrameExtractionOptions::samp_freq)
+      .def_readwrite("samp_freq", &PyClass::samp_freq)
-      .def_readwrite("frame_shift_ms", &FrameExtractionOptions::frame_shift_ms)
+      .def_readwrite("frame_shift_ms", &PyClass::frame_shift_ms)
-      .def_readwrite("frame_length_ms",
+      .def_readwrite("frame_length_ms", &PyClass::frame_length_ms)
-                     &FrameExtractionOptions::frame_length_ms)
+      .def_readwrite("dither", &PyClass::dither)
-      .def_readwrite("dither", &FrameExtractionOptions::dither)
+      .def_readwrite("preemph_coeff", &PyClass::preemph_coeff)
-      .def_readwrite("preemph_coeff", &FrameExtractionOptions::preemph_coeff)
+      .def_readwrite("remove_dc_offset", &PyClass::remove_dc_offset)
-      .def_readwrite("remove_dc_offset",
+      .def_readwrite("window_type", &PyClass::window_type)
-                     &FrameExtractionOptions::remove_dc_offset)
+      .def_readwrite("round_to_power_of_two", &PyClass::round_to_power_of_two)
-      .def_readwrite("window_type", &FrameExtractionOptions::window_type)
+      .def_readwrite("blackman_coeff", &PyClass::blackman_coeff)
-      .def_readwrite("round_to_power_of_two",
+      .def_readwrite("snip_edges", &PyClass::snip_edges)
-                     &FrameExtractionOptions::round_to_power_of_two)
-      .def_readwrite("blackman_coeff", &FrameExtractionOptions::blackman_coeff)
-      .def_readwrite("snip_edges", &FrameExtractionOptions::snip_edges)
       .def("as_dict",
-           [](const FrameExtractionOptions &self) -> py::dict {
+           [](const PyClass &self) -> py::dict { return AsDict(self); })
-             return AsDict(self);
-           })
       .def_static("from_dict",
-                  [](py::dict dict) -> FrameExtractionOptions {
+                  [](py::dict dict) -> PyClass {
                     return FrameExtractionOptionsFromDict(dict);
                   })
 #if 0
       .def_readwrite("allow_downsample",
-                     &FrameExtractionOptions::allow_downsample)
+                     &PyClass::allow_downsample)
-      .def_readwrite("allow_upsample", &FrameExtractionOptions::allow_upsample)
+      .def_readwrite("allow_upsample", &PyClass::allow_upsample)
       .def_readwrite("max_feature_vectors",
-                     &FrameExtractionOptions::max_feature_vectors)
+                     &PyClass::max_feature_vectors)
 #endif
-      .def("__str__", [](const FrameExtractionOptions &self) -> std::string {
+      .def("__str__",
-        return self.ToString();
+           [](const PyClass &self) -> std::string { return self.ToString(); })
-      });
+      .def(py::pickle(
+          [](const PyClass &self) -> py::dict { return AsDict(self); },
+          [](py::dict dict) -> PyClass {
+            return FrameExtractionOptionsFromDict(dict);
+          }));
 
   m.def("num_frames", &NumFrames, py::arg("num_samples"), py::arg("opts"),
         py::arg("flush") = true);

kaldifeat/python/csrc/mel-computations.cc

@@ -26,9 +26,15 @@ static void PybindMelBanksOptions(py::module &m) {
            [](const PyClass &self) -> std::string { return self.ToString(); })
       .def("as_dict",
            [](const PyClass &self) -> py::dict { return AsDict(self); })
-      .def_static("from_dict", [](py::dict dict) -> PyClass {
+      .def_static("from_dict",
+                  [](py::dict dict) -> PyClass {
                     return MelBanksOptionsFromDict(dict);
-      });
+                  })
+      .def(py::pickle(
+          [](const PyClass &self) -> py::dict { return AsDict(self); },
+          [](py::dict dict) -> PyClass {
+            return MelBanksOptionsFromDict(dict);
+          }));
 }
 
 void PybindMelComputations(py::module &m) { PybindMelBanksOptions(m); }

kaldifeat/python/csrc/utils.h

@@ -15,7 +15,7 @@
 
 /*
  * This file contains code about `from_dict` and
- * `to_dict` for various options in kaldifeat.
+ * `as_dict` for various options in kaldifeat.
  *
  * Regarding `from_dict`, users don't need to provide
  * all the fields in the options. If some fields

kaldifeat/python/tests/test_fbank.py

@@ -2,6 +2,7 @@
 
 # Copyright      2021  Xiaomi Corporation (authors: Fangjun Kuang)
 
+import pickle
 from pathlib import Path
 
 import torch
@@ -156,6 +157,20 @@ def test_fbank_batch():
         assert torch.allclose(features[1], features1)
 
 
+def test_pickle():
+    for device in get_devices():
+        opts = kaldifeat.FbankOptions()
+        opts.use_energy = True
+        opts.use_power = False
+        opts.device = device
+
+        fbank = kaldifeat.Fbank(opts)
+        data = pickle.dumps(fbank)
+        fbank2 = pickle.loads(data)
+
+        assert str(fbank.opts) == str(fbank2.opts)
+
+
 if __name__ == "__main__":
     test_fbank_default()
     test_fbank_htk()
@@ -164,3 +179,4 @@ if __name__ == "__main__":
     test_fbank_40_bins_no_snip_edges()
     test_fbank_chunk()
     test_fbank_batch()
+    test_pickle()

kaldifeat/python/tests/test_fbank_options.py

@@ -3,6 +3,8 @@
 # Copyright (c)  2021  Xiaomi Corporation (authors: Fangjun Kuang)
 
 
+import pickle
+
 import torch
 
 import kaldifeat
@@ -176,6 +178,21 @@ def test_from_dict_full_and_as_dict():
     assert opts3.device == torch.device("cuda", 2)
 
 
+def test_pickle():
+    opts = kaldifeat.FbankOptions()
+    opts.use_energy = True
+    opts.use_power = False
+    opts.device = torch.device("cuda", 1)
+
+    opts.frame_opts.samp_freq = 44100
+    opts.mel_opts.num_bins = 100
+
+    data = pickle.dumps(opts)
+
+    opts2 = pickle.loads(data)
+    assert str(opts) == str(opts2)
+
+
 def main():
     test_default()
     test_set_get()
@@ -184,6 +201,7 @@ def main():
     test_from_empty_dict()
     test_from_dict_partial()
     test_from_dict_full_and_as_dict()
+    test_pickle()
 
 
 if __name__ == "__main__":

kaldifeat/python/tests/test_frame_extraction_options.py

@@ -2,6 +2,8 @@
 #
 # Copyright (c)  2021  Xiaomi Corporation (authors: Fangjun Kuang)
 
+import pickle
+
 import kaldifeat
 
 
@@ -94,12 +96,23 @@ def test_from_dict_full_and_as_dict():
     assert opts3.window_type == "hanning"
 
 
+def test_pickle():
+    opts = kaldifeat.FrameExtractionOptions()
+    opts.samp_freq = 44100
+    opts.dither = 5.5
+    data = pickle.dumps(opts)
+
+    opts2 = pickle.loads(data)
+    assert str(opts) == str(opts2)
+
+
 def main():
     test_default()
     test_set_get()
     test_from_empty_dict()
     test_from_dict_partial()
     test_from_dict_full_and_as_dict()
+    test_pickle()
 
 
 if __name__ == "__main__":

kaldifeat/python/tests/test_mel_bank_options.py

@@ -2,6 +2,8 @@
 #
 # Copyright (c)  2021  Xiaomi Corporation (authors: Fangjun Kuang)
 
+import pickle
+
 import kaldifeat
 
 
@@ -82,12 +84,23 @@ def test_from_dict_full_and_as_dict():
     assert opts3.htk_mode is True
 
 
+def test_pickle():
+    opts = kaldifeat.MelBanksOptions()
+    opts.num_bins = 100
+    opts.low_freq = 22
+    data = pickle.dumps(opts)
+
+    opts2 = pickle.loads(data)
+    assert str(opts) == str(opts2)
+
+
 def main():
     test_default()
     test_set_get()
     test_from_empty_dict()
     test_from_dict_partial()
     test_from_dict_full_and_as_dict()
+    test_pickle()
 
 
 if __name__ == "__main__":

kaldifeat/python/tests/test_mfcc.py

@@ -2,6 +2,7 @@
 
 # Copyright      2021  Xiaomi Corporation (authors: Fangjun Kuang)
 
+import pickle
 from pathlib import Path
 
 import torch
@@ -46,6 +47,21 @@ def test_mfcc_no_snip_edges():
         assert torch.allclose(features.cpu(), gt, rtol=1e-1)
 
 
+def test_pickle():
+    for device in get_devices():
+        opts = kaldifeat.MfccOptions()
+        opts.device = device
+        opts.frame_opts.dither = 0
+        opts.frame_opts.snip_edges = False
+
+        mfcc = kaldifeat.Mfcc(opts)
+        data = pickle.dumps(mfcc)
+        mfcc2 = pickle.loads(data)
+
+        assert str(mfcc.opts) == str(mfcc2.opts)
+
+
 if __name__ == "__main__":
     test_mfcc_default()
     test_mfcc_no_snip_edges()
+    test_pickle()

kaldifeat/python/tests/test_mfcc_options.py

@@ -3,6 +3,8 @@
 # Copyright (c)  2021  Xiaomi Corporation (authors: Fangjun Kuang)
 
 
+import pickle
+
 import torch
 
 import kaldifeat
@@ -180,6 +182,28 @@ def test_from_dict_full_and_as_dict():
     assert opts3.device == torch.device("cuda", 10)
 
 
+def test_pickle():
+    opts = kaldifeat.MfccOptions()
+    opts.num_ceps = 222
+    opts.use_energy = False
+    opts.cepstral_lifter = 21
+    opts.htk_compat = True
+    opts.device = torch.device("cuda", 3)
+
+    opts.frame_opts.samp_freq = 44100
+    opts.frame_opts.frame_length_ms = 1
+    opts.frame_opts.dither = 0.5
+
+    opts.mel_opts.num_bins = 100
+    opts.mel_opts.low_freq = 22
+    opts.mel_opts.vtln_high = -100
+
+    data = pickle.dumps(opts)
+
+    opts2 = pickle.loads(data)
+    assert str(opts) == str(opts2)
+
+
 def main():
     test_default()
     test_set_get()
@@ -188,6 +212,7 @@ def main():
     test_from_empty_dict()
     test_from_dict_partial()
     test_from_dict_full_and_as_dict()
+    test_pickle()
 
 
 if __name__ == "__main__":

kaldifeat/python/tests/test_plp.py

@@ -2,6 +2,7 @@
 
 # Copyright      2021  Xiaomi Corporation (authors: Fangjun Kuang)
 
+import pickle
 from pathlib import Path
 
 import torch
@@ -65,7 +66,22 @@ def test_plp_htk_10_ceps():
         assert torch.allclose(features.cpu(), gt, atol=1e-1)
 
 
+def test_pickle():
+    for device in get_devices():
+        opts = kaldifeat.PlpOptions()
+        opts.device = device
+        opts.frame_opts.dither = 0
+        opts.frame_opts.snip_edges = False
+
+        plp = kaldifeat.Plp(opts)
+        data = pickle.dumps(plp)
+        plp2 = pickle.loads(data)
+
+        assert str(plp.opts) == str(plp2.opts)
+
+
 if __name__ == "__main__":
     test_plp_default()
     test_plp_no_snip_edges()
     test_plp_htk_10_ceps()
+    test_pickle()

kaldifeat/python/tests/test_plp_options.py

@@ -3,6 +3,8 @@
 # Copyright (c)  2021  Xiaomi Corporation (authors: Fangjun Kuang)
 
 
+import pickle
+
 import torch
 
 import kaldifeat
@@ -191,6 +193,27 @@ def test_from_dict_full_and_as_dict():
     assert opts3.device == torch.device("cuda", 2)
 
 
+def test_pickle():
+    opts = kaldifeat.PlpOptions()
+    opts.lpc_order = 11
+    opts.num_ceps = 1
+    opts.use_energy = False
+    opts.compress_factor = 0.5
+    opts.cepstral_lifter = 2
+    opts.device = torch.device("cuda", 1)
+
+    opts.frame_opts.samp_freq = 44100
+    opts.frame_opts.snip_edges = False
+
+    opts.mel_opts.num_bins = 100
+    opts.mel_opts.high_freq = 1
+
+    data = pickle.dumps(opts)
+
+    opts2 = pickle.loads(data)
+    assert str(opts) == str(opts2)
+
+
 def main():
     test_default()
     test_set_get()
@@ -199,6 +222,7 @@ def main():
     test_from_empty_dict()
     test_from_dict_partial()
     test_from_dict_full_and_as_dict()
+    test_pickle()
 
 
 if __name__ == "__main__":

kaldifeat/python/tests/test_spectrogram.py

@@ -2,6 +2,7 @@
 
 # Copyright      2021  Xiaomi Corporation (authors: Fangjun Kuang)
 
+import pickle
 from pathlib import Path
 
 from utils import get_devices, read_ark_txt, read_wave
@@ -50,6 +51,21 @@ def test_spectrogram_no_snip_edges():
         print(features[1, 145:148], gt[1, 145:148])  # they are different
 
 
+def test_pickle():
+    for device in get_devices():
+        opts = kaldifeat.SpectrogramOptions()
+        opts.device = device
+        opts.frame_opts.dither = 0
+        opts.frame_opts.snip_edges = False
+
+        spec = kaldifeat.Spectrogram(opts)
+        data = pickle.dumps(spec)
+        spec2 = pickle.loads(data)
+
+        assert str(spec.opts) == str(spec2.opts)
+
+
 if __name__ == "__main__":
     test_spectrogram_default()
     test_spectrogram_no_snip_edges()
+    test_pickle()

kaldifeat/python/tests/test_spectrogram_options.py

@@ -3,6 +3,8 @@
 # Copyright (c)  2021  Xiaomi Corporation (authors: Fangjun Kuang)
 
 
+import pickle
+
 import torch
 
 import kaldifeat
@@ -121,6 +123,21 @@ def test_from_dict_full_and_as_dict():
     assert str(opts3) == str(opts)
 
 
+def test_pickle():
+    opts = kaldifeat.SpectrogramOptions()
+    opts.energy_floor = 1
+    opts.raw_energy = False
+    opts.device = torch.device("cuda", 1)
+
+    opts.frame_opts.samp_freq = 44100
+    opts.frame_opts.snip_edges = False
+
+    data = pickle.dumps(opts)
+
+    opts2 = pickle.loads(data)
+    assert str(opts) == str(opts2)
+
+
 def main():
     test_default()
     test_set_get()
@@ -128,6 +145,7 @@ def main():
     test_from_empty_dict()
     test_from_dict_partial()
     test_from_dict_full_and_as_dict()
+    test_pickle()
 
 
 if __name__ == "__main__":