rnn_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 RNNModule : public Module {
public:
    RNNModule(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_xh(std::initializer_list<float> values);
    void set_weight_hh(std::initializer_list<float> values);
    void set_bias(std::initializer_list<float> values);
 
    [[nodiscard]] const Tensor& weight_xh() const { return weight_xh_; }
    [[nodiscard]] const Tensor& weight_hh() const { return weight_hh_; }
    [[nodiscard]] const Tensor& bias() const { return bias_; }
    [[nodiscard]] const Tensor& weight_xh_grad() const { return weight_xh_grad_; }
    [[nodiscard]] const Tensor& weight_hh_grad() const { return weight_hh_grad_; }
    [[nodiscard]] const Tensor& bias_grad() const { return bias_grad_; }
 
    [[nodiscard]] Tensor propagate_relevance(const Tensor& relevance_out, const LRPRuleConfig& config) override;
 
    [[nodiscard]] std::vector<NamedParamRef> named_parameters() override {
        return {
            {"weight_xh", {&weight_xh_, &weight_xh_grad_}},
            {"weight_hh", {&weight_hh_, &weight_hh_grad_}},
            {"bias", {&bias_, &bias_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_xh_;       // (input_size, hidden_size)
    Tensor weight_hh_;       // (hidden_size, hidden_size)
    Tensor bias_;            // (hidden_size,)
    Tensor weight_xh_grad_;
    Tensor weight_hh_grad_;
    Tensor bias_grad_;
    Tensor last_input_;             // (N, L, input_size)
    Tensor last_hidden_states_;     // (N, L+1, hidden_size); index 0 = h_0 = 0
    Tensor last_pre_activation_;    // (N, L, hidden_size); z_t EXCLUDING bias, cached for LRP
    int64_t last_L_ = 0;
    bool has_forwarded_ = false;
};
 
}  // namespace pulsatrix