nnet3-xvector-compute.cc File Reference

#include "base/kaldi-common.h"
#include "util/common-utils.h"
#include "nnet3/nnet-am-decodable-simple.h"
#include "base/timer.h"
#include "nnet3/nnet-utils.h"

Include dependency graph for nnet3-xvector-compute.cc:

Go to the source code of this file.

Namespaces
	kaldi
	This code computes Goodness of Pronunciation (GOP) and extracts phone-level pronunciation feature for mispronunciations detection tasks, the reference:
	kaldi::nnet3

Functions
static void	RunNnetComputation (const MatrixBase< BaseFloat > &features, const Nnet &nnet, CachingOptimizingCompiler *compiler, Vector< BaseFloat > *xvector)
int	main (int argc, char *argv[])

Function Documentation

main()

int main	(	int	argc,
		char *	argv[]
	)

Definition at line 63 of file nnet3-xvector-compute.cc.

References NnetSimpleComputationOptions::acoustic_scale, VectorBase< Real >::AddVec(), kaldi::nnet3::CollapseModel(), SequentialTableReader< Holder >::Done(), Timer::Elapsed(), ParseOptions::GetArg(), rnnlm::i, KALDI_LOG, KALDI_WARN, SequentialTableReader< Holder >::Key(), kaldi::kSetZero, SequentialTableReader< Holder >::Next(), ParseOptions::NumArgs(), NnetSimpleComputationOptions::optimize_config, Nnet::OutputDim(), ParseOptions::PrintUsage(), MatrixBase< Real >::Range(), ParseOptions::Read(), CachingOptimizingCompiler::ReadCache(), kaldi::ReadKaldiObject(), NnetSimpleComputationOptions::Register(), ParseOptions::Register(), CachingOptimizingCompilerOptions::Register(), MatrixBase< Real >::Row(), kaldi::nnet3::RunNnetComputation(), VectorBase< Real >::Scale(), kaldi::nnet3::SetBatchnormTestMode(), kaldi::nnet3::SetDropoutTestMode(), Output::Stream(), Input::Stream(), SequentialTableReader< Holder >::Value(), TableWriter< Holder >::Write(), and CachingOptimizingCompiler::WriteCache().

                                 {
  try {
    using namespace kaldi;
    using namespace kaldi::nnet3;
    typedef kaldi::int32 int32;
    typedef kaldi::int64 int64;

    const char *usage =
        "Propagate features through an xvector neural network model and write\n"
        "the output vectors.  \"Xvector\" is our term for a vector or\n"
        "embedding which is the output of a particular type of neural network\n"
        "architecture found in speaker recognition.  This architecture\n"
        "consists of several layers that operate on frames, a statistics\n"
        "pooling layer that aggregates over the frame-level representations\n"
        "and possibly additional layers that operate on segment-level\n"
        "representations.  The xvectors are generally extracted from an\n"
        "output layer after the statistics pooling layer.  By default, one\n"
        "xvector is extracted directly from the set of features for each\n"
        "utterance.  Optionally, xvectors are extracted from chunks of input\n"
        "features and averaged, to produce a single vector.\n"
        "\n"
        "Usage: nnet3-xvector-compute [options] <raw-nnet-in> "
        "<features-rspecifier> <vector-wspecifier>\n"
        "e.g.: nnet3-xvector-compute final.raw scp:feats.scp "
        "ark:nnet_prediction.ark\n"
        "See also: nnet3-compute\n";

    ParseOptions po(usage);
    Timer timer;

    NnetSimpleComputationOptions opts;
    CachingOptimizingCompilerOptions compiler_config;

    opts.acoustic_scale = 1.0; // by default do no scaling in this recipe.

    std::string use_gpu = "no";
    std::string cached_compiler_in;
    std::string cached_compiler_out;
    int32 chunk_size = -1,
      min_chunk_size = 100;
    bool pad_input = true;

    opts.Register(&po);
    compiler_config.Register(&po);

    po.Register("use-gpu", &use_gpu,
      "yes|no|optional|wait, only has effect if compiled with CUDA");
    po.Register("chunk-size", &chunk_size,
      "If set, extracts xectors from specified chunk-size, and averages.  "
      "If not set, extracts an xvector from all available features.");
    po.Register("min-chunk-size", &min_chunk_size,
      "Minimum chunk-size allowed when extracting xvectors.");
    po.Register("pad-input", &pad_input, "If true, duplicate the first and "
      "last frames of the input features as required to equal min-chunk-size.");
    po.Register("cached-compiler-in", &cached_compiler_in,
      "If set, read the cached compiler from the specified file path.");
    po.Register("cached-compiler-out", &cached_compiler_out,
      "If set, write the cached compiler to the specified file path.");

#if HAVE_CUDA==1
    CuDevice::RegisterDeviceOptions(&po);
#endif

    po.Read(argc, argv);

    if (po.NumArgs() != 3) {
      po.PrintUsage();
      exit(1);
    }

#if HAVE_CUDA==1
    CuDevice::Instantiate().SelectGpuId(use_gpu);
#endif

    std::string nnet_rxfilename = po.GetArg(1),
                feature_rspecifier = po.GetArg(2),
                vector_wspecifier = po.GetArg(3);

    Nnet nnet;
    ReadKaldiObject(nnet_rxfilename, &nnet);
    SetBatchnormTestMode(true, &nnet);
    SetDropoutTestMode(true, &nnet);
    CollapseModel(CollapseModelConfig(), &nnet);

    CachingOptimizingCompiler compiler(nnet, opts.optimize_config, compiler_config);
    
    if (!cached_compiler_in.empty()) {
        KALDI_LOG << "Reading cache from " << cached_compiler_in;
        bool cache_binary_in;
        Input ki(cached_compiler_in, &cache_binary_in);
        compiler.ReadCache(ki.Stream(), cache_binary_in);
    }

    BaseFloatVectorWriter vector_writer(vector_wspecifier);

    int32 num_success = 0, num_fail = 0;
    int64 frame_count = 0;
    int32 xvector_dim = nnet.OutputDim("output");

    SequentialBaseFloatMatrixReader feature_reader(feature_rspecifier);

    for (; !feature_reader.Done(); feature_reader.Next()) {
      std::string utt = feature_reader.Key();
      const Matrix<BaseFloat> &features (feature_reader.Value());
      if (features.NumRows() == 0) {
        KALDI_WARN << "Zero-length utterance: " << utt;
        num_fail++;
        continue;
      }
      int32 num_rows = features.NumRows(),
            feat_dim = features.NumCols(),
            this_chunk_size = chunk_size;
      if (!pad_input && num_rows < min_chunk_size) {
        KALDI_WARN << "Minimum chunk size of " << min_chunk_size
                   << " is greater than the number of rows "
                   << "in utterance: " << utt;
        num_fail++;
        continue;
      } else if (num_rows < chunk_size) {
        KALDI_LOG << "Chunk size of " << chunk_size << " is greater than "
                  << "the number of rows in utterance: " << utt
                  << ", using chunk size  of " << num_rows;
        this_chunk_size = num_rows;
      } else if (chunk_size == -1) {
        this_chunk_size = num_rows;
      }

      int32 num_chunks = ceil(
        num_rows / static_cast<BaseFloat>(this_chunk_size));
      Vector<BaseFloat> xvector_avg(xvector_dim, kSetZero);
      BaseFloat tot_weight = 0.0;

      // Iterate over the feature chunks.
      for (int32 chunk_indx = 0; chunk_indx < num_chunks; chunk_indx++) {
        // If we're nearing the end of the input, we may need to shift the
        // offset back so that we can get this_chunk_size frames of input to
        // the nnet.
        int32 offset = std::min(
          this_chunk_size, num_rows - chunk_indx * this_chunk_size);
        if (!pad_input && offset < min_chunk_size)
          continue;
        SubMatrix<BaseFloat> sub_features(
          features, chunk_indx * this_chunk_size, offset, 0, feat_dim);
        Vector<BaseFloat> xvector;
        tot_weight += offset;

        // Pad input if the offset is less than the minimum chunk size
        if (pad_input && offset < min_chunk_size) {
          Matrix<BaseFloat> padded_features(min_chunk_size, feat_dim);
          int32 left_context = (min_chunk_size - offset) / 2;
          int32 right_context = min_chunk_size - offset - left_context;
          for (int32 i = 0; i < left_context; i++) {
            padded_features.Row(i).CopyFromVec(sub_features.Row(0));
          }
          for (int32 i = 0; i < right_context; i++) {
            padded_features.Row(min_chunk_size - i - 1).CopyFromVec(sub_features.Row(offset - 1));
          }
          padded_features.Range(left_context, offset, 0, feat_dim).CopyFromMat(sub_features);
          RunNnetComputation(padded_features, nnet, &compiler, &xvector);
        } else {
          RunNnetComputation(sub_features, nnet, &compiler, &xvector);
        }
        xvector_avg.AddVec(offset, xvector);
      }
      xvector_avg.Scale(1.0 / tot_weight);
      vector_writer.Write(utt, xvector_avg);

      frame_count += features.NumRows();
      num_success++;
    }

#if HAVE_CUDA==1
    CuDevice::Instantiate().PrintProfile();
#endif
    double elapsed = timer.Elapsed();
    KALDI_LOG << "Time taken "<< elapsed
              << "s: real-time factor assuming 100 frames/sec is "
              << (elapsed*100.0/frame_count);
    KALDI_LOG << "Done " << num_success << " utterances, failed for "
              << num_fail;
    
    if (!cached_compiler_out.empty()) {
        KALDI_LOG << "Writing cache to " << cached_compiler_out;
        bool binary_write = true;
        Output ko(cached_compiler_out, &binary_write);
        compiler.WriteCache(ko.Stream(), binary_write);
    }

    if (num_success != 0) return 0;
    else return 1;
  } catch(const std::exception &e) {
    std::cerr << e.what();
    return -1;
  }
}