diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala index ff7e449653c03..05e3d0ab26657 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala @@ -369,7 +369,7 @@ class GBTClassificationModel private[ml]( }).apply(features) } else { udf((features: Vector) => { - val margin = TreeEnsembleModel.predict(features, localRootNodes, localTreeWeights) + val margin = TreeEnsembleModel.predictRaw(features, localRootNodes, localTreeWeights) if (margin > 0.0) 1.0 else 0.0 }).apply(features) } @@ -380,13 +380,13 @@ class GBTClassificationModel private[ml]( if (isDefined(thresholds)) { super.predict(features) } else { - if (margin(features) > 0.0) 1.0 else 0.0 + if (TreeEnsembleModel.predictRaw(features, _trees, _treeWeights) > 0.0) 1.0 else 0.0 } } @Since("3.0.0") override def predictRaw(features: Vector): Vector = { - val prediction: Double = margin(features) + val prediction = TreeEnsembleModel.predictRaw(features, _trees, _treeWeights) Vectors.dense(Array(-prediction, prediction)) } @@ -428,17 +428,6 @@ class GBTClassificationModel private[ml]( lazy val featureImportances: Vector = TreeEnsembleModel.featureImportances(trees, numFeatures, perTreeNormalization = false) - /** Raw prediction for the positive class. */ - private def margin(features: Vector): Double = { - var prediction = 0.0 - var i = 0 - while (i < _trees.length) { - prediction += _trees(i).rootNode.predictImpl(features).prediction * _treeWeights(i) - i += 1 - } - prediction - } - /** (private[ml]) Convert to a model in the old API */ private[ml] def toOld: OldGBTModel = { new OldGBTModel(OldAlgo.Classification, _trees.map(_.toOld), _treeWeights) @@ -470,7 +459,7 @@ object GBTClassificationModel extends MLReadable[GBTClassificationModel] { features: Vector, rootNodes: Array[Node], treeWeights: Array[Double]): Vector = { - val prediction = TreeEnsembleModel.predict(features, rootNodes, treeWeights) + val prediction = TreeEnsembleModel.predictRaw(features, rootNodes, treeWeights) Vectors.dense(-prediction, prediction) } diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala index cf990923bc7fd..26c5e2fc856dd 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala @@ -372,24 +372,8 @@ class RandomForestClassificationModel private[ml] ( } @Since("3.0.0") - override def predictRaw(features: Vector): Vector = { - // TODO: When we add a generic Bagging class, handle transform there: SPARK-7128 - // Classifies using majority votes. - // Ignore the tree weights since all are 1.0 for now. - val votes = Array.ofDim[Double](numClasses) - _trees.foreach { tree => - val classCounts = tree.rootNode.predictImpl(features).impurityStats.stats - val total = classCounts.sum - if (total != 0) { - var i = 0 - while (i < numClasses) { - votes(i) += classCounts(i) / total - i += 1 - } - } - } - Vectors.dense(votes) - } + override def predictRaw(features: Vector): Vector = + RandomForestClassificationModel.predictRaw(features, _trees, numClasses) override protected def raw2probabilityInPlace(rawPrediction: Vector): Vector = { rawPrediction match { @@ -469,13 +453,29 @@ class RandomForestClassificationModel private[ml] ( @Since("2.0.0") object RandomForestClassificationModel extends MLReadable[RandomForestClassificationModel] { + private def predictRaw( + features: Vector, + trees: Array[DecisionTreeClassificationModel], + numClasses: Int): Vector = { + val votes = Array.ofDim[Double](numClasses) + trees.foreach { tree => + val classCounts = tree.rootNode.predictImpl(features).impurityStats.stats + val total = classCounts.sum + if (total != 0) { + var i = 0 + while (i < numClasses) { + votes(i) += classCounts(i) / total + i += 1 + } + } + } + Vectors.dense(votes) + } + private def predictRaw( features: Vector, rootNodes: Array[Node], numClasses: Int): Vector = { - // TODO: When we add a generic Bagging class, handle transform there: SPARK-7128 - // Classifies using majority votes. - // Ignore the tree weights since all are 1.0 for now. val votes = Array.ofDim[Double](numClasses) rootNodes.foreach { rootNode => val classCounts = rootNode.predictImpl(features).impurityStats.stats diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala index 069cf39f19ada..c33ce2552b09a 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala @@ -280,7 +280,7 @@ class GBTRegressionModel private[ml]( if ($(predictionCol).nonEmpty) { val predUDF = udf { features: Vector => val (rootNodes, treeWeights) = bcTreeData.value - TreeEnsembleModel.predict(features, rootNodes, treeWeights) + TreeEnsembleModel.predictRaw(features, rootNodes, treeWeights) } predColNames :+= $(predictionCol) predCols :+= predUDF(col($(featuresCol))) @@ -304,17 +304,8 @@ class GBTRegressionModel private[ml]( } } - override def predict(features: Vector): Double = { - // TODO: When we add a generic Boosting class, handle transform there? SPARK-7129 - // Classifies by thresholding sum of weighted tree predictions - var prediction = 0.0 - var i = 0 - while (i < _trees.length) { - prediction += _trees(i).rootNode.predictImpl(features).prediction * _treeWeights(i) - i += 1 - } - prediction - } + override def predict(features: Vector): Double = + TreeEnsembleModel.predictRaw(features, _trees, _treeWeights) @Since("1.4.0") override def copy(extra: ParamMap): GBTRegressionModel = { diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala index 4d9a039ecaca9..d0c342b450202 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala @@ -262,12 +262,8 @@ class RandomForestRegressionModel private[ml] ( } } - override def predict(features: Vector): Double = { - // TODO: When we add a generic Bagging class, handle transform there. SPARK-7128 - // Predict average of tree predictions. - // Ignore the weights since all are 1.0 for now. - _trees.map(_.rootNode.predictImpl(features).prediction).sum / getNumTrees - } + override def predict(features: Vector): Double = + TreeEnsembleModel.predictRaw(features, _trees) / getNumTrees @Since("1.4.0") override def copy(extra: ParamMap): RandomForestRegressionModel = { diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala index d436fc896382b..f970726760b26 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala @@ -169,7 +169,32 @@ private[spark] trait TreeEnsembleModel[M <: DecisionTreeModel] { private[ml] object TreeEnsembleModel { - private[ml] def predict( + private[ml] def predictRaw[M <: DecisionTreeModel]( + features: Vector, + trees: Array[M], + treeWeights: Array[Double]): Double = { + var prediction = 0.0 + var i = 0 + while (i < trees.length) { + prediction += trees(i).rootNode.predictImpl(features).prediction * treeWeights(i) + i += 1 + } + prediction + } + + private[ml] def predictRaw[M <: DecisionTreeModel]( + features: Vector, + trees: Array[M]): Double = { + var prediction = 0.0 + var i = 0 + while (i < trees.length) { + prediction += trees(i).rootNode.predictImpl(features).prediction + i += 1 + } + prediction + } + + private[ml] def predictRaw( features: Vector, rootNodes: Array[Node], treeWeights: Array[Double]): Double = {