xgboost4j - spark评估需要RDD [(Double,Double)]

时间:2016-11-10 08:08:13

标签: scala apache-spark prediction xgboost

我尝试将xgboost4j与spark 2.0.1和Dataset API一起使用。到目前为止,我使用model.transform(testData)

以下列格式获得了预测
predictions.printSchema
root
 |-- label: double (nullable = true)
 |-- features: vector (nullable = true)
 |-- probabilities: vector (nullable = true)
 |-- prediction: double (nullable = true)


+-----+--------------------+--------------------+----------+
|label|            features|       probabilities|prediction|
+-----+--------------------+--------------------+----------+
|  0.0|[0.0,1.0,0.0,476....|[0.96766251325607...|       0.0|
|  0.0|[0.0,1.0,0.0,642....|[0.99599152803421...|       0.0|

但现在我想生成评估指标。如何将预测映射到正确的格式? XGBoost-4j by DMLC on Spark-1.6.1提出了类似的问题,但我无法让它为我工作。

val metrics = new BinaryClassificationMetrics(predictions.select("prediction", "label").rdd)
would require RDD[(Double, Double)] 

而不是predictions.select("prediction", "label"),看起来像

root
 |-- label: double (nullable = true)
 |-- prediction: double (nullable = true)

尝试将其映射到所需的元组,如:

predictions.select("prediction", "label").map{case Row(_) => (_,_)}

无法正常工作。

修改

在sparks文档中阅读更多内容我发现http://spark.apache.org/docs/latest/api/scala/index.html#org.apache.spark.ml.evaluation.BinaryClassificationEvaluator支持ml而不是ml-lib,例如数据集。到目前为止,我无法在管道中成功集成xgboost4j。

1 个答案:

答案 0 :(得分:1)

这是一个很好的例子https://github.com/dmlc/xgboost/blob/master/jvm-packages/xgboost4j-example/src/main/scala/ml/dmlc/xgboost4j/scala/example/spark/SparkModelTuningTool.scala如何在spark管道中使用xgboost4j。实际上,他们有一个XGBoostEstimator,可以在管道中很好地运行。