Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[jvm-packages] Avoid loosing precision when computing probabilities by converting to Double early (#3576)

Browse Source
This commit is contained in:
Matthew Tovbin 2018-08-13 14:05:07 -07:00 committed by Nan Zhu
parent ce0f0568a6
commit 2b7a1c5780
1 changed files with 5 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostClassifier.scala
View File

@ -411,20 +411,15 @@ class XGBoostClassificationModel private[ml](
}
val probabilityUDF = udf { probability: mutable.WrappedArray[Float] =>
if (numClasses == 2) {
Vectors.dense(Array(1 - probability(0), probability(0)).map(_.toDouble))
} else {
Vectors.dense(probability.map(_.toDouble).toArray)
}
val prob = probability.map(_.toDouble).toArray
val probabilities = if (numClasses == 2) Array(1.0 - prob(0), prob(0)) else prob
Vectors.dense(probabilities)
}
val predictUDF = udf { probability: mutable.WrappedArray[Float] =>
// From XGBoost probability to MLlib prediction
val probabilities = if (numClasses == 2) {
Array(1 - probability(0), probability(0)).map(_.toDouble)
} else {
probability.map(_.toDouble).toArray
}
val prob = probability.map(_.toDouble).toArray
val probabilities = if (numClasses == 2) Array(1.0 - prob(0), prob(0)) else prob
probability2prediction(Vectors.dense(probabilities))
}