softmax_regression benchmark

softmax_benchmark.cpp

#define STAN_OPENCL
#define STAN_OPENCL_CACHE
#define OPENCL_PLATFORM_ID 0
#define OPENCL_DEVICE_ID 0
#define CL_USE_DEPRECATED_OPENCL_1_1_APIS
#define __CL_ENABLE_EXCEPTIONS

#include <stan/math/prim/scal/meta/partials_return_type.hpp>
#include <stan/math.hpp>
#include <Eigen/Dense>
#include <vector>
#include <cstdio>
#include <chrono>

using namespace Eigen;
using namespace std;
using namespace stan;
using namespace stan::math;

namespace stan {
namespace math {
template<bool propto, typename T_x, typename T_alpha, typename T_beta>
typename return_type<T_x, T_alpha, T_beta>::type
categorical_logit_glm_simple_lpmf(const Matrix<int, Dynamic, 1>& y,
                                  const Matrix<T_x, Dynamic, Dynamic>& x,
                                  const T_alpha& alpha,
                                  const Matrix<T_beta, Dynamic, Dynamic>& beta) {
  typedef typename return_type<T_x, T_beta>::type T_x_beta;
  typedef typename return_type<T_x, T_beta, T_alpha>::type T_return;

  const size_t N_instances = x.rows();

  const auto& alpha_row = as_column_vector_or_scalar(alpha).transpose();

  Eigen::Matrix<T_return, Dynamic, Dynamic> tmp = (x.template cast<T_x_beta>() * beta.template cast<T_x_beta>()).array().rowwise() + alpha_row.array();

  T_return lpmf = 0;
  //iterate overt instances
  for (int i = 0; i < N_instances; i++) {
    lpmf += categorical_logit_lpmf<propto>(y[i], tmp.row(i).transpose().eval());
  }
  return lpmf;
}
}
}

template<typename T>
void perf_test(const char* name, T func, const size_t N_instances = 30000, const size_t N_attributes = 1000, const size_t N_classes = 20, int repeats = 10){
  Matrix<int, Dynamic, 1> y(N_instances);
  for(int i=0;i<N_instances;i++){
    y[i] = Matrix<int, Dynamic, 1>::Random(1)[0] % N_classes + 1;
  }
  MatrixXd x = MatrixXd::Random(N_instances, N_attributes);
  volatile Matrix<double,Dynamic, Dynamic> beta_ = MatrixXd::Random(N_attributes, N_classes);
  volatile Matrix<double,1, Dynamic> alpha_ = RowVectorXd::Random(N_classes);

  check_bounded("", "categorical outcome out of support", y, 1, N_classes);
  check_finite("", "Matrix of independent variables", x);
#ifdef STAN_OPENCL
#endif

  auto start = std::chrono::steady_clock::now();

  for (int i = 0; i < repeats; i++) {
    Matrix<var, 1, Dynamic> alpha = *const_cast<Matrix<double, 1, Dynamic>*>(&alpha_);
    Matrix<var, Dynamic, Dynamic> beta = *const_cast<Matrix<double, Dynamic, Dynamic>*>(&beta_);
    vector<var> grad_args(beta.data(), beta.data() + beta.size());
    for(int j=0;j<alpha.size();j++) {
      grad_args.push_back(alpha[j]);
    }

    var target = func(y, x, alpha, beta);
    volatile double no_opt = target.val();
    vector<double> grad;
    target.grad(grad_args, grad);
    for (double val : grad) {
      no_opt = val;
    }
    recover_memory();
  }

  int time = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::steady_clock::now() - start).count();
  printf("%s, %d, %d, %d, %d, %d\n", name, N_instances, N_attributes, N_classes, repeats, time);

  recover_memory();
  ChainableStack::instance().memalloc_.free_all();
}


int main(){
  for(int i=0;i<5;i++) {
    for(int instances : {1000, 2000, 4000, 6000, 8000, 10000}){
      perf_test("new", categorical_logit_glm_lpmf<false, double, Matrix<var, 1, Dynamic>, var>, instances, 300, 100);
      perf_test("old", categorical_logit_glm_simple_lpmf<false, double, Matrix<var, 1, Dynamic>, var>, instances, 300, 100);
    }
    for(int instances : {1000, 3000, 6000, 12000, 18000, 24000, 30000}){
      perf_test("new", categorical_logit_glm_lpmf<false, double, Matrix<var, 1, Dynamic>, var>, instances, 300, 3);
      perf_test("old", categorical_logit_glm_simple_lpmf<false, double, Matrix<var, 1, Dynamic>, var>, instances, 300, 3);
    }
    for(int instances : {30000, 100000, 200000, 400000, 600000, 800000, 1000000}){
      perf_test("new", categorical_logit_glm_lpmf<false, double, Matrix<var, 1, Dynamic>, var>, instances, 3, 3);
      perf_test("old", categorical_logit_glm_simple_lpmf<false, double, Matrix<var, 1, Dynamic>, var>, instances, 3, 3);
    }
  }
}