Layer.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <type_traits>
#include <variant>
 
#include "Dense.hpp"
#include "Conv1D.hpp"
#include "BatchNorm1D.hpp"
#include "Pool1D.hpp"
#include "activations/ActivationParameters.hpp"
 
namespace FlexNN::Layers {
 
class Layer {
 public:
  // Construction from each concrete type — implicit so
  // `NeuralNetwork net({Dense(...), Dense(...)})` works.
  Layer(Dense d) : var_(std::move(d)) {}
  Layer(Conv1D c) : var_(std::move(c)) {}
  Layer(BatchNorm1D b) : var_(std::move(b)) {}
  Layer(MaxPool1D m) : var_(std::move(m)) {}
  Layer(AvgPool1D a) : var_(std::move(a)) {}
 
  LayerType type() const noexcept {
    return std::visit([](auto&& v) { return v.type(); }, var_);
  }
 
  Activations::Activation activation() const noexcept {
    return std::visit([](auto&& v) { return v.activation(); }, var_);
  }
 
  Activations::ActivationParameters activationParams() const noexcept {
    return std::visit(
        [](auto&& v) -> Activations::ActivationParameters {
          using T = std::decay_t<decltype(v)>;
          if constexpr (std::is_same_v<T, Dense> || std::is_same_v<T, Conv1D>) {
            return v.activationParams();
          } else {
            return Activations::ActivationParameters{};
          }
        },
        var_);
  }
 
  std::pair<Eigen::MatrixXd, Eigen::MatrixXd> forward(
      const Eigen::MatrixXd& input) const {
    return std::visit([&](auto&& v) { return v.forward(input); }, var_);
  }
 
  Eigen::MatrixXd backward(const Eigen::MatrixXd& nextW,
                           const Eigen::MatrixXd& nextdZ,
                           const Eigen::MatrixXd& currZ) const {
    return std::visit(
        [&](auto&& v) { return v.backward(nextW, nextdZ, currZ); }, var_);
  }
 
  Eigen::MatrixXd backward(const Eigen::MatrixXd& upstream,
                           const Eigen::MatrixXd& currZ) const {
    return std::visit(
        [&](auto&& v) { return v.backward(upstream, currZ); }, var_);
  }
 
  void updateWeights(const Eigen::MatrixXd& dW, const Eigen::VectorXd& db,
                     double lr) noexcept {
    std::visit([&](auto&& v) { v.update(dW, db, lr); }, var_);
  }
  void update(const Eigen::MatrixXd& dW, const Eigen::VectorXd& db,
              double lr) noexcept {
    updateWeights(dW, db, lr);
  }
 
  Eigen::MatrixXd weightsMatrix() const {
    return std::visit(
        [](auto&& v) -> Eigen::MatrixXd {
          using T = std::decay_t<decltype(v)>;
          if constexpr (std::is_same_v<T, Dense>) {
            return v.weights();
          } else if constexpr (std::is_same_v<T, Conv1D>) {
            return v.weights();
          } else if constexpr (std::is_same_v<T, BatchNorm1D>) {
            return Eigen::MatrixXd();
          } else if constexpr (std::is_same_v<T, MaxPool1D> ||
                               std::is_same_v<T, AvgPool1D>) {
            return Eigen::MatrixXd();
          } else {
            return Eigen::MatrixXd();
          }
        },
        var_);
  }
 
  Eigen::MatrixXd propagate(const Eigen::MatrixXd& dZ) const {
    return std::visit(
        [&](auto&& v) -> Eigen::MatrixXd {
          using T = std::decay_t<decltype(v)>;
          if constexpr (std::is_same_v<T, Dense>) {
            return v.weights().transpose() * dZ;
          } else if constexpr (std::is_same_v<T, Conv1D>) {
            return v.propagate(dZ);
          } else if constexpr (std::is_same_v<T, BatchNorm1D>) {
            return v.propagate(dZ);
          } else if constexpr (std::is_same_v<T, MaxPool1D>) {
            return v.propagate(dZ);
          } else if constexpr (std::is_same_v<T, AvgPool1D>) {
            return v.propagate(dZ);
          } else {
            return dZ;
          }
        },
        var_);
  }
 
  std::pair<Eigen::MatrixXd, Eigen::VectorXd> grad(
      const Eigen::MatrixXd& dZ, const Eigen::MatrixXd& input) const {
    return std::visit(
        [&](auto&& v) -> std::pair<Eigen::MatrixXd, Eigen::VectorXd> {
          using T = std::decay_t<decltype(v)>;
          if constexpr (std::is_same_v<T, Dense>) {
            double m = static_cast<double>(dZ.cols());
            Eigen::MatrixXd dW = dZ * input.transpose() / m;
            Eigen::VectorXd db = dZ.rowwise().mean();
            return {dW, db};
          } else if constexpr (std::is_same_v<T, Conv1D>) {
            return v.grad(dZ, input);
          } else if constexpr (std::is_same_v<T, BatchNorm1D>) {
            return v.grad(dZ, input);
          } else if constexpr (std::is_same_v<T, MaxPool1D> ||
                               std::is_same_v<T, AvgPool1D>) {
            return v.grad(dZ, input);
          } else {
            return {Eigen::MatrixXd(), Eigen::VectorXd()};
          }
        },
        var_);
  }
 
  const auto& variant() const noexcept { return var_; }
  auto& variant() noexcept { return var_; }
 
  // Helpers for ModelIO/tests to query concrete type without visit boilerplate.
  const Dense* asDense() const noexcept { return std::get_if<Dense>(&var_); }
  Dense* asDense() noexcept { return std::get_if<Dense>(&var_); }
  const Conv1D* asConv1D() const noexcept { return std::get_if<Conv1D>(&var_); }
  Conv1D* asConv1D() noexcept { return std::get_if<Conv1D>(&var_); }
  const BatchNorm1D* asBatchNorm1D() const noexcept {
    return std::get_if<BatchNorm1D>(&var_);
  }
  BatchNorm1D* asBatchNorm1D() noexcept {
    return std::get_if<BatchNorm1D>(&var_);
  }
  const MaxPool1D* asMaxPool1D() const noexcept {
    return std::get_if<MaxPool1D>(&var_);
  }
  MaxPool1D* asMaxPool1D() noexcept { return std::get_if<MaxPool1D>(&var_); }
  const AvgPool1D* asAvgPool1D() const noexcept {
    return std::get_if<AvgPool1D>(&var_);
  }
  AvgPool1D* asAvgPool1D() noexcept { return std::get_if<AvgPool1D>(&var_); }
 
 private:
  std::variant<Dense, Conv1D, BatchNorm1D, MaxPool1D, AvgPool1D> var_;
};
 
} // namespace FlexNN::Layers