lstm_module.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <optional>
#include <initializer_list>
#include <vector>
 
#include "pulsatrix/module.hpp"
 
namespace pulsatrix {
 
class LSTMModule : public Module {
public:
    LSTMModule(int64_t input_size, int64_t hidden_size, DeviceBackend* backend);
 
    [[nodiscard]] Tensor backward(const Tensor& grad_output) override;
 
    [[nodiscard]] OpType op_type() const override { return OpType::Recurrent; }
 
    void set_weight_xi(std::initializer_list<float> values);
    void set_weight_hi(std::initializer_list<float> values);
    void set_bias_i(std::initializer_list<float> values);
    void set_weight_xf(std::initializer_list<float> values);
    void set_weight_hf(std::initializer_list<float> values);
    void set_bias_f(std::initializer_list<float> values);
    void set_weight_xg(std::initializer_list<float> values);
    void set_weight_hg(std::initializer_list<float> values);
    void set_bias_g(std::initializer_list<float> values);
    void set_weight_xo(std::initializer_list<float> values);
    void set_weight_ho(std::initializer_list<float> values);
    void set_bias_o(std::initializer_list<float> values);
 
    [[nodiscard]] const Tensor& weight_xi() const { return weight_xi_; }
    [[nodiscard]] const Tensor& weight_hi() const { return weight_hi_; }
    [[nodiscard]] const Tensor& bias_i() const { return bias_i_; }
    [[nodiscard]] const Tensor& weight_xf() const { return weight_xf_; }
    [[nodiscard]] const Tensor& weight_hf() const { return weight_hf_; }
    [[nodiscard]] const Tensor& bias_f() const { return bias_f_; }
    [[nodiscard]] const Tensor& weight_xg() const { return weight_xg_; }
    [[nodiscard]] const Tensor& weight_hg() const { return weight_hg_; }
    [[nodiscard]] const Tensor& bias_g() const { return bias_g_; }
    [[nodiscard]] const Tensor& weight_xo() const { return weight_xo_; }
    [[nodiscard]] const Tensor& weight_ho() const { return weight_ho_; }
    [[nodiscard]] const Tensor& bias_o() const { return bias_o_; }
 
    [[nodiscard]] const Tensor& weight_xi_grad() const { return weight_xi_grad_; }
    [[nodiscard]] const Tensor& weight_hi_grad() const { return weight_hi_grad_; }
    [[nodiscard]] const Tensor& bias_i_grad() const { return bias_i_grad_; }
    [[nodiscard]] const Tensor& weight_xf_grad() const { return weight_xf_grad_; }
    [[nodiscard]] const Tensor& weight_hf_grad() const { return weight_hf_grad_; }
    [[nodiscard]] const Tensor& bias_f_grad() const { return bias_f_grad_; }
    [[nodiscard]] const Tensor& weight_xg_grad() const { return weight_xg_grad_; }
    [[nodiscard]] const Tensor& weight_hg_grad() const { return weight_hg_grad_; }
    [[nodiscard]] const Tensor& bias_g_grad() const { return bias_g_grad_; }
    [[nodiscard]] const Tensor& weight_xo_grad() const { return weight_xo_grad_; }
    [[nodiscard]] const Tensor& weight_ho_grad() const { return weight_ho_grad_; }
    [[nodiscard]] const Tensor& bias_o_grad() const { return bias_o_grad_; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] std::vector<NamedParamRef> named_parameters() override {
        return {
            {"weight_xi", {&weight_xi_, &weight_xi_grad_}},
            {"weight_hi", {&weight_hi_, &weight_hi_grad_}},
            {"bias_i", {&bias_i_, &bias_i_grad_}},
            {"weight_xf", {&weight_xf_, &weight_xf_grad_}},
            {"weight_hf", {&weight_hf_, &weight_hf_grad_}},
            {"bias_f", {&bias_f_, &bias_f_grad_}},
            {"weight_xg", {&weight_xg_, &weight_xg_grad_}},
            {"weight_hg", {&weight_hg_, &weight_hg_grad_}},
            {"bias_g", {&bias_g_, &bias_g_grad_}},
            {"weight_xo", {&weight_xo_, &weight_xo_grad_}},
            {"weight_ho", {&weight_ho_, &weight_ho_grad_}},
            {"bias_o", {&bias_o_, &bias_o_grad_}},
        };
    }
 
 
    [[nodiscard]] std::optional<DeviceType> compute_device() const override { return backend_->device(); }
 
protected:
    [[nodiscard]] Tensor forward_impl(const Tensor& input) override;
 
private:
    int64_t input_size_;
    int64_t hidden_size_;
    DeviceBackend* backend_;
    Tensor weight_xi_;  // (input_size, hidden_size)
    Tensor weight_hi_;  // (hidden_size, hidden_size)
    Tensor bias_i_;     // (hidden_size,)
    Tensor weight_xf_;
    Tensor weight_hf_;
    Tensor bias_f_;
    Tensor weight_xg_;
    Tensor weight_hg_;
    Tensor bias_g_;
    Tensor weight_xo_;
    Tensor weight_ho_;
    Tensor bias_o_;
    Tensor weight_xi_grad_;
    Tensor weight_hi_grad_;
    Tensor bias_i_grad_;
    Tensor weight_xf_grad_;
    Tensor weight_hf_grad_;
    Tensor bias_f_grad_;
    Tensor weight_xg_grad_;
    Tensor weight_hg_grad_;
    Tensor bias_g_grad_;
    Tensor weight_xo_grad_;
    Tensor weight_ho_grad_;
    Tensor bias_o_grad_;
    Tensor last_input_;           // (N, L, input_size)
    Tensor last_hidden_states_;   // (N, L+1, hidden_size); index 0 = h_0 = 0
    Tensor last_cell_states_;     // (N, L+1, hidden_size); index 0 = c_0 = 0
    Tensor last_gate_i_;          // (N, L, hidden_size)
    Tensor last_gate_f_;          // (N, L, hidden_size)
    Tensor last_gate_g_;          // (N, L, hidden_size); the cell candidate (tanh) signal
    Tensor last_gate_o_;          // (N, L, hidden_size)
    Tensor last_cell_tanh_;       // (N, L, hidden_size); tanh(c_t), cached for backward
    Tensor last_pre_activation_g_;  // (N, L, hidden_size); g_t's z EXCLUDING bias, for LRP
    int64_t last_L_ = 0;
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix