BatchNorm1D.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <utility>
 
#include <Eigen/Dense>
 
#include "LayerTypes.hpp"
#include "../activations/Activation.hpp"
 
namespace FlexNN::Layers {
 
class BatchNorm1D {
 public:
  explicit BatchNorm1D(BatchNormParams p,
                       Activations::Activation act = Activations::Activation::None);
  explicit BatchNorm1D(int numFeatures, float epsilon = 1e-5f,
                       Activations::Activation act = Activations::Activation::None);
 
  LayerType type() const noexcept { return LayerType::BatchNorm1D; }
  Activations::Activation activation() const noexcept { return act_; }
  const BatchNormParams& params() const noexcept { return params_; }
 
  // Accessors for ModelIO and tests (gamma/beta/mean/var each [F])
  const Eigen::VectorXd& gamma() const noexcept { return gamma_; }
  const Eigen::VectorXd& beta() const noexcept { return beta_; }
  const Eigen::VectorXd& runningMean() const noexcept { return runningMean_; }
  const Eigen::VectorXd& runningVar() const noexcept { return runningVar_; }
  void setGamma(const Eigen::VectorXd& g) { gamma_ = g; }
  void setBeta(const Eigen::VectorXd& b) { beta_ = b; }
  void setRunningMean(const Eigen::VectorXd& m) { runningMean_ = m; }
  void setRunningVar(const Eigen::VectorXd& v) { runningVar_ = v; }
 
  std::pair<Eigen::MatrixXd, Eigen::MatrixXd> forward(
      const Eigen::MatrixXd& input) const;
 
  std::pair<Eigen::MatrixXd, Eigen::MatrixXd> forward(
      const Eigen::MatrixXd& input, bool training) const;
 
  // Backward with upstream (for hidden layers)
  Eigen::MatrixXd backward(const Eigen::MatrixXd& upstream,
                           const Eigen::MatrixXd& currZ) const;
  // Legacy 3-arg for compatibility (delegates to upstream version)
  Eigen::MatrixXd backward(const Eigen::MatrixXd& nextW,
                           const Eigen::MatrixXd& nextdZ,
                           const Eigen::MatrixXd& currZ) const;
 
  Eigen::MatrixXd propagate(const Eigen::MatrixXd& dZ) const;
 
  std::pair<Eigen::MatrixXd, Eigen::VectorXd> grad(
      const Eigen::MatrixXd& dZ, const Eigen::MatrixXd& input) const;
 
  void update(const Eigen::MatrixXd& dW, const Eigen::VectorXd& db,
              double lr) noexcept;
  void updateWeights(const Eigen::MatrixXd& dW, const Eigen::VectorXd& db,
                     double lr) noexcept {
    update(dW, db, lr);
  }
 
 private:
  BatchNormParams params_{};
  Activations::Activation act_ = Activations::Activation::None;
  Eigen::VectorXd gamma_;       
  Eigen::VectorXd beta_;        
  mutable Eigen::VectorXd runningMean_; 
  mutable Eigen::VectorXd runningVar_;  
};
 
} // namespace FlexNN::Layers