XGBoost4J-Spark 教程

XGBoost4J-Spark 项目旨在通过将 XGBoost 适配到 Apache Spark 的 MLLIB 框架中,实现 XGBoost 与 Apache Spark 的无缝集成。通过这种集成,用户不仅可以使用 XGBoost 高性能的算法实现,还可以利用 Spark 强大的数据处理引擎进行:

  • 特征工程:特征提取、转换、降维和选择等。

  • 流水线(Pipelines):构建、评估和调优机器学习流水线

  • 持久化:持久化和加载机器学习模型,甚至整个流水线

本教程将涵盖使用 XGBoost4J-Spark 构建端到端机器学习流水线的完整过程。我们将讨论:

  • 使用 Spark 预处理数据以适配 XGBoost4J-Spark 的数据接口

  • 使用 XGBoost4J-Spark 训练 XGBoost 模型

  • 使用 Spark 进行 XGBoost 模型服务化(预测)

  • 使用 XGBoost4J-Spark 构建机器学习流水线

  • 在生产环境中运行 XGBoost4J-Spark

使用 XGBoost4J-Spark 构建机器学习应用

参考 XGBoost4J-Spark 依赖

在开始了解如何使用 XGBoost4J-Spark 之前,您应该首先参考 从 Maven 仓库安装,以便将 XGBoost4J-Spark 添加为项目的依赖。我们同时提供稳定版本和快照版本。

注意

XGBoost4J-Spark 需要 Apache Spark 3.0+

XGBoost4J-Spark 现在需要 Apache Spark 3.0+。最新版本的 XGBoost4J-Spark 大量使用了 org.apache.spark.ml.param.shared 的功能,以提供与 Spark MLLIB 框架的紧密集成,而这些功能在较早版本的 Spark 中并不完全可用。

此外,请确保直接从 Apache 官方网站 安装 Spark。上游 XGBoost 不保证能与第三方分发版的 Spark(如 Cloudera Spark)协同工作。 请咨询相应的第三方以获取其分发版本的 XGBoost。

警告

Spark 4.0.0 与 XGBoost4J-Spark 不兼容

Apache Spark 4.0.0 对 org.apache.spark.ml.param.Param 类构造函数引入了破坏性变更(SPARK-52259),这会在实例化任何 XGBoost 估算器(estimator)时导致 NoSuchMethodError

java.lang.NoSuchMethodError: 'void org.apache.spark.ml.param.Param.<init>(
    org.apache.spark.ml.util.Identifiable, java.lang.String, java.lang.String, scala.Function1)'

此问题影响所有使用 Param 的第三方机器学习库。该问题在 Spark 4.0.1Spark 4.1.0 中已得到修复。如果您使用的是 Spark 4.x,请升级到 Spark 4.0.1 或更高版本

数据准备

如前所述,XGBoost4J-Spark 无缝集成了 Spark 和 XGBoost。通过这种集成,用户可以利用 Spark 这一便捷且强大的数据处理框架,对训练/测试数据集应用各种类型的转换。

在本节中,我们以 鸢尾花(Iris) 数据集为例,展示如何使用 Spark 转换原始数据集,使其适配 XGBoost 的数据接口。

Iris 数据集以 CSV 格式提供。每个实例包含 4 个特征:“花萼长度”(sepal length)、“花萼宽度”(sepal width)、“花瓣长度”(petal length)和“花瓣宽度”(petal width)。此外,它还包含 “class” 列,即标签,具有三个可能的值:“Iris Setosa”(山鸢尾)、“Iris Versicolour”(变色鸢尾)和“Iris Virginica”(维吉尼亚鸢尾)。

使用 Spark 内置读取器读取数据集

数据转换的第一步是将数据集加载为 Spark 的结构化数据抽象:DataFrame。

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.types.{DoubleType, StringType, StructField, StructType}

val spark = SparkSession.builder().getOrCreate()
val schema = new StructType(Array(
  StructField("sepal length", DoubleType, true),
  StructField("sepal width", DoubleType, true),
  StructField("petal length", DoubleType, true),
  StructField("petal width", DoubleType, true),
  StructField("class", StringType, true)))
val rawInput = spark.read.schema(schema).csv("input_path")

在第一行中,我们创建了一个 SparkSession 实例,这是所有处理 DataFrame 的 Spark 程序的入口。 schema 变量定义了封装 Iris 数据的 DataFrame 的模式(schema)。通过显式设置该模式,我们可以定义列名及其类型;否则,列名将是 Spark 默认派生的名称,如 _col0 等。最后,我们可以使用 Spark 内置的 CSV 读取器,将 Iris CSV 文件加载为名为 rawInput 的 DataFrame。

Spark 还包含许多用于其他格式的内置读取器。最新版本的 Spark 支持 CSV、JSON、Parquet 和 LIBSVM。

转换原始鸢尾花数据集

为了使 Iris 数据集能够被 XGBoost 识别,我们需要:

  1. 将 String(字符串)类型的标签(即 “class”)转换为 Double 类型的标签。

  2. 将特征列组合成一个向量(Vector),以适配 Spark ML 框架的数据接口。

要将 String 类型的标签转换为 Double,我们可以使用 Spark 内置的特征转换器 StringIndexer

import org.apache.spark.ml.feature.StringIndexer
val stringIndexer = new StringIndexer().
  setInputCol("class").
  setOutputCol("classIndex").
  fit(rawInput)
val labelTransformed = stringIndexer.transform(rawInput).drop("class")

通过新创建的 StringIndexer 实例:

  1. 我们设置输入列,即包含 String 类型标签的列。

  2. 我们设置输出列,即包含转换后 Double 类型标签的列。

  3. 然后,我们用输入 DataFrame rawInputfit StringIndexer,以便 Spark 内部能够获取诸如非重复值总数等信息。

现在我们有了一个准备好应用到输入 DataFrame 的 StringIndexer。为了执行 StringIndexer 的转换逻辑,我们对输入 DataFrame rawInput 进行 transform 转换。为了保持 DataFrame 的简洁,我们丢弃 “class” 列,仅保留特征列和转换后的 Double 类型标签列(如上述代码片段的最后一行所示)。

fittransform 是 MLLIB 中的两个核心操作。简单来说,fit 会生成一个“转换器”(transformer,例如 StringIndexer),而每个转换器通过对 DataFrame 应用 transform 方法来添加包含已转换特征/标签或预测结果的新列。要深入了解 fittransform,您可以在此处找到更多详细信息。

类似地,我们可以使用另一个转换器 VectorAssembler,将特征列 “sepal length”、“sepal width”、“petal length” 和 “petal width” 组合成一个向量。

import org.apache.spark.ml.feature.VectorAssembler
val vectorAssembler = new VectorAssembler().
  setInputCols(Array("sepal length", "sepal width", "petal length", "petal width")).
  setOutputCol("features")
val xgbInput = vectorAssembler.transform(labelTransformed).select("features", "classIndex")

现在,我们得到了一个仅包含两列的 DataFrame:“features” 列包含以向量形式表示的 “sepal length”、“sepal width”、“petal length” 和 “petal width”;“classIndex” 列则包含 Double 类型的标签。像这样(包含向量特征和数值标签)的 DataFrame 可以直接输入到 XGBoost4J-Spark 的训练引擎中。

处理缺失值

XGBoost 默认支持缺失值(如这里所述)。如果输入是 SparseVector(稀疏向量),XGBoost 会将 SparseVector 中缺失的所有值视为缺失值。您还可以指定 XGBoost 将数据集中的某个特定值视为缺失值。默认情况下,XGBoost 将 NaN 视为代表缺失的值。

将缺失值(例如 -999)设置给 XGBoostClassifier 中的 “missing” 参数的示例

import ml.dmlc.xgboost4j.scala.spark.XGBoostClassifier
val xgbParam = Map("eta" -> 0.1f,
      "missing" -> -999,
      "objective" -> "multi:softprob",
      "num_class" -> 3,
      "num_round" -> 100,
      "num_workers" -> 2)
val xgbClassifier = new XGBoostClassifier(xgbParam).
      setFeaturesCol("features").
      setLabelCol("classIndex")

注意

缺失值

如果特征是向量类型,单特征实例可能是一个 SparseVector,其中 “0” 会被视作缺失值。为了获得正确的模型,XGBoost4j-Spark 会通过还原 “0” 将 SparseVector 转换为数组。然而,我们不能将 0 默认为缺失值,因为 0 可能具有实际意义。因此在这种情况下,即使 XGBoost4j-Spark 默认将缺失值设置为 Float.NaN,用户也需要显式指定缺失值。

训练

XGBoost 支持回归、分类和排序。虽然我们在本教程中使用 Iris 数据集来展示如何使用 XGBoost4J-Spark 解决多分类问题,但在回归和排序中的用法与分类非常相似。

要训练一个用于分类的 XGBoost 模型,我们首先需要创建一个 XGBoostClassifier

import ml.dmlc.xgboost4j.scala.spark.XGBoostClassifier
val xgbParam = Map("eta" -> 0.1f,
      "max_depth" -> 2,
      "objective" -> "multi:softprob",
      "num_class" -> 3)
val xgbClassifier = new XGBoostClassifier(xgbParam).
      setNumRound(100).
      setNumWorkers(2).
      setFeaturesCol("features").
      setLabelCol("classIndex")

训练 XGBoost 模型时可用的参数可以在这里找到。在 XGBoost4J-Spark 中,我们不仅支持默认的参数集,还支持这些参数的小驼峰命名法(camel-case)变体,以保持与 Spark MLLIB 参数的一致性。

具体来说,此页面中的每个参数在 XGBoost4J-Spark 中都有其对应的驼峰命名形式。例如,要设置每棵树的 max_depth,您可以像我们在上面的代码片段中所做的那样传递参数(将 max_depth 包装在 Map 中),或者您可以通过 XGBoostClassifier 中的 setter 方法来进行设置

val xgbClassifier = new XGBoostClassifier().
  setFeaturesCol("features").
  setLabelCol("classIndex")
xgbClassifier.setMaxDepth(2)

在我们设置好 XGBoostClassifier 参数以及特征/标签列之后,我们可以通过用输入 DataFrame 来 fit XGBoostClassifier,从而构建一个转换器:XGBoostClassificationModel。这个 fit 操作本质上就是训练过程,生成的模型随后可用于预测。

val xgbClassificationModel = xgbClassifier.fit(xgbInput)

提前停止

早停(Early stopping)是一项用于防止不必要训练迭代的功能。通过指定 num_early_stopping_rounds 或直接对 XGBoostClassifier 或 XGBoostRegressor 调用 setNumEarlyStoppingRounds,我们可以定义:当评估指标偏离最佳迭代达到多少轮时,提前终止训练迭代。

当涉及自定义评估指标时,除了 num_early_stopping_rounds 外,您还需要定义 maximize_evaluation_metrics 或调用 setMaximizeEvaluationMetrics 来指定您在训练中是想要最大化还是最小化该指标。对于内置的评估指标,XGBoost4J-Spark 会自动选择优化方向。

例如,我们需要最大化评估指标(将 maximize_evaluation_metrics 设置为 true),并将 num_early_stopping_rounds 设置为 5。如果第 10 次迭代的评估指标是目前为止最大的。在随后的迭代中,如果没有评估指标超过第 10 次迭代(最佳迭代)的值,训练将在第 15 次迭代时提前停止。

使用评估数据集进行训练

您还可以在训练期间使用评估数据集监控模型的性能。只需在 XGBoostClassifier、XGBoostRegressor 或 XGBoostRanker 上调用 setEvalDataset 即可。

预测

XGBoost4j-Spark 支持两种模型服务化方式:批量预测和单实例预测。

批量预测

当我们获得模型(XGBoostClassificationModel、XGBoostRegressionModel 或 XGBoostRankerModel)后,它会接收一个 DataFrame,读取包含特征向量的列,对每个特征向量进行预测,并在默认情况下输出一个包含以下列的新 DataFrame:

  • XGBoostClassificationModel 将输出每个可能标签的边际值(margins,对应 rawPredictionCol)、概率(probabilities,对应 probabilityCol)和最终预测标签(predictionCol)。

  • XGBoostRegressionModel 将输出预测标签(predictionCol)。

  • XGBoostRankerModel 将输出预测标签(predictionCol)。

批量预测要求用户以 DataFrame 的形式传入测试集。XGBoost4J-Spark 为 DataFrame 的每个分区(partition)启动一个 XGBoost 工作器(worker)进行并行预测,并以批量方式生成整个 DataFrame 的预测结果。

val xgbClassificationModel = xgbClassifier.fit(xgbInput)
val results = xgbClassificationModel.transform(testSet)

通过上述代码片段,我们得到了一个名为 result 的结果 DataFrame,其中包含了边际值、每个类别的概率以及每个实例的预测结果。

+-----------------+----------+--------------------+--------------------+----------+
|         features|classIndex|       rawPrediction|         probability|prediction|
+-----------------+----------+--------------------+--------------------+----------+
|[5.1,3.5,1.4,0.2]|       0.0|[3.45569849014282...|[0.99579632282257...|       0.0|
|[4.9,3.0,1.4,0.2]|       0.0|[3.45569849014282...|[0.99618089199066...|       0.0|
|[4.7,3.2,1.3,0.2]|       0.0|[3.45569849014282...|[0.99643349647521...|       0.0|
|[4.6,3.1,1.5,0.2]|       0.0|[3.45569849014282...|[0.99636095762252...|       0.0|
|[5.0,3.6,1.4,0.2]|       0.0|[3.45569849014282...|[0.99579632282257...|       0.0|
|[5.4,3.9,1.7,0.4]|       0.0|[3.45569849014282...|[0.99428516626358...|       0.0|
|[4.6,3.4,1.4,0.3]|       0.0|[3.45569849014282...|[0.99643349647521...|       0.0|
|[5.0,3.4,1.5,0.2]|       0.0|[3.45569849014282...|[0.99579632282257...|       0.0|
|[4.4,2.9,1.4,0.2]|       0.0|[3.45569849014282...|[0.99618089199066...|       0.0|
|[4.9,3.1,1.5,0.1]|       0.0|[3.45569849014282...|[0.99636095762252...|       0.0|
|[5.4,3.7,1.5,0.2]|       0.0|[3.45569849014282...|[0.99428516626358...|       0.0|
|[4.8,3.4,1.6,0.2]|       0.0|[3.45569849014282...|[0.99643349647521...|       0.0|
|[4.8,3.0,1.4,0.1]|       0.0|[3.45569849014282...|[0.99618089199066...|       0.0|
|[4.3,3.0,1.1,0.1]|       0.0|[3.45569849014282...|[0.99618089199066...|       0.0|
|[5.8,4.0,1.2,0.2]|       0.0|[3.45569849014282...|[0.97809928655624...|       0.0|
|[5.7,4.4,1.5,0.4]|       0.0|[3.45569849014282...|[0.97809928655624...|       0.0|
|[5.4,3.9,1.3,0.4]|       0.0|[3.45569849014282...|[0.99428516626358...|       0.0|
|[5.1,3.5,1.4,0.3]|       0.0|[3.45569849014282...|[0.99579632282257...|       0.0|
|[5.7,3.8,1.7,0.3]|       0.0|[3.45569849014282...|[0.97809928655624...|       0.0|
|[5.1,3.8,1.5,0.3]|       0.0|[3.45569849014282...|[0.99579632282257...|       0.0|
+-----------------+----------+--------------------+--------------------+----------+

单实例预测

XGBoostClassificationModel、XGBoostRegressionModel 或 XGBoostRankerModel 也支持对单实例进行预测。它接受单个 Vector 作为特征,并输出预测标签。

但是,由于 XGBoost 内部开销,单实例预测的开销较高,请谨慎使用!

val features = xgbInput.head().getAs[Vector]("features")
val result = xgbClassificationModel.predict(features)

模型持久化

模型和流水线持久化

数据科学家产生一个机器学习模型,并将其交给工程团队部署到生产环境中。反之,一个训练好的模型也可能被数据科学家在数据探索过程中使用(例如作为基准线)。因此,支持模型持久化至关重要,这可以使模型能够跨使用场景和编程语言进行复用。

XGBoost4j-Spark 支持将 XGBoostClassifier/XGBoostClassificationModel、XGBoostRegressor/XGBoostRegressionModel 以及 XGBoostRanker/XGBoostRankerModel 保存到文件系统或从中加载。它还支持保存和加载包含这些估算器(estimator)和模型的机器学习流水线(ML pipeline)。

我们可以将 XGBoostClassificationModel 保存到文件系统

val xgbClassificationModelPath = "/tmp/xgbClassificationModel"
xgbClassificationModel.write.overwrite().save(xgbClassificationModelPath)

然后可以在另一个会话中加载该模型

import ml.dmlc.xgboost4j.scala.spark.XGBoostClassificationModel

val xgbClassificationModel2 = XGBoostClassificationModel.load(xgbClassificationModelPath)
xgbClassificationModel2.transform(xgbInput)

注意

除了将模型导出为原始格式外,用户还可以将模型导出为 JSON 或 UBJ 格式。

val xgbClassificationModelPath = "/tmp/xgbClassificationModel"
xgbClassificationModel.write.overwrite().option("format", "json").save(xgbClassificationModelPath)

关于机器学习流水线的保存和加载,请参阅下一节。

与 XGBoost 的其他语言绑定进行交互

在使用 XGBoost4j-Spark 在海量数据集上训练模型后,有时我们希望在单机上进行模型服务化,或者将其与其他单节点库集成以进行进一步处理。

保存模型后,我们可以直接使用单节点 Python 版 XGBoost 加载此模型。

val xgbClassificationModelPath = "/tmp/xgbClassificationModel"
xgbClassificationModel.write.overwrite().save(xgbClassificationModelPath)
import xgboost as xgb
bst = xgb.Booster({'nthread': 4})
bst.load_model("/tmp/xgbClassificationModel/data/model")

注意

XGBoost4J-Spark 与其他绑定之间的一致性问题

XGBoost4J-Spark 与 XGBoost 的其他语言绑定之间存在一个一致性问题。

当用户使用以下代码片段让 Spark 加载 LIBSVM 格式的训练/测试数据时

spark.read.format("libsvm").load("trainingset_libsvm")

Spark 默认数据集使用的是从 1 开始的索引(即特征索引从 1 开始)。然而,当您使用 XGBoost 的其他绑定(例如 XGBoost 的 Python API)进行预测时,XGBoost 默认假设数据集使用的是从 0 开始的索引(即特征索引从 0 开始)。这给那些使用 Spark 训练模型,却在其他 XGBoost 绑定中用相同格式数据集进行预测的用户带来了隐患。解决方案是在使用例如 Python API 进行预测之前,将数据集转换为从 0 开始的索引,或者在使用 DMatrix 加载时在文件路径后追加 ?indexing_mode=1。例如在 Python 中:

xgb.DMatrix('test.libsvm?indexing_mode=1')

使用 XGBoost4J-Spark 构建机器学习流水线

基础机器学习流水线

Spark 机器学习流水线可以将多个算法或功能组合成单个流水线。它涵盖了从特征提取、转换、选择到模型训练和预测的完整过程。XGBoost4j-Spark 使得将 XGBoost 无缝嵌入到此类流水线中成为可能。以下示例展示了如何构建一个由 Spark MLlib 特征转换器和 XGBoostClassifier 估算器组成的流水线。

我们仍然使用 Iris 数据集和 rawInput DataFrame。首先我们需要将数据集分割为训练集和测试集。

val Array(training, test) = rawInput.randomSplit(Array(0.8, 0.2), 123)

然后我们构建包含 4 个阶段的机器学习流水线:

  • 将所有特征组合成单个向量列。

  • 将 String 类型标签转换为索引化后的 Double 类型标签。

  • 使用 XGBoostClassifier 训练分类模型。

  • 将索引化后的 Double 类型标签转换回原始 String 类型标签。

我们在前面的章节中已经展示了前三个步骤,最后一步是通过一个名为 IndexToString 的新转换器完成的

val labelConverter = new IndexToString()
.setInputCol("prediction")
.setOutputCol("realLabel")
.setLabels(stringIndexer.labels)

我们需要在 Spark ML 框架中将这些步骤组织为一个流水线(Pipeline),并对整个流水线进行评估以获得 PipelineModel

import org.apache.spark.ml.feature._
import org.apache.spark.ml.Pipeline

val pipeline = new Pipeline()
    .setStages(Array(assembler, stringIndexer, booster, labelConverter))
val model = pipeline.fit(training)

获取 PipelineModel 后,我们可以对测试数据集进行预测,并评估模型的准确率。

import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator

val prediction = model.transform(test)
val evaluator = new MulticlassClassificationEvaluator()
val accuracy = evaluator.evaluate(prediction)

带超参数调优的流水线

最大限度发挥 XGBoost 威力的最关键操作是为模型选择最优参数。手动调参是一个繁琐且耗费人力的过程。借助最新版本的 XGBoost4J-Spark,我们可以利用 Spark 的模型选择工具来自动执行此过程。

以下示例展示了利用 CrossValidation(交叉验证)和 MulticlassClassificationEvaluator(多分类评估器)搜索两个 XGBoost 参数 max_deptheta 的最佳组合的代码片段。(请参阅 XGBoost 参数。)由 MulticlassClassificationEvaluator 评估得出准确率最高的模型将被选中,并用于生成测试集的预测结果。

import org.apache.spark.ml.tuning._
import org.apache.spark.ml.PipelineModel
import ml.dmlc.xgboost4j.scala.spark.XGBoostClassificationModel

val paramGrid = new ParamGridBuilder()
    .addGrid(booster.maxDepth, Array(3, 8))
    .addGrid(booster.eta, Array(0.2, 0.6))
    .build()
val cv = new CrossValidator()
    .setEstimator(pipeline)
    .setEvaluator(evaluator)
    .setEstimatorParamMaps(paramGrid)
    .setNumFolds(3)

val cvModel = cv.fit(training)

val bestModel = cvModel.bestModel.asInstanceOf[PipelineModel].stages(2)
    .asInstanceOf[XGBoostClassificationModel]
bestModel.extractParamMap()

在生产环境中运行 XGBoost4J-Spark

XGBoost4J-Spark 是使 XGBoost 能够更轻松地落地生产环境的重要步骤之一。在本节中,我们将介绍在生产环境中运行 XGBoost4J-Spark 的三个关键功能。

并行/分布式训练

庞大的训练数据集是生产环境中最显著的特征之一。为了确保 XGBoost 的训练规模能够随着数据量的大小进行扩展,XGBoost4J-Spark 桥接了 Spark 的分布式/并行处理框架与 XGBoost 的并行/分布式训练机制。

在 XGBoost4J-Spark 中,每个 XGBoost 工作器(worker)都被包装在一个 Spark 任务(task)中,并且 Spark 内存空间中的训练数据集会以对用户透明的方式提供给 XGBoost 工作器。

在构建 XGBoostClassifier 的代码片段中,我们设置了参数 num_workers(或 numWorkers)。该参数控制我们在训练 XGBoostClassificationModel 时想要拥有的并行工作器(worker)数量。

注意

关于 OpenMP 优化

默认情况下,我们为每个 XGBoost 工作器分配一个 CPU 核心。因此,每个 XGBoost 工作器内部的 OpenMP 优化不会生效,训练的并行化是通过同时运行多个工作器(即 Spark 任务)来实现的。

如果您确实需要 OpenMP 优化,则必须:

  1. 在创建 XGBoostClassifier/XGBoostRegressor 时,将 nthread 设置为大于 1 的值

  2. 将 Spark 中的 spark.task.cpus 设置为与 nthread 相同的值

Gang 调度

XGBoost 使用 AllReduce 算法在训练期间同步每个工作器的统计数据(例如直方图值)。因此,XGBoost4J-Spark 要求在训练运行前,所有 nthread * numWorkers 个核心都必须可用。

在许多用户共享同一个集群的生产环境中,很难保证您的 XGBoost4J-Spark 应用程序每次运行都能获得所有请求的资源。默认情况下,当需要更多可用资源时,XGBoost 中的通信层将阻塞整个应用程序。此过程通常会带来不必要的资源浪费,因为它会占用已经准备就绪的资源并尝试申请更多资源。此外,这通常会静默发生,不会引起用户的注意。

XGBoost4J-Spark 允许用户设置一个从集群中申请资源的超时阈值。如果应用程序在此时间段内无法获得足够的资源,则应用程序将报错失败,而不是因长时间挂起而浪费资源。要启用此功能,您可以在 XGBoostClassifier/XGBoostRegressor/XGBoostRanker 中进行设置

xgbClassifier.setRabitTrackerTimeout(60000L)

或者在构建 XGBoostClassifier 时在 xgbParamMap 中传入 rabit_tracker_timeout

val xgbParam = Map("eta" -> 0.1f,
   "max_depth" -> 2,
   "objective" -> "multi:softprob",
   "num_class" -> 3,
   "num_round" -> 100,
   "num_workers" -> 2,
   "rabit_tracker_timeout" -> 60000L)
val xgbClassifier = new XGBoostClassifier(xgbParam).
    setFeaturesCol("features").
    setLabelCol("classIndex")

如果 XGBoost4J-Spark 无法获得运行两个 XGBoost 工作器所需的足够资源,则应用程序将失败。用户可以通过外部机制来监控应用程序的状态并在发生此类情况时获得通知。

训练期间的检查点(Checkpoint)

在生产环境中,临时性故障也屡见不鲜。为了简化 XGBoost 的设计,如果任何一个分布式工作器出现故障,我们就会停止训练。然而,如果在经过长时间训练后由于故障而失败,那将会极大地浪费资源。

我们支持在训练期间创建检查点,以有助于从故障中更高效地恢复。要启用此功能,您可以使用 setCheckpointInterval 设置每隔多少次迭代构建一次检查点,并使用 setCheckpointPath 设置检查点的存储路径

xgbClassifier.setCheckpointInterval(2)
xgbClassifier.setCheckpointPath("/checkpoint_path")

另一种等效的方法是在 XGBoostClassifier 的构造函数中传入参数

val xgbParam = Map("eta" -> 0.1f,
   "max_depth" -> 2,
   "objective" -> "multi:softprob",
   "num_class" -> 3,
   "num_round" -> 100,
   "num_workers" -> 2,
   "checkpoint_path" -> "/checkpoints_path",
   "checkpoint_interval" -> 2)
val xgbClassifier = new XGBoostClassifier(xgbParam).
    setFeaturesCol("features").
    setLabelCol("classIndex")

如果在这 100 轮中训练失败,下一次训练运行将通过读取 /checkpoints_path 中的最新检查点文件,从构建该检查点时的迭代次数开始继续,直到发生下一次故障或达到指定的 100 轮。

外部内存

3.0 版本新增。

警告

该功能处于实验阶段。

这里我们指的是基于迭代器的外部内存,而不是使用特殊 URL 参数的那种。自 3.0 版本起,XGBoost-Spark 对基于 GPU 的外部内存训练提供了实验性支持(请参阅 XGBoost4J-Spark-GPU 教程)。当与基于 GPU 的训练结合使用时,数据首先缓存在磁盘上,然后暂存在 CPU 内存中。关于外部内存训练的通用概念和最佳实践,请参阅 使用 XGBoost 外部内存版本。此外,请参阅估算器参数 useExternalMemory 的文档字符串。结合 Spark 估算器使用

val xgbClassifier = new XGBoostClassifier(xgbParam)
    .setFeaturesCol(featuresNames)
    .setLabelCol(labelName)
    .setUseExternalMemory(true)
    .setDevice("cuda")  // CPU is not yet supported