full_batch_minimizer.hpp

Go to the documentation of this file.

#pragma once
 
#include "../common.hpp"
#include "../iteration_recorder.hpp"
#include <Eigen/Eigen>
#include <cmath>
#include <functional>
#include <limits>
 
extern "C" {
extern double __enzyme_autodiff(void *, ...);
extern int enzyme_dup;
extern int enzyme_const;
extern int enzyme_out;
}
 
namespace cpu_mlp {
 
template <typename V, typename M> class FullBatchMinimizer {
public:
  virtual ~FullBatchMinimizer() = default;
 
  virtual V solve(V x, VecFun<V, double> &f, GradFun<V> &Gradient) = 0;
 
  virtual void setInitialHessian(const M & /*hess*/) {}
 
  virtual void setHessian(const HessFun<V, M> & /*hessFun*/) {}
 
  template <auto LossFn, typename DataType> V solve_with_enzyme(V x, DataType *data) {
    VecFun<V, double> f = [data](const V &w) -> double { return LossFn(const_cast<double *>(w.data()), data); };
    GradFun<V> Gradient = [data](const V &w) -> V {
      V grad = V::Zero(w.size());
      __enzyme_autodiff((void *)LossFn, enzyme_dup, const_cast<double *>(w.data()), grad.data(), enzyme_const, data);
      return grad;
    };
    return this->solve(x, f, Gradient);
  }
 
  void setMaxIterations(int max_iters) { _max_iters = max_iters; }
  void setMaxLineIters(int max_line) { max_line_iters = max_line; }
 
  void setArmijoMaxIter(int max_armijo) { max_line_iters = max_armijo; }
 
  void setTolerance(double tol) { _tol = tol; }
 
  unsigned int iterations() const { return _iters; }
 
  double tolerance() const { return _tol; }
  void setRecorder(::IterationRecorder<CpuBackend> *recorder) { recorder_ = recorder; }
 
protected:
  unsigned int _max_iters = 1000;
  unsigned int _iters = 0;
  double _tol = 1e-10;
  ::IterationRecorder<CpuBackend> *recorder_ = nullptr; 
 
  // Line Search Parameters
  double c1 = 1e-4;
  double c2 = 0.9;
  double rho = 0.5;
  double max_line_iters = 50;
 
  double line_search(V x, V p, VecFun<V, double> &f, GradFun<V> &Gradient) {
    double f_old = f(x);
    double grad_f_old = Gradient(x).dot(p);
    double inf = std::numeric_limits<double>::infinity();
    double alpha_min = 0.0;
    double alpha_max = inf;
    double alpha = 1.0;
 
    for (int i = 0; i < max_line_iters; ++i) {
      V x_new = x + alpha * p;
      double f_new = f(x_new);
 
      if (f_new > f_old + c1 * alpha * grad_f_old) {
        alpha_max = alpha;
        alpha = rho * (alpha_min + alpha_max);
        continue;
      }
 
      double grad_f_new_dot_p = Gradient(x_new).dot(p);
 
      if (grad_f_new_dot_p < c2 * grad_f_old) {
        alpha_min = alpha;
        if (alpha_max == inf)
          alpha *= 2;
        else
          alpha = rho * (alpha_min + alpha_max);
        continue;
      }
      return alpha;
    }
    return alpha;
  }
};
 
} // namespace cpu_mlp