nnet3-offline-xvector-compute.cc

// nnet3bin/nnet3-xvector-compute.cc

// Copyright 2017   Johns Hopkins University (author: Daniel Povey)
//           2017   Johns Hopkins University (author: Daniel Garcia-Romero)
//           2017   David Snyder
//           2019   xmuspeech (author: Snowdar)

// See ../../COPYING for clarification regarding multiple authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//  http://www.apache.org/licenses/LICENSE-2.0
//
// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
// MERCHANTABLITY OR NON-INFRINGEMENT.
// See the Apache 2 License for the specific language governing permissions and
// limitations under the License.


#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"

namespace kaldi {
namespace nnet3 {

// Computes an xvector from a chunk of speech features.
static void RunNnetComputation(const MatrixBase<BaseFloat> &features,
    const Nnet &nnet, CachingOptimizingCompiler *compiler,
    Vector<BaseFloat> *xvector) {
  ComputationRequest request;
  request.need_model_derivative = false;
  request.store_component_stats = false;
  request.inputs.push_back(
    IoSpecification("input", 0, features.NumRows()));
  IoSpecification output_spec;
  output_spec.name = "output";
  output_spec.has_deriv = false;
  output_spec.indexes.resize(1);
  request.outputs.resize(1);
  request.outputs[0].Swap(&output_spec);
  std::shared_ptr<const NnetComputation> computation(std::move(compiler->Compile(request)));
  Nnet *nnet_to_update = NULL;  // we're not doing any update.
  NnetComputer computer(NnetComputeOptions(), *computation,
                  nnet, nnet_to_update);
  CuMatrix<BaseFloat> input_feats_cu(features);
  computer.AcceptInput("input", &input_feats_cu);
  computer.Run();
  CuMatrix<BaseFloat> cu_output;
  computer.GetOutputDestructive("output", &cu_output);
  xvector->Resize(cu_output.NumCols());
  xvector->CopyFromVec(cu_output.Row(0));
}

// Computes an xvector from a chunk of speech features.
static void RunNnetComputationOffline(const MatrixBase<BaseFloat> &features,
    const Nnet &nnet, NnetComputation &computation,
    Vector<BaseFloat> *xvector) {

    Nnet *nnet_to_update = NULL;  // we're not doing any update.
    NnetComputer computer(NnetComputeOptions(), computation,
        nnet, nnet_to_update);
    CuMatrix<BaseFloat> input_feats_cu(features);
    computer.AcceptInput("input", &input_feats_cu);
    computer.Run();
    CuMatrix<BaseFloat> cu_output;
    computer.GetOutputDestructive("output", &cu_output);
    xvector->Resize(cu_output.NumCols());
    xvector->CopyFromVec(cu_output.Row(0));
}

} // namespace nnet3
} // namespace kaldi

int main(int argc, char *argv[]) {
  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-offline-xvector-compute [options] <raw-nnet-in> "
        "<features-rspecifier> <vector-wspecifier>\n"
        "e.g.: nnet3-offline-xvector-compute --compiled-path-table=compiled.table 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", path_table="";
    int32 max_chunk_size = -1,
      min_chunk_size = 100;
    bool pad_input = true;
    bool binary = 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("max-chunk-size", &max_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("compiled-path-table", &path_table, "A compiled nnet table of path, which used to extract xvectors offline.");
    po.Register("binary", &binary, "If true, read the compiled net with binary mod, which should be consistent with the compiled net file.");

    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);

    BaseFloatVectorWriter vector_writer(vector_wspecifier);

    typedef unordered_map<int32, NnetComputation> HashType;
    HashType computation;
    if (!path_table.empty()) {
        Timer t;
        SequentialTokenReader table_reader("ark:" + path_table);
        KALDI_LOG << "Compiled nets are exist, so use them to extract xvectors offline.";
        for (; !table_reader.Done();table_reader.Next()) {
            std::ifstream path;
            
            int32 chunk;
            ConvertStringToInteger(table_reader.Key(), &chunk);
            
            if (chunk<min_chunk_size || chunk>max_chunk_size) {
                continue;
            }
            
            if (computation.count(chunk) == 0 && chunk > 0) {
                KALDI_LOG << "Load compiled net for " << chunk;
                std::string str = table_reader.Value();
                path.open(str.c_str());
                computation[chunk].Read(path, binary);
                path.close();
            }
        }
        KALDI_LOG << t.Elapsed() << "s used in reading compiled nets.";
    }

    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 = max_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 < max_chunk_size) {
          KALDI_LOG << "Chunk size of " << max_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 (max_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);
          if (computation.count(min_chunk_size) == 0) {
              KALDI_LOG << "There is no " << min_chunk_size << " compiled net, so compile it temporarily.";
              RunNnetComputation(padded_features, nnet, &compiler, &xvector);
          }
          else {
              RunNnetComputationOffline(padded_features, nnet, computation[min_chunk_size], &xvector);
          }
        }
        else if (computation.count(offset) == 0) {
            KALDI_LOG << "There is no " << offset << " compiled net, so compile it temporarily.";
            RunNnetComputation(sub_features, nnet, &compiler, &xvector);
        }
        else {
            RunNnetComputationOffline(sub_features, nnet, computation[offset], &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 (num_success != 0) return 0;
    else return 1;
  } catch(const std::exception &e) {
    std::cerr << e.what();
    return -1;
  }
}