lstm_layer.h

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-only
// Copyright (C) 2026 Artificial Intelligence Techniques, S.L. and Roberto López González
 
#pragma once
 
#include "opennn/network/layers/layer.h"
#include "opennn/network/operators/operator.h"
#include "opennn/network/operators/activation_operator.h"
#include "opennn/network/operators/cudnn_rnn.h"
 
namespace opennn
{
 
struct LSTMOperator : Operator, CudnnRnnState
{
    enum ForwardSlot
    {
        InputSlot = 0,
        ForgetGateSlot,
        InputGateSlot,
        CandidateGateSlot,
        OutputGateSlot,
        CellStateSlot,
        HiddenStateSlot,
        CellActivationSlot,
        CudnnInputSequenceSlot,
        CudnnOutputSequenceSlot,
        OutputSlot
    };
 
    enum BackwardSlot
    {
        OutputDeltaSlot = 0,
        InputDeltaSlot,
        HiddenDeltaScratchSlot,
        CellDeltaScratchSlot,
        ForgetDeltaScratchSlot,
        InputDeltaScratchSlot,
        CandidateDeltaScratchSlot,
        OutputDeltaScratchSlot,
        CudnnOutputDeltaScratchSlot,
        CudnnInputDeltaScratchSlot
    };
 
    Index input_features  = 0;
    Index output_features = 0;
    Index time_steps      = 0;
 
    bool return_sequences = false;
 
    ActivationFunction activation_function = ActivationFunction::Tanh;
    ActivationFunction recurrent_activation_function = ActivationFunction::Sigmoid;
 
    TensorView forget_bias;
    TensorView input_bias;
    TensorView candidate_bias;
    TensorView output_bias;
 
    TensorView forget_weights;
    TensorView input_weights;
    TensorView candidate_weights;
    TensorView output_weights;
 
    TensorView forget_recurrent_weights;
    TensorView input_recurrent_weights;
    TensorView candidate_recurrent_weights;
    TensorView output_recurrent_weights;
 
    TensorView forget_bias_gradient;
    TensorView input_bias_gradient;
    TensorView candidate_bias_gradient;
    TensorView output_bias_gradient;
 
    TensorView forget_weight_gradient;
    TensorView input_weight_gradient;
    TensorView candidate_weight_gradient;
    TensorView output_weight_gradient;
 
    TensorView forget_recurrent_weight_gradient;
    TensorView input_recurrent_weight_gradient;
    TensorView candidate_recurrent_weight_gradient;
    TensorView output_recurrent_weight_gradient;
 
    void set(Index,
             Index,
             Index,
             ActivationFunction new_activation_function = ActivationFunction::Tanh,
             ActivationFunction new_recurrent_activation_function = ActivationFunction::Sigmoid,
             Type new_compute_dtype = Type::FP32);
 
    // The four gates are always taken in this order, and every site that
    // touches them -- specs, linking, initialization, the cuDNN pack -- used to
    // spell the quartet out by hand. That was 34 copies of the same four names.
    enum Gate { Forget, Input, Candidate, Output, GateCount };
 
    // A TensorView is itself a view, so reading the quartet never needs a
    // mutable pointer; only parameter_slots does, and it pairs the members up
    // by hand so a parameter can never drift away from its gradient.
    using GateViews = array<const TensorView*, GateCount>;
 
    GateViews gate_biases() const            { return {&forget_bias, &input_bias, &candidate_bias, &output_bias}; }
    GateViews gate_weights() const           { return {&forget_weights, &input_weights, &candidate_weights, &output_weights}; }
    GateViews gate_recurrent_weights() const { return {&forget_recurrent_weights, &input_recurrent_weights,
                                                       &candidate_recurrent_weights, &output_recurrent_weights}; }
 
    GateViews gate_bias_gradients() const    { return {&forget_bias_gradient, &input_bias_gradient,
                                                       &candidate_bias_gradient, &output_bias_gradient}; }
    GateViews gate_weight_gradients() const  { return {&forget_weight_gradient, &input_weight_gradient,
                                                       &candidate_weight_gradient, &output_weight_gradient}; }
    GateViews gate_recurrent_weight_gradients() const { return {&forget_recurrent_weight_gradient,
                                                                &input_recurrent_weight_gradient,
                                                                &candidate_recurrent_weight_gradient,
                                                                &output_recurrent_weight_gradient}; }
 
    vector<TensorSpec> parameter_specs() const override;
    vector<ParameterSlot> parameter_slots() override;
 
    void set_parameters_random() override;
    void set_parameters_glorot() override;
    void set_parameters_pytorch() override;
 
    void forward_propagate(ForwardPropagation&, size_t, ForwardPropagationMode) override;
    void back_propagate(ForwardPropagation&, BackPropagation&, size_t) const override;
 
private:
    void apply(const TensorView&,
               TensorView&,
               TensorView&,
               TensorView&,
               TensorView&,
               TensorView&,
               TensorView&,
               TensorView&,
               TensorView&,
               bool training) const;
 
    void apply_delta(const TensorView&,
                     const TensorView&,
                     TensorView&,
                     TensorView&,
                     TensorView&,
                     TensorView&,
                     TensorView&,
                     TensorView&,
                     TensorView&,
                     const TensorView&,
                     const TensorView&,
                     const TensorView&,
                     const TensorView&,
                     const TensorView&,
                     const TensorView&,
                     const TensorView&) const;
 
    bool apply_onednn(const TensorView&,
                      TensorView&,
                      TensorView&,
                      TensorView&,
                      TensorView&,
                      Buffer&,
                      bool is_training,
                      uint64_t parameters_version) const;
 
    bool apply_delta_onednn(const TensorView&,
                            const TensorView&,
                            const TensorView&,
                            const TensorView&,
                            TensorView&,
                            TensorView&,
                            TensorView&,
                            const Buffer&,
                            Buffer&,
                            bool return_seq) const;
 
    void apply_gpu(const TensorView&,
                   TensorView&,
                   TensorView&,
                   TensorView&,
                   TensorView&,
                   Buffer&,
                   bool return_seq,
                   bool is_training,
                   uint64_t parameters_version) const;
 
    void apply_delta_gpu(const TensorView&,
                         const TensorView&,
                         const TensorView&,
                         const TensorView&,
                         const TensorView&,
                         TensorView&,
                         TensorView&,
                         TensorView&,
                         const Buffer&,
                         Buffer&,
                         bool return_seq) const;
 
    CudnnRnnShapeSlot& ensure_cudnn_setup_(Index batch_size, bool for_training) const;
    void pack_weights_to_cudnn_(Buffer&, uint64_t) const;
    void unpack_gradients_from_cudnn_(Buffer&) const;
 
};
 
class LSTM final : public Layer
{
public:
 
    LSTM(const Shape& = {},
                        const Shape& = {},
                        const string& = "Tanh",
                        const string& = "Sigmoid",
                        const string& = "lstm_layer");
 
    Shape get_input_shape()  const noexcept override { return input_shape; }
    Shape get_output_shape() const override
    {
        return return_sequences ? Shape{get_time_steps(), output_features}
                                : Shape{output_features};
    }
 
    Index get_time_steps()      const noexcept { return input_shape.get_rank() == 2 ? input_shape[0] : Index(0); }
    Index get_input_features()  const noexcept { return input_shape.get_rank() == 2 ? input_shape[1] : Index(0); }
    Index get_output_features() const noexcept { return output_features; }
 
    bool get_return_sequences() const noexcept { return return_sequences; }
    void set_return_sequences(bool value);
 
    const TensorView& get_forget_bias()    const noexcept { return lstm_op.forget_bias; }
 
    const ActivationFunction& get_activation_function() const noexcept { return lstm_op.activation_function; }
    const ActivationFunction& get_recurrent_activation_function() const noexcept { return lstm_op.recurrent_activation_function; }
    ActivationFunction get_output_activation() const noexcept override { return lstm_op.activation_function; }
 
    vector<TensorSpec> get_forward_specs(Index)  const override;
    vector<TensorSpec> get_backward_specs(Index) const override;
 
    void set(const Shape& = {},
             const Shape& = {},
             const string& = "Tanh",
             const string& = "Sigmoid",
             const string& = "lstm_layer");
 
    bool accepts_input_rank(Index rank) const override { return is_one_of(rank, 2); }
    bool is_recurrent() const noexcept override { return true; }
 
    void apply_input_shape(const Shape&) override;
    void set_output_shape(const Shape&) override;
    void on_compute_dtype_changed() override { configure_operators(); }
 
    void set_activation_function(const string&);
    void set_recurrent_activation_function(const string&);
 
    void read_JSON_body(const Json*) override;
    void write_JSON_body(JsonWriter&) const override;
 
    string write_expression(const vector<string>&,
                            const vector<string>&) const override;
 
private:
 
    Index output_features = 0;
    bool  return_sequences = false;
 
    LSTMOperator lstm_op;
 
    void configure_operators();
};
 
}