DirectKPPCNetwork.h

Go to the documentation of this file.

#pragma once
#include <vector>
#include <memory>
#include <random>
#include <deepity/layers/DirectKPPCLayer.h>
#include <deepity/utils/MemoryArena.h>
#include <deepity/utils/DeviceMemoryArena.h>
#include <deepity/backend/Backend.h>

 
namespace Deep
{
    class DirectKPPCNetwork
    {
    public:
        explicit DirectKPPCNetwork(int batchSize, DeviceType device = DeviceType::DEVICE_CPU) noexcept;
        ~DirectKPPCNetwork() = default;

        DirectKPPCNetwork(const DirectKPPCNetwork &) = delete;
        DirectKPPCNetwork &operator=(const DirectKPPCNetwork &) = delete;

        void AddLayer(size_t size, size_t nextSize, size_t terminalSize,
                      float lr, float ir, float fl, float lmbda,
                      ActivationType aType, ActivationType dType);

        void RandomizeWeights(std::mt19937 &rng);
        void ResetState() noexcept;
        void Clamp(const std::vector<float> &input);

        void ProjectForward() noexcept;

        float CalculateTerminalError() noexcept;

        void DirectFeedbackUpdate() noexcept;

        float Step(bool computeEnergy = true) noexcept;

        void UpdateWeights() noexcept;

        void SetOptimizer(OptimizerType o) noexcept
        {
            for (auto &layer : layers)
                layer->SetOptimizer(o);
        }

        void SetPsiOptimizer(OptimizerType o) noexcept
        {
            for (auto &layer : layers)
                layer->SetPsiOptimizer(o);
        }

        void SetLearningRate(float lr) noexcept
        {
            for (auto &layer : layers)
                layer->SetLearningRate(lr);
        }

        void SetFeedbackRate(float fl) noexcept
        {
            for (auto &layer : layers)
                layer->SetFeedbackRate(fl);
        }

        DirectKPPCLayer *GetTerminalLayer() noexcept { return layers.back().get(); }
        std::vector<std::unique_ptr<DirectKPPCLayer>> &GetLayers() noexcept { return layers; }
        const std::vector<std::unique_ptr<DirectKPPCLayer>> &GetLayers() const noexcept { return layers; }
        int GetBatchSize() const noexcept { return batchSize; }
        DeviceType GetDevice() const noexcept { return device; }

        float TrainStep(const std::vector<float> &x, const std::vector<float> &y,
                        int inferenceSteps = 1);

        std::vector<float> Predict(const std::vector<float> &x, int inferenceSteps);

        void Compile();
 
    private:
        std::vector<std::unique_ptr<DirectKPPCLayer>> layers;

        std::unique_ptr<IComputeBackend> backend;
        DeviceType device;

        std::unique_ptr<MemoryArena> cpuArena;
#if defined(DEEPITY_USE_CUDA)
        std::unique_ptr<DeviceMemoryArena> gpuArena;
#endif

        int batchSize;
 
        // @brief CUDA Graph:
        bool graphCaptured = false;
        int capturedInferenceSteps = -1;
    };
}