optimizers.hpp

Go to the documentation of this file.

#pragma once
 
#include <memory>
#include <vector>
#include <unordered_map>
#include "utils/autograd.hpp"
#include "utils/matrix.hpp"
 
namespace dl::optimization {
    using utils::Variable;
    using utils::VariableD;
    using utils::VariableF;
    using utils::Matrix;
    using utils::MatrixD;
    using utils::MatrixF;
 
    template<typename T>
    class AutogradOptimizer {
    public:
        explicit AutogradOptimizer(std::vector<Variable<T>*> parameters)
            : parameters_(parameters) {}
        
        virtual ~AutogradOptimizer() = default;
        
        virtual void step() = 0;
        
        virtual void zero_grad() {
            for (auto* param : parameters_) {
                param->zero_grad();
            }
        }
        
        virtual T get_lr() const = 0;
        
        virtual void set_lr(T lr) = 0;
        
    protected:
        std::vector<Variable<T>*> parameters_;
    };
 
    template<typename T>
    class SGD : public AutogradOptimizer<T> {
    public:
        SGD(std::vector<Variable<T>*> parameters, 
            T lr, 
            T momentum = 0.0, 
            T weight_decay = 0.0, 
            bool nesterov = false);
        
        void step() override;
        
        T get_lr() const override { return lr_; }
        
        void set_lr(T lr) override { lr_ = lr; }
        
    private:
        T lr_;
        T momentum_;
        T weight_decay_;
        bool nesterov_;
        
        // Momentum buffers for each parameter
        std::vector<Matrix<T>> momentum_buffers_;
        
        void initialize_momentum_buffers();
    };
 
    template<typename T>
    class Adam : public AutogradOptimizer<T> {
    public:
        Adam(std::vector<Variable<T>*> parameters,
             T lr = 1e-3,
             T beta1 = 0.9,
             T beta2 = 0.999,
             T eps = 1e-8,
             T weight_decay = 0.0);
        
        void step() override;
        
        T get_lr() const override { return lr_; }
        
        void set_lr(T lr) override { lr_ = lr; }
        
    private:
        T lr_;
        T beta1_;
        T beta2_;
        T eps_;
        T weight_decay_;
        
        // State for each parameter
        std::vector<Matrix<T>> exp_avg_;        // First moment estimate
        std::vector<Matrix<T>> exp_avg_sq_;     // Second moment estimate
        size_t step_count_;
        
        void initialize_state();
    };
 
    template<typename T>
    class AdamW : public AutogradOptimizer<T> {
    public:
        AdamW(std::vector<Variable<T>*> parameters,
              T lr = 1e-3,
              T beta1 = 0.9,
              T beta2 = 0.999,
              T eps = 1e-8,
              T weight_decay = 1e-2);
        
        void step() override;
        
        T get_lr() const override { return lr_; }
        
        void set_lr(T lr) override { lr_ = lr; }
        
    private:
        T lr_;
        T beta1_;
        T beta2_;
        T eps_;
        T weight_decay_;
        
        // State for each parameter
        std::vector<Matrix<T>> exp_avg_;        // First moment estimate
        std::vector<Matrix<T>> exp_avg_sq_;     // Second moment estimate
        size_t step_count_;
        
        void initialize_state();
    };
 
    template<typename T>
    class RMSprop : public AutogradOptimizer<T> {
    public:
        RMSprop(std::vector<Variable<T>*> parameters,
                T lr = 1e-2,
                T alpha = 0.99,
                T eps = 1e-8,
                T weight_decay = 0.0,
                T momentum = 0.0);
        
        void step() override;
        
        T get_lr() const override { return lr_; }
        
        void set_lr(T lr) override { lr_ = lr; }
        
    private:
        T lr_;
        T alpha_;
        T eps_;
        T weight_decay_;
        T momentum_;
        
        // State for each parameter
        std::vector<Matrix<T>> square_avg_;     // Moving average of squared gradients
        std::vector<Matrix<T>> momentum_buffer_; // Momentum buffer (if momentum > 0)
        
        void initialize_state();
    };
 
    template<typename T>
    class LRScheduler {
    public:
        explicit LRScheduler(AutogradOptimizer<T>* optimizer) : optimizer_(optimizer) {}
        virtual ~LRScheduler() = default;
        
        virtual void step() = 0;
        
        T get_lr() const { return optimizer_->get_lr(); }
        
    protected:
        AutogradOptimizer<T>* optimizer_;
    };
 
    template<typename T>
    class StepLR : public LRScheduler<T> {
    public:
        StepLR(AutogradOptimizer<T>* optimizer, size_t step_size, T gamma = 0.1)
            : LRScheduler<T>(optimizer), step_size_(step_size), gamma_(gamma), 
              last_epoch_(0), base_lr_(optimizer->get_lr()) {}
        
        void step() override;
        
    private:
        size_t step_size_;
        T gamma_;
        size_t last_epoch_;
        T base_lr_;
    };
 
    // Type aliases for convenience
    using SGDD = SGD<double>;
    using SGDF = SGD<float>;
    using AdamD = Adam<double>;
    using AdamF = Adam<float>;
    using AdamWD = AdamW<double>;
    using AdamWF = AdamW<float>;
    using RMSpropD = RMSprop<double>;
    using RMSpropF = RMSprop<float>;
    using StepLRD = StepLR<double>;
    using StepLRF = StepLR<float>;
 
} // namespace dl::optimization