layers.hpp

Go to the documentation of this file.

#pragma once
 
#include <memory>
#include <vector>
#include <functional>
#include "utils/autograd.hpp"
#include "utils/matrix.hpp"
 
namespace dl::layers {
    using utils::Variable;
    using utils::VariableD;
    using utils::VariableF;
    using utils::Matrix;
    using utils::MatrixD;
    using utils::MatrixF;
 
    template<typename T>
    class Module {
    public:
        virtual ~Module() = default;
        
        virtual Variable<T> forward(const Variable<T>& input) = 0;
        
        virtual std::vector<Variable<T>*> parameters() = 0;
        
        virtual void zero_grad() {
            for (auto* param : parameters()) {
                param->zero_grad();
            }
        }
        
        virtual void train(bool training = true) {
            training_ = training;
        }
        
        virtual void eval() {
            train(false);
        }
        
        bool is_training() const { return training_; }
        
    protected:
        bool training_ = true;
    };
 
    template<typename T>
    class Linear : public Module<T> {
    public:
        Linear(size_t in_features, size_t out_features, bool bias = true);
        
        Variable<T> forward(const Variable<T>& input) override;
        
        std::vector<Variable<T>*> parameters() override;
        
        // Getters for parameters
        Variable<T>& weight() { return weight_; }
        Variable<T>& bias() { return bias_; }
        const Variable<T>& weight() const { return weight_; }
        const Variable<T>& bias() const { return bias_; }
        
    private:
        Variable<T> weight_;  // Shape: (out_features, in_features)
        Variable<T> bias_;    // Shape: (out_features,)
        bool has_bias_;
        size_t in_features_;
        size_t out_features_;
        
        void initialize_parameters();
    };
 
    template<typename T>
    class ReLU : public Module<T> {
    public:
        Variable<T> forward(const Variable<T>& input) override;
        std::vector<Variable<T>*> parameters() override { return {}; }
    };
 
    template<typename T>
    class Sigmoid : public Module<T> {
    public:
        Variable<T> forward(const Variable<T>& input) override;
        std::vector<Variable<T>*> parameters() override { return {}; }
    };
 
    template<typename T>
    class Tanh : public Module<T> {
    public:
        Variable<T> forward(const Variable<T>& input) override;
        std::vector<Variable<T>*> parameters() override { return {}; }
    };
 
    template<typename T>
    class Dropout : public Module<T> {
    public:
        explicit Dropout(T p = 0.5) : p_(p) {}
        
        Variable<T> forward(const Variable<T>& input) override;
        std::vector<Variable<T>*> parameters() override { return {}; }
        
    private:
        T p_;  // Dropout probability
    };
 
    template<typename T>
    class Sequential : public Module<T> {
    public:
        void add_module(std::shared_ptr<Module<T>> module);
        
        Variable<T> forward(const Variable<T>& input) override;
        
        std::vector<Variable<T>*> parameters() override;
        
        void zero_grad() override;
        
        void train(bool training = true) override;
        
    private:
        std::vector<std::shared_ptr<Module<T>>> modules_;
    };
 
    // Type aliases for convenience
    using LinearD = Linear<double>;
    using LinearF = Linear<float>;
    using ReLUD = ReLU<double>;
    using ReLUF = ReLU<float>;
    using SigmoidD = Sigmoid<double>;
    using SigmoidF = Sigmoid<float>;
    using TanhD = Tanh<double>;
    using TanhF = Tanh<float>;
    using DropoutD = Dropout<double>;
    using DropoutF = Dropout<float>;
    using SequentialD = Sequential<double>;
    using SequentialF = Sequential<float>;
 
} // namespace dl::layers