DirectKPPCLayer.h

Go to the documentation of this file.

#pragma once
 
#include <deepity/backend/IComputeBackend.h>
#include <deepity/layers/Layer.h>
#include <deepity/utils/Activations.h>
#include <deepity/utils/AdamOptimizer.h>
#include <deepity/utils/DeviceMemoryArena.h>
#include <deepity/utils/MemoryArena.h>
#include <map>
#include <memory>
#include <random>
#include <vector>

 
namespace Deep {
class DirectKPPCLayer : public Layer {
protected:
  size_t size;         
  size_t nextSize;      
  size_t terminalSize;  
  size_t batchSize;     
 
  float lr;     
  float ir;     
  float fl;     
  float lmbda;  
 
  bool isClamped = false;   
  bool muCacheValid = false; 
 
  DirectKPPCLayer *layerAbove = nullptr;   
  DirectKPPCLayer *layerBelow = nullptr;   
  DirectKPPCLayer *terminalLayer = nullptr; 
 
  ActivationType activationType;  
  ActivationFn activation;        
  DerivativeFn activationDerivative;      
  DerivativeFn2 activationDerivativeInto; 
 
  OptimizerType opt = OptimizerType::SGD;    
  OptimizerType optPsi = OptimizerType::SGD; 
  int t = 0;     
  int tPsi = 0;  

  using BackendDeleter = void (*)(IComputeBackend *);
  std::unique_ptr<IComputeBackend, BackendDeleter> backend; 
 
  float *z = nullptr; 
  float *e = nullptr; 
 
  float *W = nullptr;       
  float *b = nullptr;       
  float *mu = nullptr;      
  float *cachedMu = nullptr; 
  float *Psi = nullptr;     
  float *proj = nullptr;    
 
  float *grad_W = nullptr; 
  float *grad_b = nullptr; 
  float *m_W = nullptr;    
  float *v_W = nullptr;    
  float *m_b = nullptr;    
  float *v_b = nullptr;    
 
  float *grad_Psi = nullptr; 
  float *m_Psi = nullptr;    
  float *v_Psi = nullptr;    
 
  int *t_device = nullptr;      
  float *lr_device = nullptr;   
  int *tPsi_device = nullptr;   
  float *fl_device = nullptr;   
 
  float *zF = nullptr;             
  float *zFDeriv = nullptr;        
  float *feedbackScratch = nullptr; 

  float *biasGradScratch = nullptr;
 
  std::unique_ptr<MemoryArena> localArena; 
 
public:
  DirectKPPCLayer(size_t size, size_t nextSize, size_t terminalSize,
                  size_t batchSize, float learningRate, float inferenceRate,
                  float feedback, float lmbda, ActivationType aType,
                  ActivationType dType, IComputeBackend *backend = nullptr);
 
  ~DirectKPPCLayer() override = default;

  template <typename ArenaT> void BindMemory(ArenaT &arena);
  size_t GetRequiredFloats() const noexcept;
  void RandomizeWeights(std::mt19937 &seedGenerator) noexcept;
 
  std::map<std::string, TensorDescriptor> GetStateDict() const override;

  void SetLayerAbove(DirectKPPCLayer *l) noexcept { layerAbove = l; }
  void SetLayerBelow(DirectKPPCLayer *l) noexcept { layerBelow = l; }
  void SetTerminalLayer(DirectKPPCLayer *l) noexcept { terminalLayer = l; }

  float CalculateState() noexcept override { return CalculateState(true); }
  float CalculateState(bool needEnergy) noexcept;

  void ComputeMuOnly() noexcept;
  void UpdateState() noexcept override;
  void UpdateWeights() noexcept override;
  void DirectFeedbackUpdate() noexcept;

  void ClampState(const std::vector<float> &inputData) noexcept;
  void UnclampState() noexcept;
  void ResetState() noexcept;

  bool IsClamped() const noexcept { return isClamped; }

  void SetOptimizer(OptimizerType o) noexcept { opt = o; }
  void SetPsiOptimizer(OptimizerType o) noexcept { optPsi = o; }
  void SetLearningRate(float learningRate) noexcept;
  void SetInferenceRate(float inferenceRate) noexcept { ir = inferenceRate; }
  void SetFeedbackRate(float feedbackRate) noexcept;
  void SetLambda(float lmbda) noexcept { this->lmbda = lmbda; }
 
  float *GetBeliefs() noexcept override { return z; }
  const float *GetErrors() const noexcept override { return e; }
  const float *GetMu() const noexcept { return mu; }
  const float *GetWeights() const noexcept { return W; }
  const float *GetDirectFeedbackWeights() const noexcept { return Psi; }
  const float *GetBiases() const noexcept { return b; }
 
  size_t GetBatchSize() const noexcept override { return batchSize; }
  size_t GetInputSize() const noexcept override { return size; }
  size_t GetOutputSize() const noexcept override { return nextSize; }
  size_t GetTerminalSize() const noexcept { return terminalSize; }
};
} // namespace Deep