// Copyright (c) 2017 Graphcore Ltd. All rights reserved. /** \file * * Support for Long short-term memory cells. * */ #ifndef popnn_Lstm_hpp #define popnn_Lstm_hpp #include #include #include #include #include namespace popnn { namespace lstm { /** * Get the default order of the gates in a basic LSTM cell. * The default order is: * [Forget gate, Input gate, Candidate, Output Gate]. */ const std::vector getDefaultBasicLstmCellOrder(); /** Structure representing the parameters of the LSTM. */ struct LstmParams { rnn::RnnParams rnn; /// The datatype of the LSTM. /// \deprecated Use rnn::RnnParams.dataType instead. poplar::Type dataType; /// The batch size. /// \deprecated Use rnn::RnnParams.batchSize instead. std::size_t batchSize; /// The number of time steps in the sequence of the LSTM. /// \deprecated Use rnn::RnnParams.maxTimeSteps instead. std::size_t timeSteps; /// The number of neurons before and after each layer of the LSTM. /// If the LSTM consists of N layers, then this should be a vector /// of size N+1. The first element is the input size and each subsequent /// element is the output size of the LSTM layer. /// \deprecated Use rnn::RnnParams.layerSizes instead. std::vector layerSizes; /// If true the Lstm function returns the entire sequence of outputs, /// otherwise it returns just the final output. bool outputFullSequence = true; /// If this parameter is set to false then the LSTM will skip the /// calculation of weighted inputs (only useful for benchmarking). bool doInputWeightCalc = true; /// If this parameter is set to false then the LSTM will skip the /// calculation of the gradients of the inputs. bool calcInputGradients = true; /// If this parameter is set to true then the LSTM will preserve the internal /// state at the last valid timestep for each batch entry when using /// variable timesteps. This may come with a speed penalty. /// See RnnParams for information about variable timesteps. /// Use of this parameter does not affect the retained intermediate results in /// the case of training. bool preserveFinalState = false; /// The weights and biases for all of the layers being processed are /// concatenated in the outermost dimension of the weights and biases tensors. /// This option allows you to specify the order of the gates in that outermost /// dimension. The default order can be obtained with /// getDefaultBasicLstmCellOrder(). std::vector cellOrder = getDefaultBasicLstmCellOrder(); /// Activation function. NonLinearityType activation = NonLinearityType::TANH; /// Recurrent activation function. NonLinearityType recurrentActivation = NonLinearityType::SIGMOID; LstmParams(poplar::Type dataType, std::size_t batchSize, std::size_t timeSteps, std::vector layerSizes, NonLinearityType activation = NonLinearityType::TANH, NonLinearityType recurrentActivation = NonLinearityType::SIGMOID); LstmParams(poplar::Type dataType, std::size_t batchSize, std::size_t maxTimeSteps, const poplar::Tensor &timeSteps, std::vector layerSizes, NonLinearityType activation = NonLinearityType::TANH, NonLinearityType recurrentActivation = NonLinearityType::SIGMOID); }; /** * Structure holding the state of a LSTM cell, or the gradients for the state * (depending on the context). */ struct LstmState { poplar::Tensor output; poplar::Tensor cellState; poplar::Tensor getAsTensor() const; }; /** * Predict what matrix multiplications will be needed for the given parameters * and return a list of corresponding matmul parameters and options. */ std::vector> getMatMulPrePlanParameters(LstmParams params, poplar::OptionFlags opts); uint64_t getBasicLstmCellFwdFlops(const LstmParams ¶ms); uint64_t getBasicLstmCellBwdFlops(const LstmParams ¶ms); uint64_t getBasicLstmCellWuFlops(const LstmParams ¶ms); /** Create an input tensor of shape [\p numSteps, \p batchSize, \p inputSize] * that is optimally mapped to multiply the whole input sequence in a single * matrix multiply operation. * * **LSTM options** * * * `availableMemoryProportion` Decimal between 0 and 1 (inclusive) * * See poplin::createWeights(). * * * `inferenceOnly` (true, false) [=false] * * Sets convolution pass to \c INFERENCE_FWD if true; \c TRAINING_FWD * otherwise. See the \c pass option in poplin::createWeights(). * * * `partialsType` (half, float) [=float] * * See poplin::createWeights(). * * * `weightAccumulatorsType` (half, float) [=data type of lstm] * * Data type of the weight accumulators for the LSTM's weight matrices * and biases * * * `preCalcWeights` (true, false) [=false] * * If true, use one big matrix multiply before the recurrent calculation to * perform the part of the calculation that only depends on the input * sequence. * * * `recomputationMode` (none, cellAndTanh, full) [=none] * * * **none**: No recomputation is done in the backwards pass. * * * **cellAndTanh**: A small amount of recomputation is done in the * backwards pass, yielding some reduction in memory footprint for the * layer. * * * **full**: Recompute everything from the forward pass. Saves the most * memory at the cost of an extra forward pass of cycles. * * \param graph Graph to which the LSTM cell belongs. * \param params The LSTM parameters. * \param debugContext Debug information. * \param options Any implementation/debug options for the LSTM. * \param planningCache A poplin matrix multiply planning cache. * * \return A tensor created in the graph, of shape * [\p timeSteps, \p batchSize, \p inputSize]. */ poplar::Tensor createInput(poplar::Graph &graph, const LstmParams ¶ms, const poplar::DebugContext &debugContext, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); /** Create the initial output that can be combined with the initial cell state * using an LstmState. This then can be fed into the LSTM cell at the first * timestep. * * \param graph Graph to which the LSTM cell belongs. * \param params The LSTM parameters. * \param debugContext Debug information. * \param options Any implementation/debug options for the LSTM. See * createInput(). * \param planningCache A poplin matrix multiply planning cache. * * \return A tensor which is the cell state for the forward operation of the * LSTM cell. */ poplar::Tensor createInitialOutput(poplar::Graph &graph, const LstmParams ¶ms, const poplar::DebugContext &debugContext, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); /** Create the initial cell state that can be combined with the initial output * using an LstmState. This then can be fed into the LSTM cell at the first * timestep. * * \param graph Graph to which the LSTM cell belongs. * \param params The LSTM parameters. * \param debugContext Debug information. * \param options Any implementation/debug options for the LSTM. See * createInput(). * \param planningCache A poplin matrix multiply planning cache. * * \return A tensor which is the cell state for the forward operation of the * LSTM cell. */ poplar::Tensor createInitialCellState(poplar::Graph &graph, const LstmParams ¶ms, const poplar::DebugContext &debugContext, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); /** Creates the initial state (both output and cell state) that is fed into * the LSTM cell at the first timestep. It can be initialised by writing the * appropriate member or using zeroInitialState(). * * \param graph Graph to which the LSTM cell belongs. * \param params The LSTM parameters. * \param debugContext Debug information. * \param options Any implementation/debug options for the LSTM. See * createInput(). * \param planningCache A poplin matrix multiply planning cache. * * \return A tensor which is the state for the forward operation of the LSTM * cell. */ LstmState createInitialState(poplar::Graph &graph, const LstmParams ¶ms, const poplar::DebugContext &debugContext, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); /** Initialise the forward state of an LSTM with zeros. * * \param graph Graph to which the LSTM cell belongs. * \param initialState The initial state to zero. * \param prog The program to extend with the initialization * code. * \param debugContext Optional debug information. */ void zeroInitialState(poplar::Graph &graph, const LstmState &initialState, poplar::program::Sequence &prog, const poplar::DebugContext &debugContext = {}); /** * Structure holding all the parameters of an LSTM cell, or the * deltas for those parameters (depending on the context). */ struct LstmWeights { poplar::Tensor inputWeights; poplar::Tensor outputWeights; poplar::Tensor biases; }; /** Create the weights kernel used to weight the input of an LSTM. * Returns the \c inputWeights and \c outputWeights. */ std::pair createWeightsKernel(poplar::Graph &graph, const LstmParams ¶ms, const poplar::DebugContext &debugContext, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); /** Create the weights biases. */ poplar::Tensor createWeightsBiases(poplar::Graph &graph, const LstmParams ¶ms, const poplar::DebugContext &debugContext, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); /** Create the weights (both kernel and biases) used to weight the input of an * LSTM. */ LstmWeights createWeights(poplar::Graph &graph, const LstmParams ¶ms, const poplar::DebugContext &debugContext, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); /** Calculate the result of applying an LSTM across a sequence. * * The LSTM is run for rnn::RnnParams.maxTimeSteps, each with a batch of size \p * batchSize, an input size of \p inputSize and output size of \p outputSize. * The total number of units within each LSTM cell is \c lstmUnits = * \c BASIC_LSTM_CELL_NUM_UNITS. * * \param graph Graph to which the LSTM cell belongs. * \param params The parameters of the LSTM. * \param stateInit Initial state for the LSTM. * \param in The input tensor to the LSTM, of shape * [\p timeSteps, \p batchSize, \p inputSize]. * \param weights The LSTM weights structure. * \param[out] intermediates Intermediate results that are retained in the * the forward pass of training for use in the * backward pass. This argument should be set to * null if we are only doing inference. * \param fwdProg Program sequence. * \param debugContext Optional debug information. * \param options LSTM implementation options. * See createInput(). * \param planningCache The matmul planning cache. * * \return The output of the LSTM and the final cell state. * Depending on the `outputFullSequence` parameter, the output tensor is * either the output of the last timestep in the shape * [\p batchSize, \p outputSize] or it is the sequence of outputs for * every timestep in the shape [\p timeSteps, \p batchSize, * \p outputSize]. */ std::pair lstmFwd(poplar::Graph &graph, const LstmParams ¶ms, const LstmState &stateInit, const poplar::Tensor &in, const LstmWeights &weights, poplar::Tensor *intermediates, poplar::program::Sequence &fwdProg, const poplar::DebugContext &debugContext = {}, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); /** * Run LSTM backward pass. The backward pass executes in reverse order as * compared to the forward pass. If the forward steps for a LSTM layer are sf = * {0, 1, 2, ..., S - 1} then the backward steps run for sb = {S - 1, S - 2, * .... , 1, 0}. * * Note 1: If the time step limit is variable, the entries above the given time * step limit must be explicitly set to zero in `fwdIntermediates`, in * order for the weights to be correctly updated. * Note 2: if the time limit is variable, the initialising state gradients at * the last time step are internally zero initialised and the * `lastStepStateGrad` parameter is ignored. * * \param graph Graph to which the LSTM cell belongs. * \param params The parameters of the LSTM. * \param prog Program sequence. * \param fwdStateInit Forward state tensor for initial step. * \param fwdIntermediates Intermediates results from the forward pass. * \param weights The LSTM weights structure. * \param input The input tensor to the LSTM of shape: * [timesteps, batch, inputSize]. * \param output The output tensor from the forward pass. Depending * on the outputFullSequence parameter this is either * the output for the last timestep or it is a * sequence of outputs for each timestep. * \param outputGrad The gradients of the output. Depending on the * \p outputFullSequence parameter this is either the * gradient of the output for the last timestep or * it is a sequence output gradients for each * timestep. * \param *lastStepStateGrad The gradient of the state at the last step. May be * null if state gradient is to be zero initialised. * \param[out] *inputSeqGrad The gradients of the inputs - may be null if * this information is not required. * \param[out] *bwdIntermediates Intermediates gradients that are retained in * the backward pass of training for use in the * weight update. This argument should be set to null * if you do not need to calculate weight deltas. * \param debugContext Optional debug information. * \param options LSTM implementation options. See createInput(). * \param planningCache The matmul planning cache. * * \return The gradient of the initial state. */ LstmState lstmBwd(poplar::Graph &graph, const LstmParams ¶ms, poplar::program::Sequence &prog, const LstmState &fwdStateInit, const poplar::Tensor &fwdIntermediates, const LstmWeights &weights, const poplar::Tensor &input, const poplar::Tensor &output, const poplar::Tensor &outputGrad, const LstmState *lastStepStateGrad, poplar::Tensor *inputGrad, poplar::Tensor *bwdIntermediates, const poplar::DebugContext &debugContext = {}, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); /** \deprecated * **deprecated** Use previously defined popnn::lstmBwd() instead. * * Run LSTM backward pass. The backward pass executes in reverse order as * compared to the forward pass. If the forward steps for a LSTM layer are sf = * {0, 1, 2, ..., S - 1} then the backward steps run for sb = {S - 1, S - 2, * .... , 1, 0}. * * Note 1: If the time step limit is variable, the entries above the given time * step limit must be explicitly set to zero in `fwdIntermediates`, in * order for the weights to be correctly updated. * Note 2: if the time limit is variable, the initialising cell state gradient * at the last time step is internally zero initialised and the * `lastCellStateGrad` parameter is ignored. * * \param graph Graph to which the LSTM cell belongs. * \param params The parameters of the LSTM. * \param prog Program sequence. * \param fwdStateInit Forward state tensor for initial step. * \param fwdIntermediates Intermediates results from the forward pass. * \param weights The LSTM weights structure. * \param input The input tensor to the LSTM, of shape * [\p timeSteps, \p batchSize, \p inputSize]. * \param output The output tensor from the forward pass. Depending * on the \p outputFullSequence parameter, this is * either the output for the last timestep or it is a * sequence of outputs for each timestep. * \param outputGrad The gradients of the output. Depending on the * \p outputFullSequence parameter, this is either the * gradient of the output for the last timestep or * it is a sequence output gradients for each * timestep. * \param lastCellStateGrad The gradient of the last cell state - may be * null if there is no incoming gradient. * \param[out] *inputSeqGrad The gradients of the inputs - may be null if * this information is not required. * \param[out] *bwdIntermediates Intermediates gradients that are retained in * the backward pass of training for use in the * weight update. This argument should be set to null * if you do not need to calculate weight deltas. * \param debugContext Optional debug information. * \param options LSTM implementation options. * See createInput(). * \param planningCache The matmul planning cache. * * \return The gradient of the initial state. */ LstmState lstmBwd(poplar::Graph &graph, const LstmParams ¶ms, poplar::program::Sequence &prog, const LstmState &fwdStateInit, const poplar::Tensor &fwdIntermediates, const LstmWeights &weights, const poplar::Tensor &input, const poplar::Tensor &output, const poplar::Tensor &outputGrad, const poplar::Tensor *lastCellStateGrad, poplar::Tensor *inputGrad, poplar::Tensor *bwdIntermediates, const poplar::DebugContext &debugContext = {}, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); /** * Run a standalone weight update pass. Takes intermediates and gradients * from the backward pass and calculates and returns weight deltas. * * Note: If the timestep limit is variable, the entries above the given time * step limit must be explicitly set to zero in \p fwdIntermediates, in * order for the weights to be correctly updated. * * \param graph Graph to which the LSTM cell belongs. * \param params The parameters of the LSTM. * \param prog Program sequence to add operations to. * \param fwdStateInit Forward state tensor for initial step. * \param fwdIntermediates Intermediate results from the forward pass. * \param bwdIntermediates Intermediate results from the backward pass. * \param weights The LSTM weights structure. * \param input The input tensor to the LSTM, of shape * [\p timeSteps, \p batchSize, \p inputSize]. * \param output The output tensor from the forward pass. Depending * on the \p outputFullSequence parameter, this is * either the output for the last timestep or it is a * sequence of outputs for each timestep. * \param debugContext Optional debug information. * \param options LSTM implementation options. * See createInput(). * \param planningCache The matmul planning cache. * * \return A set of weight gradients to sum with weights. */ LstmWeights lstmWU(poplar::Graph &graph, const LstmParams ¶ms, poplar::program::Sequence &prog, const LstmState &fwdStateInit, const poplar::Tensor &fwdIntermediates, const poplar::Tensor &bwdIntermediates, const LstmWeights &weights, const poplar::Tensor &input, const poplar::Tensor &output, const poplar::DebugContext &debugContext = {}, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); /** * Run a combined LSTM backward and weight update pass. Use this combined * backward and weight update pass in preference to lstmBwd() and lstmWU() * separately, in order to allow the most efficient implementation to be * chosen if you do not need to split the operation. * * Note 1: If the time step limit is variable, the entries above the given time * step limit must be explicitly set to zero in `fwdIntermediates`, in * order for the weights to be correctly updated. * Note 2: if the time limit is variable, the initialising state gradients at * the last time step are internally zero initialised and the * `lastStepStateGrad` parameter is ignored. * * \param graph Graph to which the LSTM cell belongs. * \param params The parameters of the LSTM. * \param prog Program sequence. * \param fwdStateInit Forward state tensor for initial step. * \param fwdIntermediates Intermediates results from the forward pass. * \param weights The LSTM weights structure. * \param input The input tensor to the LSTM of shape: * [timesteps, batch, inputSize]. * \param output The output tensor from the forward pass. Depending * on the outputFullSequence parameter this is either * the output for the last timestep or it is a * sequence of outputs for each timestep. * \param outputGrad The gradients of the output. Depending on the * \p outputFullSequence parameter this is either the * gradient of the output for the last timestep or * it is a sequence output gradients for each * timestep. * \param *lastStepStateGrad The gradient of the state at the last step. May be * null if state gradient is to be zero initialised. * \param[out] *inputSeqGrad The gradients of the inputs. May be null if * this information is not required. * \param weightsGrad A set of weight deltas to sum with weights. * \param debugContext Optional debug information. * \param options LSTM implementation options. See createInput(). * \param planningCache The matmul planning cache. * * \return The gradient of the initial state. */ LstmState lstmBwdWithWU(poplar::Graph &graph, const LstmParams ¶ms, poplar::program::Sequence &prog, const LstmState &fwdStateInit, const poplar::Tensor &fwdIntermediates, const LstmWeights &weights, const poplar::Tensor &input, const poplar::Tensor &output, const poplar::Tensor &outputGrad, const LstmState *lastStepStateGrad, poplar::Tensor *inputGrad, LstmWeights &weightsGrad, const poplar::DebugContext &debugContext = {}, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); /** \deprecated * **deprecated** Use previously defined popnn::lstmBwdWithWU() instead. * * Run a combined LSTM backward and weight update pass. Use this combined * backward and weight update pass in preference to `lstmBwd` and `lstmWU` * separately in order to allow the most efficient implementation to be * chosen if you do not need to split the operation. * * Note 1: If the time step limit is variable, the entries above the given time * step limit must be explicitly set to zero in `fwdIntermediates`, in * order for the weights to be correctly updated. * Note 2: if the time limit is variable, the initialising cell state gradient * at the last time step is internally zero initialised and the * `lastCellStateGrad` parameter is ignored. * * \param graph Graph to which the LSTM cell belongs. * \param params The parameters of the LSTM. * \param prog Program sequence. * \param fwdStateInit Forward state tensor for initial step. * \param fwdIntermediates Intermediates results from the forward pass. * \param weights The LSTM weights structure. * \param input The input tensor to the LSTM, of shape * [\p timeSteps, \p batchSize, \p inputSize]. * \param output The output tensor from the forward pass. Depending * on the \p outputFullSequence parameter, this is * either the output for the last timestep or it is a * sequence of outputs for each timestep. * \param outputGrad The gradients of the output. Depending on the * \p outputFullSequence parameter, this is either the * gradient of the output for the last timestep or * it is a sequence output gradients for each * timestep. * \param lastCellStateGrad The gradient of the last cell state - may be * null if there is no incoming gradient. * \param[out] *inputSeqGrad The gradients of the inputs. May be null if * this information is not required. * \param weightsGrad A set of weight deltas to sum with weights. * \param debugContext Optional debug information. * \param options LSTM implementation options. * See createInput(). * \param planningCache The matmul planning cache. * * \return The gradient of the initial state. */ LstmState lstmBwdWithWU(poplar::Graph &graph, const LstmParams ¶ms, poplar::program::Sequence &prog, const LstmState &fwdStateInit, const poplar::Tensor &fwdIntermediates, const LstmWeights &weights, const poplar::Tensor &input, const poplar::Tensor &output, const poplar::Tensor &outputGrad, const poplar::Tensor *lastCellStateGrad, poplar::Tensor *inputGrad, LstmWeights &weightsGrad, const poplar::DebugContext &debugContext = {}, const poplar::OptionFlags &options = {}, poplin::PlanningCache *planningCache = nullptr); } // namespace lstm } // namespace popnn #endif // popnn_Lstm_hpp