nnet-optimize-utils.h

Go to the documentation of this file.

// nnet3/nnet-optimize-utils.h

// Copyright 2015    Johns Hopkins University (author: Daniel Povey)

// 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.

#ifndef KALDI_NNET3_NNET_OPTIMIZE_UTILS_H_
#define KALDI_NNET3_NNET_OPTIMIZE_UTILS_H_

#include <mutex>
#include <list>
#include "nnet3/nnet-compile.h"
#include "nnet3/nnet-analyze.h"


namespace kaldi {
namespace nnet3 {


struct NnetOptimizeOptions;  // Forward declaration.

class VariableMergingOptimizer {
 public:
  VariableMergingOptimizer(const NnetOptimizeOptions &config,
                           const Nnet &nnet,
                           NnetComputation *computation);
  // Note: you can call this only once.  If it returns true, it means it has
  // merged variables.  In this case, you have the option to instantiate another
  // copy of the class and try again with that other copy.
  bool MergeVariables();

 private:
  std::pair<bool,bool> MayBeMerged(int32 command, int32 s1, int32 s2) const;

  // Merges to matrices, whether left merge or right merge.  s_to_keep and
  // s_to_discard are the submatrix-indexes we will keep and discard
  // respectively (these are s1 and s2 in some order.
  void DoMerge(int32 command_index, int32 s_to_keep, int32 m_to_discard);

  void MarkAsDirty(int32 s);

  void Initialize();

  const NnetOptimizeOptions &config_;
  const Nnet &nnet_;
  NnetComputation *computation_;

  Analyzer analyzer_;

  // lists of submatrices that correspond to each matrix.
  std::vector<std::vector<int32> > matrix_to_submatrix_;

  // for each variable (as defined by analyzer_.variables), true if
  // we have already performed a merge on it.
  std::vector<bool> variable_dirty_;

  bool already_called_merge_variables_;
};

void ExtendMatrices(NnetComputation *computation);


void ConsolidateModelUpdate(const Nnet &nnet,
                            NnetComputation *computation);




// Class DerivativeTimeLimiter is used inside LimitDerivativeTimes().
// Its function is to modify the computation so that we don't work
// with derivatives outside of a specified range of t values; this is
// useful, for instance, in BLSTMs where you might have a fair amount of
// left and right context in the training examples but don't want to
// propagate the derivatives to there.
//
// We require that the computation have debug info set up
// (!matrix_debug_info.empty()) and that this be the first
// optimization you perform.  This means that the debug_info will
// be accurate and that all matrices will be initialized with
// zero contents.
class DerivativeTimeLimiter {
 public:
  DerivativeTimeLimiter(const Nnet &nnet,
                        int32 min_deriv_time,
                        int32 max_deriv_time,
                        NnetComputation *computation);

  void LimitDerivTimes();

 private:

  // sets up matrix_prune_info_.
  void ComputeMatrixPruneInfo();

  // sets up subatrix_map_ and submatrix_map_if_deriv_.
  void ComputeSubmatrixMaps();

  // modifies all the commands as appropriate to reflect that some derivative
  // values are zero (i.e. save any computation we can, based on this
  // assumption).
  void ModifyCommands();

  // this function, called after we've modified the commands to operate on
  // submatrices of the original matrices, works out for which of the matrices
  // we can actually limit their extent in time, and changes the way the
  // matrices are allocated (it may remove some matrices entirely).
  void PruneMatrices();

  // this function modifies commands of type kPropagate to set the memo indexes
  // to zero if the memo indexes appear in the list memos_to_delete_.  It's
  // because if a backprop command has been deleted, the propagate command
  // should no longer store a memo.
  void RemoveUnusedMemos();


  // called from PruneMatrices only for matrices that are derivatives,
  // not inputs or outputs of the computation, and which are partly
  // inside the time range, this function returns true if we can
  // limit the size of the matrix (because variables outside the
  // desired range are never accessed), and false otherwise.
  inline bool CanLimitMatrix(const Analyzer &analyzer,
                             int32 matrix_index) const;

  // called from PruneMatrices after it has figured out which matrices we need
  // to limit to a row-range, this function changes computation->submatrices and
  // computation->matrices in the way required to do that.
  inline void LimitMatrices(const std::vector<bool> &will_limit);

  // does the processing for a command of type kMatrixCopy or kMatrixAdd.
  void MapSimpleMatrixCommand(NnetComputation::Command *c);

  // does the processing for a command of type kCopyRows or kAddRows, where
  // 1st and 2nd args are submatrix indexes and the 3rd arg is a vector of
  // row-indexes.
  void MapIndexesCommand(NnetComputation::Command *c);

  // does the processing for a command of type kAddRowsMulti, kAddToRowsMulti,
  // kCopyRowsMulti or kCopyToRowsMulti, 1st arg is submatrix index that the
  // command is called with, and 2nd arg is 'indexes_multi' index (which
  // contains pairs (source-submatrix, source-row).
  void MapIndexesMultiCommand(NnetComputation::Command *c);

  // does the processing for a command of type kAddRowRanges.
  void MapAddRowRangesCommand(NnetComputation::Command *c);

  // Modifies this command to take into account prune_info_.  At this point we
  // don't actually reduce the size of the matrices, we simply make the commands
  // operate on submatrices of the original matrices where possible- or
  // delete them completely if their output is all zeros or for other reasons
  // we detect that they would be no-ops.
  // Note: this calls computation_->NewSubMatrix, and will generate duplicates
  // of the same submatrix which we'll later remove in RemoveOrphanMatrices.
  void ModifyCommand(NnetComputation::Command *command);

  // this will detect which matrices we can reduce the allocated size of,
  // and reduce their size.
  void ResizeMatrices();

  // Requires that we have mapped 'initial_submatrix' to 'new_submatrix' in
  // an operation that may have removed some data on the left and/or the
  // right (but still they point to the same underlying matrix).  Outputs
  // to 'left_prune' and 'right_prune' the number of rows we have
  // removed on the left and on the right respectively.
  inline void GetPruneValues(int32 initial_submatrix,
                             int32 new_submatrix,
                             int32 *left_prune,
                             int32 *right_prune) const;

  // This helper function, used while mapping commands, returns true if the
  // Cindex represented by the pair (submatrix, row_index) has a 't' value
  // within the range [min_deriv_time_, max_deriv_time_].
  bool RowIsKept(int32 submatrix,
                 int32 row_index) const;


  struct MatrixPruneInfo {
    bool is_deriv;  // true if the matrix represents a derivative (copied from
                    // the debug-info; repeated here for convenience).
    bool fully_inside_range;  // True if the matrix is completely inside the time range
                             // specified.
    bool partly_inside_range;  // true if the matrix is partly (but not fully)
                               // inside the time range specified.
    int32 row_begin;  // if partly_inside_range, the first row that's within the time range (i.e. for which
                      // min_deriv_time_ <= t < max_deriv_time_.
    int32 row_end;    // if partly_inside_range, one plus the last row that's within
                      // the time range.
  };


  const Nnet &nnet_;

  int32 min_deriv_time_;
  int32 max_deriv_time_;

  // the computation; we require it to have debug info set up
  // (otherwise you shouldn't be instantiating this class).
  NnetComputation *computation_;

  // for each matrix index > 0, the index of a submatrix that consists of
  // the entirety of that matrix.
  std::vector<int32> whole_submatrices_;

  std::vector<MatrixPruneInfo> matrix_prune_info_;

  // for each submatrix in the original range of computation_->submatrices,
  // submatrix_map_ maps it to itself if the submatrix is completely inside the
  // time-range, or to zero if it's completely outside the time-range, or to a
  // newly created submatrix-index if it's partly inside the time-range.
  std::vector<int32> submatrix_map_;

  // submatrix_map_if_deriv_ contains the quantity:
  // IsDerivative(s) ? submatrix_map_[s] : s,
  // where IsDerivative(s) is true if s is part of a matrix that (according to its
  // debug info) represents a derivative.
  // this comes up so frequently that storing it separately seemed like a good idea.
  std::vector<int32> submatrix_map_if_deriv_;

  std::vector<MatrixPruneInfo> prune_info_;

  // List of indexes of memos that will no longer be stored because the backprop
  // commands using them were deleted.
  std::unordered_set<int32> memos_to_delete_;
};


// This utility function, used in code that calls LimitDerivativeTimes(), returns
// the largest time 't' in any of the 'outputs' in the computation request,
// or crashes if there are no outputs (or no cindexes in those outputs).
int32 MaxOutputTimeInRequest(const ComputationRequest &request);

// This is the top-level interface to limit the times on which derivatives are
// computed (e.g. for truncated BPTT); internally it uses class
// DerivativeLimiter.  Will do nothing if min_deriv_time and max_deriv_time are
// their default -inf,+inf values.
void LimitDerivativeTimes(const Nnet &nnet,
                          int32 min_deriv_time,
                          int32 max_deriv_time,
                          NnetComputation *computation);

bool RequestIsDecomposable(const ComputationRequest &request,
                           ComputationRequest *mini_request,
                           int32 *num_n_values);


void ExpandComputation(const Nnet &nnet,
                       const MiscComputationInfo &misc_info,
                       const NnetComputation &computation,
                       bool need_debug_info,
                       int32 num_n_values,
                       NnetComputation *expanded_computation);



bool ReplaceRowWithMatrixOps(NnetComputation *computation);

bool SnipRowOps(NnetComputation *computation);


bool SplitRowOps(NnetComputation *computation);

void RenumberComputation(NnetComputation *computation);


void RemoveNoOps(NnetComputation *computation);

void IdentifySubmatrixArgs(NnetComputation::Command *command,
                           std::vector<int32*> *submatrix_args);

bool MatrixIsUnused(const Analyzer &analyzer,
                    const NnetComputation &computation,
                    int32 m);

void RemoveCommandsForUnusedMatrix(const Analyzer &analyzer,
                                   int32 m,
                                   NnetComputation *computation);


void IdentifySubmatrixArgs(std::vector<NnetComputation::Command> *commands,
                           std::vector<int32*> *submatrix_args);

void IdentifySubmatrixArgsInComputation(NnetComputation *computation,
                                        std::vector<int32*> *submatrix_args);


void IdentifyIndexesMultiArgs(std::vector<NnetComputation::Command> *commands,
                              std::vector<int32*> *indexes_multi_args);

void IdentifyIndexesArgs(std::vector<NnetComputation::Command> *commands,
                         std::vector<int32*> *indexes_args);

void IdentifyIndexesArgs(std::vector<NnetComputation::Command> *commands,
                         std::vector<int32*> *indexes_args);

void IdentifyIndexesRangesArgs(std::vector<NnetComputation::Command> *commands,
                               std::vector<int32*> *indexes_ranges_args);

void InsertCommands(
    std::vector<std::pair<int32, NnetComputation::Command> > *commands,
    NnetComputation *computation);

void OptimizeMemoryCompression(const Nnet &nnet,
                               int32 memory_compression_level,
                               NnetComputation *computation);


void OptimizeLoopedComputation(const Nnet &nnet,
                               NnetComputation *computation);


void FixGotoLabel(NnetComputation *computation);


class ComputationCache {
 public:
  ComputationCache(int32 cache_capacity);

  // Note: if something fails in Read(), or the written cache was from an older
  // format, it will just leave the cache empty.
  void Read(std::istream &is, bool binary);

  void Write(std::ostream &os, bool binary) const;


  // Searches for the computation corresponding to this computation, and returns
  // it if cached, or NULL (as std::shared_ptr) if not.  (We need shared_ptr to
  // handle multi-threaded operation, so that if the computation is ejected from
  // the cache by another thread, it won't be deleted while still in use).  This
  // function also moves this computation to the end of the
  // most-recently-accessed queue, which is why it's not const.
  std::shared_ptr<const NnetComputation> Find(const ComputationRequest &request);


  // Inserts the computation into the cache-- this is assumed to be the
  // computation for the computation-request 'request'.  Returns a shared_ptr
  // which can be used to access the object.  This function takes ownership of
  // 'computation'.
  std::shared_ptr<const NnetComputation> Insert(const ComputationRequest &request,
                                                const NnetComputation *computation);

  ~ComputationCache();

  // Checks the stored computation for correctness.
  void Check(const Nnet &nnet) const;
 private:

  std::mutex mutex_;  // Read/write mutex.

  int32 cache_capacity_;

  // The access queue for keeping track of the freshness of computation.
  // Most-recently-accessed computation is at the end, and
  // least-recently-accessed computaiton is at the beginning.  Together with
  // computation_cache_, this forms a most-recently-used (MRU) cache for
  // Computations, indexed by ComputationRequest. The pointers are owned in
  // computation_cache_.
  typedef std::list<const ComputationRequest*> AqType;
  AqType access_queue_;

  // Map from computation-request to pair of (computation, and position in
  // access_queue_). Used for fast lookup of previously compiled computations.
  // All pointers are owned here.
  typedef unordered_map<const ComputationRequest*,
                        std::pair<std::shared_ptr<const NnetComputation>, AqType::iterator>,
                        ComputationRequestHasher,
                        ComputationRequestPtrEqual> CacheType;
  CacheType computation_cache_;
};




} // namespace nnet3
} // namespace kaldi


#endif