lstm.hpp

Go to the documentation of this file.

#ifndef MLPACK_METHODS_ANN_LAYER_LSTM_HPP
#define MLPACK_METHODS_ANN_LAYER_LSTM_HPP

#include <mlpack/prereqs.hpp>
#include <limits>

#include "layer.hpp"

namespace mlpack {
namespace ann  {

template<typename MatType = arma::mat>
class LSTMType : public RecurrentLayer<MatType>
{
 public:
  LSTMType();

  LSTMType(const size_t outSize);

  LSTMType* Clone() const { return new LSTMType(*this); }

  LSTMType(const LSTMType& other);
  LSTMType(LSTMType&& other);
  LSTMType& operator=(const LSTMType& other);
  LSTMType& operator=(LSTMType&& other);

  virtual ~LSTMType() { }

  void SetWeights(typename MatType::elem_type* weightsPtr);

  void Forward(const MatType& input, MatType& output);

  void Backward(const MatType& input, const MatType& gy, MatType& g);

  /*
   * Calculate the gradient using the output delta and the input activation.
   *
   * @param input The input parameter used for calculating the gradient.
   * @param error The calculated error.
   * @param gradient The calculated gradient.
   */
  void Gradient(const MatType& input,
                const MatType& error,
                MatType& gradient);

  void ClearRecurrentState(const size_t bpttSteps, const size_t batchSize);

  const MatType& Parameters() const { return weights; }
  MatType& Parameters() { return weights; }

  size_t WeightSize() const
  {
    return (4 * outSize * inSize + 7 * outSize + 4 * outSize * outSize);
  }

  void ComputeOutputDimensions()
  {
    inSize = std::accumulate(this->inputDimensions.begin(),
        this->inputDimensions.end(), 0);
    this->outputDimensions = std::vector<size_t>(this->inputDimensions.size(),
        1);

    // The LSTM layer flattens its input.
    this->outputDimensions[0] = outSize;
  }

  template<typename Archive>
  void serialize(Archive& ar, const uint32_t /* version */);

 private:
  size_t inSize;

  size_t outSize;

  MatType weights;

  MatType output2GateInputWeight;

  MatType input2GateInputWeight;

  MatType input2GateInputBias;

  MatType cell2GateInputWeight;

  MatType output2GateForgetWeight;

  MatType input2GateForgetWeight;

  MatType input2GateForgetBias;

  MatType cell2GateForgetWeight;

  MatType output2GateOutputWeight;

  MatType input2GateOutputWeight;

  MatType input2GateOutputBias;

  MatType cell2GateOutputWeight;

  // Below here are recurrent state matrices.

  MatType inputGate;

  MatType forgetGate;

  MatType hiddenLayer;

  MatType outputGate;

  MatType input2HiddenWeight;

  MatType input2HiddenBias;

  MatType output2HiddenWeight;

  arma::Cube<typename MatType::elem_type> cell;

  // These members store recurrent state.

  arma::Cube<typename MatType::elem_type> inputGateActivation;

  arma::Cube<typename MatType::elem_type> forgetGateActivation;

  arma::Cube<typename MatType::elem_type> outputGateActivation;

  arma::Cube<typename MatType::elem_type> hiddenLayerActivation;

  arma::Cube<typename MatType::elem_type> cellActivation;

  MatType forgetGateError;

  MatType outputGateError;

  arma::Cube<typename MatType::elem_type> outParameter;

  MatType inputCellError;

  MatType inputGateError;

  MatType hiddenError;
}; // class LSTMType

// Convenience typedefs.

// Standard LSTM layer.
typedef LSTMType<arma::mat> LSTM;

} // namespace ann
} // namespace mlpack

// Include implementation.
#include "lstm_impl.hpp"

#endif