Spark/Scala环境下基于xgboost4j-spark的XGBoost MultiOutputRegressor咨询

xgboost4j-spark官方没有提供专属的多输出回归封装，但可以通过两种方式实现和你提供的Python sklearn版MultiOutputRegressor等效的功能，完全适配Spark/Scala分布式环境。

方案1：直接使用Spark MLlib自带的MultiOutputRegressor包装 #

Spark MLlib的org.apache.spark.ml.regression.MultiOutputRegressor是通用的多输出包装器，支持包装所有实现了Spark ML Regressor接口的模型，而xgboost4j-spark的XGBoostRegressor完全符合该接口规范，可直接使用。
对应的Scala实现代码如下：

import org.apache.spark.ml.regression.MultiOutputRegressor
import ml.dmlc.xgboost4j.scala.spark.XGBoostRegressor

// 初始化单输出XGBoost回归器，参数和你提供的Python版本完全对齐
val xgbRegressor = new XGBoostRegressor()
  .setMaxDepth(1)
  .setEta(0.01)
  .setSilent(1)
  .setSubsample(0.8)
  .setRegLambda(1.515)
  .setRegAlpha(0.0017)
  .setMinChildWeight(7)
  .setColsampleBytree(0.85)
  .setGamma(0.01)
  .setObjective("reg:linear")
  .setTreeMethod("approx")
  .setBooster("gbtree")
  .setNumRound(100) // 对应Python版的n_estimators参数
  .setFeaturesCol("features") // 特征列，需要提前用VectorAssembler把所有自变量合并为向量列

// 初始化多输出回归器
val multiOutputReg = new MultiOutputRegressor()
  .setRegressor(xgbRegressor)
  .setLabelCol("labels") // 标签列，需要提前把所有因变量合并为向量列

// 训练模型
val model = multiOutputReg.fit(trainDataset)

该方案的特点：

• 实现简单，和sklearn的MultiOutputRegressor逻辑完全一致，每个标签独立训练一个XGBoost模型，天然支持分布式运行
• 缺点是所有标签共享同一套XGBoost超参数，无法为不同标签配置差异化参数

方案2：自定义多输出封装（支持差异化参数） #

如果需要为不同的因变量设置不同的XGBoost超参数，可以手动实现封装逻辑：

// 自定义多输出XGBoost模型类
case class MultiOutputXGBModel(models: List[(String, XGBoostRegressionModel)]) {
  def transform(dataset: Dataset[_]): DataFrame = {
    models.foldLeft(dataset.toDF()) { (df, modelInfo) =>
      val (labelName, model) = modelInfo
      model.transform(df).withColumnRenamed("prediction", s"pred_$labelName")
    }
  }
}

// 训练逻辑
val labelCols = List("dep_1", "dep_2", "dep_3", "dep_4")
val trainedModels = labelCols.map { label =>
  // 可针对每个label单独配置不同的超参数
  val xgb = new XGBoostRegressor()
    .setMaxDepth(1)
    .setEta(0.01)
    // 其余超参数配置同上
    .setLabelCol(label)
    .setFeaturesCol("features")
  (label, xgb.fit(trainDataset))
}

val multiModel = MultiOutputXGBModel(trainedModels)
// 预测时会自动生成所有标签的预测结果列
val predictResult = multiModel.transform(testDataset)

注意事项 #

• 以上两种实现均和sklearn的MultiOutputRegressor逻辑一致，默认假设多个标签之间独立，没有利用标签间的关联关系。如果需要利用标签相关性优化效果，需要自定义XGBoost的多目标损失函数，在xgboost4j-spark中可通过setCustomObj方法传入自定义损失实现。
• 输入数据集需要提前用VectorAssembler将所有自变量合并为features向量列，如果用方案1还需要将所有因变量合并为labels向量列。

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