使用terra+randomForest二分类时如何仅预测单一类别概率

问题 #

我正在为数百个物种构建物种分布模型，使用包含45个特征、总计约110GiB的VRT栅格栈。这些模型预测适宜栖息地的概率（0-1尺度，<0.5为不适宜，>0.5为适宜），每个物种的randomForest模型仅使用不到20个特征。目前将预测概率写入.tif文件需4-9小时，取决于数据范围和特征数量。

当前直接写入磁盘的代码：

terra::predict(
   x, y, cpkgs="randomForest", type = 'prob', ncores = parallel::detectCores(), filename = 'example.tif', ...)

我希望仅预测单一类别的概率，避免默认输出两类概率的冗余。目前通过先预测再子集化的方式获取目标概率：

rfp <- terra::subset( 
         terra::predict(logo, rfm2, cores=2, type="prob", cpkgs="randomForest"),
 2)
# terra::writeRaster(rfp)

但想知道是否可以通过仅返回单一类别概率，来加快内存预测或直接写入磁盘的速度？考虑到数据量较大，即使计算两类概率很快，也可能带来显著差异。

（附：猜测更严重的瓶颈是VRT数据从HDD读取到内存，目前正在迁移到SSD。）

最小可复现示例：

library(terra)
library(randomForest)

logo <- rast(system.file("ex/logo.tif", package="terra"))   
names(logo) <- c("red", "green", "blue")
p <- matrix(c(48, 48, 48, 53, 50, 46, 54, 70, 84, 85, 74, 84, 95, 85, 
              66, 42, 26, 4, 19, 17, 7, 14, 26, 29, 39, 45, 51, 56, 46, 38, 31, 
              22, 34, 60, 70, 73, 63, 46, 43, 28), ncol=2)

a <- matrix(c(22, 33, 64, 85, 92, 94, 59, 27, 30, 64, 60, 33, 31, 9,
              99, 67, 15, 5, 4, 30, 8, 37, 42, 27, 19, 69, 60, 73, 3, 5, 21,
              37, 52, 70, 74, 9, 13, 4, 17, 47), ncol=2)

xy <- rbind(cbind(1, p), cbind(0, a))

# extract predictor values for points
e <- extract(logo, xy[,2:3])

# combine with response (excluding the ID column)
v <- data.frame(cbind(pa=xy[,1], e))

rm(p, a, xy, e)

### with two output variables (probabilities for each class)
v$pa <- as.factor(v$pa)
rfm2 <- randomForest(formula=pa~., data=v)
rfp <- predict(logo, rfm2, cores=2, type="prob", cpkgs="randomForest")

plot(rfp) # showing probabilities for both classes
plot(sum(rfp)) # obviously in binary classification the probabilities sum to 1
plot(subset(rfp, 2)) # interested in only one of the probabilities!!

解决方案 #

1. 自定义预测函数，仅返回目标类别概率 #

通过包装randomForest的预测逻辑，让terra::predict直接输出单一类别概率，避免生成冗余波段，减少内存占用和写入时间。

示例代码：

# 自定义预测函数，仅返回目标类别概率（此处取第2列对应类别1的概率）
custom_predict <- function(model, data) {
  pred_prob <- predict(model, data, type = "prob")
  return(data.frame(pred = pred_prob[, 2]))
}

# 使用自定义函数预测，直接生成单波段结果并写入磁盘
rfp_single <- terra::predict(
  logo, 
  rfm2, 
  fun = custom_predict,
  cores = 2, 
  cpkgs = "randomForest",
  filename = "single_class_prob.tif",
  overwrite = TRUE
)

plot(rfp_single)

2. 预测时直接指定输出目标波段 #

如果不想自定义函数，可在terra::predict中通过layers参数直接指定仅写入目标波段，跳过先生成双波段再子集化的步骤：

terra::predict(
  logo, 
  rfm2, 
  type = "prob", 
  cores = 2, 
  cpkgs = "randomForest",
  filename = "single_class_prob.tif",
  layers = 2,  # 仅写入第2个波段
  overwrite = TRUE
)

额外优化建议 #

• 迁移到SSD确实能显著提升VRT栅格的读取速度，这是缓解当前瓶颈的核心方案之一。
• 可尝试调整terra::predict的window参数，采用分块处理进一步降低内存压力，尤其针对超大规模栅格。

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