adam_optimizer.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <unordered_map>
 
#include "pulsatrix/module.hpp"
 
namespace pulsatrix {
 
class AdamOptimizer {
public:
    explicit AdamOptimizer(float learning_rate, DeviceBackend* backend, float beta1 = 0.9f, float beta2 = 0.999f,
                            float eps = 1e-8f);
 
    void step(Module& module);
 
    void zero_grad(Module& module);
 
    [[nodiscard]] float learning_rate() const { return learning_rate_; }
 
    void set_learning_rate(float learning_rate) { learning_rate_ = learning_rate; }
 
private:
    struct AdamState {
        Tensor m;
        Tensor v;
        int64_t t = 0;
    };
 
    float learning_rate_;
    DeviceBackend* backend_;
    float beta1_;
    float beta2_;
    float eps_;
    std::unordered_map<const Tensor*, AdamState> state_;
};
 
}  // namespace pulsatrix