support loading kaldi model in python (kaldi-asr#3976)

* support load kaldi model in python

* add some component

* split one file to multi component wrap files

* fix some bugs and add test mdl

* add testmode func in batchnorm pybind

* change StatsSum StatsSumsq to Mean Var

* make const

Author: fanlu

src/nnet3/nnet-convolutional-component.h

@@ -553,6 +553,8 @@ class TdnnComponent: public UpdatableComponent {
 
   CuMatrixBase<BaseFloat> &LinearParams() { return linear_params_; }
 
+  const CuMatrix<BaseFloat> &Linearparams() const { return linear_params_; }
+
   // This allows you to resize the vector in order to add a bias where
   // there previously was none-- obviously this should be done carefully.
   CuVector<BaseFloat> &BiasParams() { return bias_params_; }

src/nnet3/nnet-normalize-component.cc

@@ -641,6 +641,24 @@ void BatchNormComponent::Write(std::ostream &os, bool binary) const {
   WriteToken(os, binary, "</BatchNormComponent>");
 }
 
+CuVector<BaseFloat> BatchNormComponent::Mean() const {
+  CuVector<BaseFloat> mean(stats_sum_);
+  if (count_ != 0) {
+    mean.Scale(1.0 / count_);
+  }
+  return mean;
+}
+
+CuVector<BaseFloat> BatchNormComponent::Var() const {
+  CuVector<BaseFloat> mean(stats_sum_), var(stats_sumsq_);
+  if (count_ != 0) {
+    mean.Scale(1.0 / count_);
+    var.Scale(1.0 / count_);
+    var.AddVecVec(-1.0, mean, mean, 1.0);
+  }
+  return var;
+}
+
 void BatchNormComponent::Scale(BaseFloat scale) {
   if (scale == 0) {
     count_ = 0.0;

src/nnet3/nnet-normalize-component.h

@@ -224,6 +224,11 @@ class BatchNormComponent: public Component {
   const CuVector<BaseFloat> &Offset() const { return offset_; }
   const CuVector<BaseFloat> &Scale() const { return scale_; }
 
+  CuVector<BaseFloat> Mean() const;
+  CuVector<BaseFloat> Var() const;
+  double Count() const { return count_; }
+  BaseFloat Eps() const { return epsilon_; }
+
  private:
 
   struct Memo {

src/nnet3/nnet-simple-component.h

@@ -971,6 +971,7 @@ class LinearComponent: public UpdatableComponent {
   BaseFloat OrthonormalConstraint() const { return orthonormal_constraint_; }
   CuMatrixBase<BaseFloat> &Params() { return params_; }
   const CuMatrixBase<BaseFloat> &Params() const { return params_; }
+  const CuMatrix<BaseFloat> &Params2() const { return params_; }
  private:
 
   // disallow assignment operator.

src/pybind/Makefile

@@ -90,7 +90,12 @@ matrix/sparse_matrix_pybind.cc \
 nnet3/nnet3_pybind.cc \
 nnet3/nnet_chain_example_pybind.cc \
 nnet3/nnet_common_pybind.cc \
+nnet3/nnet_component_itf_pybind.cc \
+nnet3/nnet_convolutional_component_pybind.cc \
 nnet3/nnet_example_pybind.cc \
+nnet3/nnet_nnet_pybind.cc \
+nnet3/nnet_normalize_component_pybind.cc \
+nnet3/nnet_simple_component_pybind.cc \
 tests/test_dlpack_subvector.cc \
 util/kaldi_holder_pybind.cc \
 util/kaldi_io_pybind.cc \

src/pybind/kaldi/io_util.py

@@ -82,3 +82,19 @@ def read_transition_model(rxfilename):
     ki.Close()
 
     return trans_model
+
+
+def read_nnet3_model(rxfilename):
+    '''Read nnet model from an rxfilename.
+    '''
+    ki = kaldi_pybind.Input()
+    is_opened, is_binary = ki.Open(rxfilename, read_header=True)
+    if not is_opened:
+        raise FileNotOpenException('Failed to open {}'.format(rxfilename))
+
+    nnet = kaldi_pybind.nnet3.Nnet()
+    nnet.Read(ki.Stream(), is_binary)
+
+    ki.Close()
+
+    return nnet

src/pybind/nnet3/Makefile

@@ -1,4 +1,5 @@
 
 test:
 	python3 ./nnet_chain_example_pybind_test.py
+	python3 ./nnet_nnet_pybind_test.py
 

src/pybind/nnet3/final.mdl

@@ -2,6 +2,7 @@
 
 // Copyright 2019   Mobvoi AI Lab, Beijing, China
 //                  (author: Fangjun Kuang, Yaguang Hu, Jian Wang)
+// Copyright 2020   JD AI, Beijing, China (author: Lu Fan)
 
 // See ../../../COPYING for clarification regarding multiple authors
 //
@@ -22,12 +23,22 @@
 
 #include "nnet3/nnet_chain_example_pybind.h"
 #include "nnet3/nnet_common_pybind.h"
+#include "nnet3/nnet_component_itf_pybind.h"
+#include "nnet3/nnet_convolutional_component_pybind.h"
 #include "nnet3/nnet_example_pybind.h"
+#include "nnet3/nnet_nnet_pybind.h"
+#include "nnet3/nnet_normalize_component_pybind.h"
+#include "nnet3/nnet_simple_component_pybind.h"
 
 void pybind_nnet3(py::module& _m) {
   py::module m = _m.def_submodule("nnet3", "nnet3 pybind for Kaldi");
 
   pybind_nnet_common(m);
+  pybind_nnet_component_itf(m);
+  pybind_nnet_convolutional_component(m);
   pybind_nnet_example(m);
   pybind_nnet_chain_example(m);
+  pybind_nnet_nnet(m);
+  pybind_nnet_normalize_component(m);
+  pybind_nnet_simple_component(m);
 }