能否在mlr3pipelines中基于分组执行特征缩放操作？

问题描述 #

能否基于分组应用mlr3管道处理？例如，mlr3pipelines文档中提供了基础的变量缩放代码：

library(mlr3)
library(mlr3pipelines)
task = tsk("iris")
pop = po("scalemaxabs")
pop$train(list(task))[[1]]$data()

我们给鸢尾花数据添加month分组列后：

library(mlr3)
library(mlr3pipelines)
task = tsk("iris")
dt = task$data()
dt[, month := c(rep(1, 50), rep(2, 50), rep(3, 50))]
task = as_task_classif(dt, target = "Species", id = "iris")

希望按month列对预测变量单独缩放（类似以下data.table的实现逻辑）：

task$data()[, lapply(.SD, function(x) as.vector(scale(x))), .SDcols = names(dt)[2:5], by = month]

是否能在mlr3管道图中完成该操作？

解决方案 #

完全可以，以下两种方法可以实现分组缩放需求：

方法1：使用内置po("group")处理器 #

po("group")可以按指定列拆分任务，对每个分组的子任务应用后续管道操作，最后自动合并回原任务结构。示例代码如下：

library(mlr3)
library(mlr3pipelines)

# 创建带分组列的分类任务
task = tsk("iris")
dt = task$data()
dt[, month := c(rep(1, 50), rep(2, 50), rep(3, 50))]
task = as_task_classif(dt, target = "Species", id = "iris")

# 构建分组缩放管道：拆分分组 → 缩放 → 合并分组
pipe = po("group", group_col = "month") %>>% 
  po("scalemaxabs") %>>% 
  po("ungroup")

# 训练管道并获取处理后的任务
processed_task = pipe$train(list(task))[[1]]
processed_task$data()

方法2：自定义分组缩放PipeOp #

如果需要更灵活的缩放逻辑（比如使用scale()进行标准化而非scalemaxabs），可以自定义PipeOp结合data.table的分组操作：

library(mlr3)
library(mlr3pipelines)
library(data.table)

# 创建带分组列的分类任务
task = tsk("iris")
dt = task$data()
dt[, month := c(rep(1, 50), rep(2, 50), rep(3, 50))]
task = as_task_classif(dt, target = "Species", id = "iris")

# 自定义分组缩放PipeOp
PipeOpGroupScale = R6::R6Class("PipeOpGroupScale",
  inherit = PipeOpTaskPreproc,
  public = list(
    initialize = function(id = "group_scale", param_vals = list()) {
      super$initialize(id, param_vals = param_vals)
    }
  ),
  private = list(
    .train_task = function(task) {
      dt = task$data()
      # 按month分组对所有特征列执行标准化
      scaled_features = dt[, lapply(.SD, function(x) as.vector(scale(x))), 
                           .SDcols = task$feature_names, by = month]
      # 合并目标列与处理后的特征列
      full_dt = merge(scaled_features, dt[, .(get(task$target_names), month)], 
                      by = c("month", task$feature_names), all = TRUE)
      setnames(full_dt, "V1", task$target_names)
      # 转换回分类任务
      as_task_classif(full_dt, target = task$target_names, id = task$id)
    },
    .predict_task = function(task) {
      # 预测阶段复用训练阶段的分组缩放逻辑
      dt = task$data()
      scaled_features = dt[, lapply(.SD, function(x) as.vector(scale(x))), 
                           .SDcols = task$feature_names, by = month]
      full_dt = merge(scaled_features, dt[, .(get(task$target_names), month)], 
                      by = c("month", task$feature_names), all = TRUE)
      setnames(full_dt, "V1", task$target_names)
      as_task_classif(full_dt, target = task$target_names, id = task$id)
    }
  )
)

# 实例化自定义PipeOp并运行
po_group_scale = PipeOpGroupScale$new()
processed_task = po_group_scale$train(list(task))[[1]]
processed_task$data()

说明：方法1适合使用内置处理器的标准化场景，代码简洁；方法2适合需要自定义缩放逻辑的场景，灵活性更高。

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