0

我正在加快我在 R 中编写的程序的速度。代码涉及在多维数组上重复计算 LogSumExp，即计算 s_lnj = exp(u_lnj) / (1 + sum_k exp(u_lnk))。我试图提高速度的代码的基本 R 版本如下：

log_sum_exp_func <- function(vec){
  max_vec <- max(vec)
  return(max_vec + log(sum(exp(vec-max_vec))))
}

compute_share_from_utils_func <- function(u_lnj){
  ### get dimensions
  L <- dim(u_lnj)[1]; n_poly <- dim(u_lnj)[2]; J <- dim(u_lnj)[3]
  
  ### compute denominator of share, 1 + sum exp utils
  den_ln <- 1 + exp(apply(u_lnj, c(1,2), log_sum_exp_func))
  den_lnj <- array(rep(den_ln, J), dim = c(L, n_poly, J))
  
  ### take ratio of utils and denominator
  s_lnj <- exp(u_lnj) / den_lnj
  return(s_lnj)
}

我尝试使用 xtensor 和 Rcpp 来加快速度，但遇到了几个问题。我写的Rcpp代码如下

// [[Rcpp::depends(xtensor)]]
// [[Rcpp::plugins(cpp14)]]
#include <numeric>                    // Standard library import for std::accumulate
#define STRICT_R_HEADERS              // Otherwise a PI macro is defined in R
#include "xtensor/xmath.hpp"          // xtensor import for the C++ universal functions
#include "xtensor/xarray.hpp"
#include "xtensor/xio.hpp"
#include "xtensor/xview.hpp"
#include "xtensor-r/rarray.hpp"       // R bindings
#include <Rcpp.h>

using namespace Rcpp;

// [[Rcpp::export]]
double cxxlog_sum_exp_vec(xt::rarray<double>& m)
{
  auto shape_m = m.shape();
  double maxvec = xt::amax(m)[0];
  xt::rarray<double> arr_maxvec = maxvec * xt::ones<double>(shape_m);
  xt::rarray<double> vec_min_max = m - arr_maxvec;
  xt::rarray<double> exp_vec_min_max = xt::exp(vec_min_max);
  double sum_exp = xt::sum(exp_vec_min_max)[0];
  double log_sum_exp = std::log(sum_exp);
  return log_sum_exp + maxvec; 
}

// [[Rcpp::export]]
xt::rarray<double> cxxshare_from_utils(xt::rarray<double>& u_lnj)
{
  int L = u_lnj.shape(0);
  int N = u_lnj.shape(1);
  int J = u_lnj.shape(2);
  xt::rarray<double> res = xt::ones<double>({L,N,J});
  for (std::size_t l = 0; l < u_lnj.shape()[0]; ++l)
  {
    for (std::size_t n = 0; n < u_lnj.shape()[1]; ++n)
    {
      xt::rarray<double> utils_j = xt::view(u_lnj, l, n, xt::all());
      double inv_lse = 1 / (1 + std::exp(cxxlog_sum_exp_vec(utils_j)));
      for (std::size_t j = 0; j < J; ++j)
      {
        res(l, n, j) = std::exp(u_lnj(l, n, j)) * inv_lse;
      }
    }
  }
  return res;
}

Rcpp 实现似乎确实产生了与基本 R 代码相同的结果，但是每当输入数组的维度增加时，它似乎都会遇到问题。如果我运行，我的 R 会话将失败

L <- 100
n <- 100
J <- 200
u_lnj <- array(rnorm(L*n*J,0,2), dim = c(L, n, J))
test <- cxxshare_from_utils(u_lnj)

但是例如，对于 L、n、J = 10、10、20，代码运行良好。此外，log_sum_exp 的 C++ 实现似乎并没有比基础 R 版本好太多。

编辑：我无法弄清楚我使用 xtensor 的方式有什么问题。但我确实使用以下 RcppArmadillo 代码加快了速度。这个版本的缺点是它可能不像依赖 Log Sum Exp 的基本 R 函数那样鲁棒性溢出。

#include <RcppArmadillo.h>
// [[Rcpp::depends(RcppArmadillo)]]
// [[Rcpp::plugins(cpp14)]]
// [[Rcpp::export]]
arma::cube cxxarma_share_from_utils(arma::cube u_lnj) {
  
  // Extract the different dimensions
  
  // Normal Matrix dimensions
  unsigned int L = u_lnj.n_rows;
  unsigned int N = u_lnj.n_cols;
  
  // Depth of Array
  unsigned int J = u_lnj.n_slices;
  
  //resulting cube
  arma::cube s_lnj = arma::exp(u_lnj);
  for (unsigned int l = 0; l < L; l++) {
    
    for (unsigned int n = 0; n < N; n++) {
      
      double den = 1 / (1 + arma::accu(s_lnj.subcube(arma::span(l), arma::span(n), arma::span())));
      
      for (unsigned int j = 0; j < J; j++) {
        
        s_lnj(l, n, j) = s_lnj(l, n, j) * den;
      }
    }
  }
  return s_lnj;
}
4

0 回答 0