使用foreach并行化posterior_epred时，如何导出BRMS模型中的自定义函数？

问题根源 #

BRMS模型创建时，公式中的自定义函数（比如scale01）会绑定到函数定义的环境（通常是主环境）。但并行集群的每个节点拥有独立的工作环境，默认不会加载主环境中的自定义函数；即便在foreach循环内重新定义函数，BRMS调用公式时仍会查找模型创建时的环境，最终导致函数找不到的错误。

解决方案 #

方法1：将自定义函数导出到集群节点

使用clusterExport()把scale01函数传递给所有并行节点，确保每个节点的环境中都存在该函数：

test_data <- data.frame(category=c(rep("a",10),
                                   rep("b",10),
                                   rep("c",10)),
                        var1 = runif(30,min=0,max=20))

scale01 <- function(x) (x-min(x))/(max(x)-min(x))

test_brm <- brm(category ~ scale01(var1),
                family=categorical(refcat="a"),
                data=test_data)

prediction_data <- data.frame(id=rep(c(1:500),2),
                              var1=runif(1000,min=0,max=20))

cl <- makeCluster(parallel::detectCores())
registerDoSNOW(cl)
# 将scale01函数导出到所有集群节点
clusterExport(cl, "scale01")

# 优化：用unique(id)避免重复处理同一组数据
preds <- foreach(id = unique(prediction_data$id), .combine=rbind,
                 packages=c("brms","matrixStats")) %dopar% {
                   colMeans(posterior_epred(test_brm,
                                            newdata=prediction_data[which(prediction_data$id==id),]))
                 }

stopCluster(cl)

方法2：使用foreach的.export参数

直接在foreach调用中指定要导出的函数，无需额外调用clusterExport()：

cl <- makeCluster(parallel::detectCores())
registerDoSNOW(cl)

preds <- foreach(id = unique(prediction_data$id), .combine=rbind,
                 packages=c("brms","matrixStats"),
                 .export = "scale01") %dopar% {
                   colMeans(posterior_epred(test_brm,
                                            newdata=prediction_data[which(prediction_data$id==id),]))
                 }

stopCluster(cl)

方法3：提前缩放数据（推荐）

避免在BRMS公式中使用自定义函数，提前对训练数据和预测数据完成缩放，从根源上解决环境依赖问题：

test_data <- data.frame(category=c(rep("a",10),
                                   rep("b",10),
                                   rep("c",10)),
                        var1 = runif(30,min=0,max=20))

# 提取训练数据的缩放参数（基于训练集的min/max，而非新数据）
train_min <- min(test_data$var1)
train_max <- max(test_data$var1)

# 提前缩放训练数据
test_data$var1_scaled <- (test_data$var1 - train_min)/(train_max - train_min)

# 用缩放后的变量训练模型
test_brm <- brm(category ~ var1_scaled,
                family=categorical(refcat="a"),
                data=test_data)

prediction_data <- data.frame(id=rep(c(1:500),2),
                              var1=runif(1000,min=0,max=20))
# 提前缩放预测数据
prediction_data$var1_scaled <- (prediction_data$var1 - train_min)/(train_max - train_min)

cl <- makeCluster(parallel::detectCores())
registerDoSNOW(cl)

preds <- foreach(id = unique(prediction_data$id), .combine=rbind,
                 packages=c("brms","matrixStats")) %dopar% {
                   colMeans(posterior_epred(test_brm,
                                            newdata=prediction_data[which(prediction_data$id==id),]))
                 }

stopCluster(cl)

注意事项 #

• 原代码中foreach(id = prediction_data$id)会重复处理同一个id多次，改用unique(prediction_data$id)可以大幅提升效率。
• 缩放函数的min/max应基于训练数据而非新数据，这是标准化的正确做法，避免数据泄露。

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