// nnet3/discriminative-training.h
  
  // Copyright      2012-2015    Johns Hopkins University (author: Daniel Povey)
  // Copyright      2014-2015    Vimal Manohar
  
  
  // 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_DISCRIMINATIVE_TRAINING_H_
  #define KALDI_NNET3_DISCRIMINATIVE_TRAINING_H_
  
  #include "base/kaldi-common.h"
  #include "util/common-utils.h"
  #include "fstext/fstext-lib.h"
  #include "tree/context-dep.h"
  #include "lat/kaldi-lattice.h"
  #include "matrix/kaldi-matrix.h"
  #include "hmm/transition-model.h"
  #include "nnet3/discriminative-supervision.h"
  #include "lat/lattice-functions.h"
  #include "cudamatrix/cu-matrix-lib.h"
  
  namespace kaldi {
  namespace discriminative {
  
  /* Options for discriminative training
   *
   * Legend:
   * mmi - Maximum Mutual Information
   * mpfe - Minimum Phone Frame Error
   * smbr - State Minimum Bayes Risk
   *
   */
  struct DiscriminativeOptions {
    std::string criterion; // one of {"mmi", "mpfe", "smbr"}
                           // If the criterion does not match the supervision
                           // object, the derivatives may not be very accurate
    BaseFloat acoustic_scale; // e.g. 0.1
    bool drop_frames; // for MMI, true if we ignore frames where alignment
                      // pdf-id is not in the lattice.
    bool one_silence_class;  // Affects MPFE and SMBR objectives 
    BaseFloat boost; // for MMI, boosting factor (would be Boosted MMI)... e.g. 0.1.
    
    std::string silence_phones_str; // colon-separated list of integer ids of silence phones,
                                    // for MPFE and SMBR objectives
  
    // Cross-entropy regularization constant.  (e.g. try 0.1).  If nonzero,
    // the network is expected to have an output named 'output-xent', which
    // should have a softmax as its final nonlinearity.
    BaseFloat xent_regularize;
  
    // l2 regularization constant on the 'chain' output; the actual term added to
    // the objf will be -0.5 times this constant times the squared l2 norm.
    // (squared so it's additive across the dimensions).  e.g. try 0.0005.
    BaseFloat l2_regularize;
    
    // Options for debugging discriminative training
    
    // Accumulates gradients wrt nnet outputs
    bool accumulate_gradients;  
    
    // Accumulates nnet output
    bool accumulate_output;     
    
    // Applicable for debugging discriminative training when accumulate_gradients
    // or accumulate_output is true 
    int32 num_pdfs;             
  
    DiscriminativeOptions(): criterion("smbr"), 
                             acoustic_scale(0.1),
                             drop_frames(false),
                             one_silence_class(false),
                             boost(0.0), 
                             xent_regularize(0.0), 
                             l2_regularize(0.0),
                             accumulate_gradients(false), 
                             accumulate_output(false),
                             num_pdfs(0) { }
  
    void Register(OptionsItf *opts) {
      opts->Register("criterion", &criterion, "Criterion, 'mmi'|'mpfe'|'smbr', "
                     "determines the objective function to use.  Should match "
                     "option used when we created the examples.");
      opts->Register("acoustic-scale", &acoustic_scale, "Weighting factor to "
                     "apply to acoustic likelihoods.");
      opts->Register("drop-frames", &drop_frames, "For MMI, if true we drop frames "
                     "with no overlap of num and den pdf-ids");
      opts->Register("boost", &boost, "Boosting factor for boosted MMI (e.g. 0.1)");
      opts->Register("one-silence-class", &one_silence_class, "If true, newer "
                     "behavior which will tend to reduce insertions "
                     "when using MPFE or SMBR objective");
      opts->Register("silence-phones", &silence_phones_str,
                     "For MPFE or SMBR objectives, colon-separated list of "
                     "integer ids of silence phones, e.g. 1:2:3");
      opts->Register("l2-regularize", &l2_regularize, "l2 regularization "
                     "constant for 'chain' output "
                     "of the neural net.");
      opts->Register("xent-regularize", &xent_regularize, "Cross-entropy "
                     "regularization constant for sequence training.  If "
                     "nonzero, the network is expected to have an output "
                     "named 'output-xent', which should have a softmax as "
                     "its final nonlinearity.");
      opts->Register("accumulate-gradients", &accumulate_gradients,
                     "Accumulate gradients wrt nnet output "
                     "for debugging discriminative training");
      opts->Register("accumulate-output", &accumulate_output,
                     "Accumulate nnet output "
                     "for debugging discriminative training");
      opts->Register("num-pdfs", &num_pdfs,
                     "Number of pdfs; "
                     "applicable when accumulate-output or accumulate-gradients "
                     "is true for discriminative training");
    }
  };
  
  struct DiscriminativeObjectiveInfo {
    double tot_t;          // total number of frames
    double tot_t_weighted; // total number of frames times weight.
    double tot_objf;      // for 'mmi', the (weighted) denominator likelihood; for
                          // everything else, the objective function.
    double tot_num_count; // total count of numerator posterior 
    double tot_den_count; // total count of denominator posterior 
    double tot_num_objf;  // for 'mmi', the (weighted) numerator likelihood; for
                          // everything else 0
    
    double tot_l2_term;   // l2 regularization objective
    // l2 regularization constant on the 'chain' output; the actual term added to
    // the objf will be -0.5 times this constant times the squared l2 norm.
    // (squared so it's additive across the dimensions).  e.g. try 0.0005.
  
    // Options for debugging discriminative training
    
    // Accumulates gradients wrt nnet outputs
    bool accumulate_gradients;  
    
    // Accumulates nnet output
    bool accumulate_output;     
    
    // Applicable for debugging discriminative training when accumulate_gradients
    // or accumulate_output is true 
    int32 num_pdfs;             
  
    // Used to accumulates gradients wrt nnet outputs
    // when accumulate_gradients is true
    CuVector<double> gradients;
    // Used to accumulates output when accumulate_output is true
    CuVector<double> output;
  
    // Print statistics for the criterion
    void Print(const std::string &criterion, 
               bool print_avg_gradients = false, 
               bool print_avg_output = false) const;
  
    // Print all accumulated statistics for debugging
    void PrintAll(const std::string &criterion) const {
      Print(criterion, true, true);
    }
  
    // Print the gradient wrt nnet output accumulated for a pdf
    void PrintAvgGradientForPdf(int32 pdf_id) const;
  
    // Add stats from another object
    void Add(const DiscriminativeObjectiveInfo &other);
  
    // Returns the objective function value for the criterion
    inline double TotalObjf(const std::string &criterion) const {
      if (criterion == "mmi") return (tot_num_objf - tot_objf);
      return tot_objf;
    }
  
    // Returns true if accumulate_gradients is true 
    // and the gradients vector has been resized to store the 
    // accumulated gradients
    inline bool AccumulateGradients() const {
      return accumulate_gradients && gradients.Dim() > 0;
    }
  
    // Returns true if accumulate_output is true 
    // and the output vector has been resized to store the 
    // accumulated nnet output 
    inline bool AccumulateOutput() const {
      return accumulate_output && output.Dim() > 0;
    }
  
    // Empty constructor
    DiscriminativeObjectiveInfo();
  
    // Constructor preparing to gradients or output to be accumulated
    DiscriminativeObjectiveInfo(int32 num_pdfs);
  
    // Constructor from config options
    DiscriminativeObjectiveInfo(const DiscriminativeOptions &opts);
    
    // Reset statistics
    void Reset();
    
    void Configure(const DiscriminativeOptions &opts);
  };
  
  /**
     This function does forward-backward on the numerator and denominator 
     lattices and computes derivates wrt to the output for the specified 
     objective function.
  
     @param [in] opts        Struct containing options
     @param [in] tmodel       Transition model
     @param [in] log_priors   Vector of log-priors for pdfs
     @param [in] supervision  The supervision object, containing the numerator
                              and denominator paths. The denominator is 
                              always a lattice. The numerator is an alignment.
     @param [in] nnet_output  The output of the neural net; dimension must equal
                            ((supervision.num_sequences * supervision.frames_per_sequence) by
                              tmodel.NumPdfs()).
  
     @param [out] stats       Statistics accumulated during training such as 
                              the objective function and the total weight.
     @param [out] xent_output_deriv  If non-NULL, then the xent objective derivative
                             (which equals a posterior from the numerator forward-backward,
                             scaled by the supervision weight) is written to here.  This will
                             be used in the cross-entropy regularization code.  
  */
  void ComputeDiscriminativeObjfAndDeriv(
      const DiscriminativeOptions &opts,
      const TransitionModel &tmodel,
      const CuVectorBase<BaseFloat> &log_priors,
      const DiscriminativeSupervision &supervision,
      const CuMatrixBase<BaseFloat> &nnet_output,
      DiscriminativeObjectiveInfo *stats,
      CuMatrixBase<BaseFloat> *nnet_output_deriv,
      CuMatrixBase<BaseFloat> *xent_output_deriv);
  
  }  // namespace discriminative
  }  // namespace kaldi
  
  #endif  // KALDI_NNET3_DISCRIMINATIVE_TRAINING_H_