如何使用bigstatsr R包对两个大型数据集并行估计参数并优化计算开销

步骤1：将数据集转换为FBM格式 #

bigstatsr通过内存映射格式FBM存储超大矩阵，无需将全部数据加载到内存，且按列存储的特性刚好适配你每次取一列自变量、一列因变量的计算需求：

library(bigstatsr)
library(pracma)

# 若原始数据为csv，不要用read.csv加载，直接用big_read读取为FBM，避免内存溢出
# test_A_fbm <- big_read("自变量数据集路径", type = "double")
# test_B_fbm <- big_read("因变量数据集路径", type = "double")

# 若数据已加载为data.frame，直接转换为FBM
test_A_fbm <- as_FBM(test_A)
test_B_fbm <- as_FBM(test_B)

# 提前获取变量数量，无需生成expand.grid
n_A <- ncol(test_A_fbm)
n_B <- ncol(test_B_fbm)
total_tasks <- n_A * n_B

步骤2：重写适配FBM的计算函数 #

调整函数调用逻辑，避免全局变量依赖，直接从FBM中读取需要的列：

# 自研方法保持不变
Proposed_method<- function(Data, Beta) 
{ 
  n = dim(Data)[1]
  Median <- t(apply(Data,2,median))
  Dist <- sqrt(rowSums((Data - as.matrix(rep(1,dim(Data)[1]))%*%Median)^2))
  Data0 <- as.matrix(Data[which(Dist <= as.numeric(quantile(Dist, p=.45, na.rm = TRUE))),])
  Yo <- as.matrix(Data0[,dim(Data0)[2]])
  Xo <- as.matrix(Data0[,-dim(Data0)[2]])
  Gama0 <- as.numeric(pinv(crossprod(Xo, Xo))%*%crossprod(Xo, Yo))
  Sigma2o <- var(Yo)
  Y <- as.matrix(Data[,dim(Data)[2]])
  X <- as.matrix(Data[,-dim(Data)[2]])
  
  DiffTol = 0.0001;
  DiffNorm = +10000;
  Iter = 0;
  while (DiffNorm > DiffTol)
  {
    Const <- sqrt(2*pi*Sigma2o)
    devmat <- (Y-X%*%Gama0)
    Squaremat <- as.matrix(apply(devmat, c(1,2), function(x) x^2))
    Gauss <- exp(-Squaremat/(2*as.numeric(Sigma2o)))/as.numeric(Const)
    Wbeta <- exp(-(Beta*((Y-X%*%Gama0)*(Y-X%*%Gama0)))/(2*as.numeric(Sigma2o)))
    ONE1 <- rep(1,dim(X)[2]);
    Xb <- (X*(Wbeta%*%ONE1)) 
    Gama <- as.numeric(pinv(crossprod(X, Xb))%*%crossprod(Xb, Y)) 
    hedprod <- (Y-X%*%Gama)*(Y-X%*%Gama) 
    tWbeta <- as.matrix(t(Wbeta)) 
    One_1 <- as.matrix(rep(1,dim(X)[1])) 
    Sigma2 <- (tWbeta%*%hedprod)*pinv(tWbeta%*%One_1)
    
    LHb<-(sum(Gauss^Beta)/n-1)/Beta
    LH<-prod(Gauss)
    Norm2 <- ((sum(Gama*Gama))^0.5 + abs(Sigma2))
    DiffNorm <-((sum((Gama-Gama0)*(Gama-Gama0)))^0.5 + abs(Sigma2 - Sigma2o))/Norm2
    Gama0 = Gama
    Sigma2o=Sigma2
    Iter = Iter + 1 
  }
  return(list(Gama=Gama,Sigma2=Sigma2,Wt=Wbeta,LHb=LHb,LH=LH))
}

# 适配FBM的单任务计算函数
test_function_fbm <- function(A_fbm, B_fbm, a_idx, b_idx, Beta = 0.1, Omit = 2){
  c1 <- A_fbm[, a_idx]
  c2 <- B_fbm[, b_idx]
  Data <- data.frame(1, XX=c1, YY=c2)
  nn <- nrow(Data)
  
  ResL1 <- Proposed_method(Data, Beta)
  ResL0 <- Proposed_method(as.matrix(Data[,-Omit]), Beta)
  LR0 <- (-nn)*log(ResL1$Sigma2/ResL0$Sigma2)
  
  Proposed_estimator <- (ResL1$Gama)[2]
  Proposed_pvalue <- as.numeric(pchisq(q=LR0, df=1, lower.tail = FALSE))
  
  model_lm <- lm(YY ~ XX, Data)
  est_lm <- as.numeric(model_lm$coefficients)[2]
  pvalue_lm <- as.numeric(summary(model_lm)$coeffi[,4][2])
  
  return(c(
    c1 = colnames(A_fbm)[a_idx],
    c2 = colnames(B_fbm)[b_idx],
    lm.estimator = est_lm, 
    lm.pvalue = pvalue_lm, 
    Proposed_estimator = Proposed_estimator,
    Proposed_pvalue = Proposed_pvalue
  ))
}

步骤3：无expand.grid的并行任务调度 #

通过数学映射反向计算每个任务对应的变量索引，完全避免生成超大组合表的内存开销，配合bigstatsr原生并行框架实现多进程计算：

# 设置并行核心数，根据机器配置调整
ncores <- nb_cores()

# 分块大小，每次处理1000个任务，避免内存溢出
block_size <- 1000
n_blocks <- ceiling(total_tasks / block_size)

# 并行计算
output_list <- big_parallelize(
  X = test_A_fbm,
  p.FUN = function(X, ind, A_fbm, B_fbm, n_A, block_size) {
    start <- (ind[1] - 1) * block_size + 1
    end <- min(ind[length(ind)] * block_size, n_A * ncol(B_fbm))
    res <- vector("list", end - start + 1)
    k <- 1
    for (task_id in start:end) {
      # 反向映射得到变量索引，无需expand.grid
      a_idx <- (task_id - 1) %% n_A + 1
      b_idx <- (task_id - 1) %/% n_A + 1
      res[[k]] <- test_function_fbm(A_fbm, B_fbm, a_idx, b_idx)
      k <- k + 1
    }
    return(do.call(rbind, res))
  },
  p.combine = "rbind",
  ncores = ncores,
  nblocks = n_blocks,
  A_fbm = test_A_fbm,
  B_fbm = test_B_fbm,
  n_A = n_A,
  block_size = block_size
)

# 转换为最终结果格式
output_final <- as.data.frame(output_list, stringsAsFactors = FALSE)
num_cols <- c("lm.estimator", "lm.pvalue", "Proposed_estimator", "Proposed_pvalue")
output_final[num_cols] <- lapply(output_final[num_cols], as.numeric)

核心优化说明 #

• 完全规避expand.grid生成超大组合表的内存开销，通过数学映射生成任务索引，内存占用恒定
• FBM采用共享内存映射，所有并行进程共享同一份数据，不会每个进程重复复制数据集，内存开销降低数个数量级
• 分块计算+原生并行框架，调度开销低，运行速度随核心数线性提升

内容的提问来源于stack exchange，提问作者user8129108