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

Make sure 'thresholds' are considered when executing predict method (#3577)

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

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

@ -251,11 +251,11 @@ class XGBoostClassificationModel private[ml](
override def predict(features: Vector): Double = {
import DataUtils._
val dm = new DMatrix(XGBoost.removeMissingValues(Iterator(features.asXGB), $(missing)))
val probability = _booster.predict(data = dm)(0)
val probability = _booster.predict(data = dm)(0).map(_.toDouble)
if (numClasses == 2) {
math.round(probability(0))
} else {
Vectors.dense(probability.map(_.toDouble)).argmax
probability2prediction(Vectors.dense(probability))
}
}