diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
index 8c4cec132665354108b251cefa484f29e989dd13..7f0397f6bd65a543e129f889a6390ac81e3e40a0 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/classification/DecisionTreeClassifier.scala
@@ -169,6 +169,25 @@ final class DecisionTreeClassificationModel private[ml] (
s"DecisionTreeClassificationModel (uid=$uid) of depth $depth with $numNodes nodes"
}
+ /**
+ * Estimate of the importance of each feature.
+ *
+ * This generalizes the idea of "Gini" importance to other losses,
+ * following the explanation of Gini importance from "Random Forests" documentation
+ * by Leo Breiman and Adele Cutler, and following the implementation from scikit-learn.
+ *
+ * This feature importance is calculated as follows:
+ * - importance(feature j) = sum (over nodes which split on feature j) of the gain,
+ * where gain is scaled by the number of instances passing through node
+ * - Normalize importances for tree to sum to 1.
+ *
+ * Note: Feature importance for single decision trees can have high variance due to
+ * correlated predictor variables. Consider using a [[RandomForestClassifier]]
+ * to determine feature importance instead.
+ */
+ @Since("2.0.0")
+ lazy val featureImportances: Vector = RandomForest.featureImportances(this, numFeatures)
+
/** (private[ml]) Convert to a model in the old API */
private[ml] def toOld: OldDecisionTreeModel = {
new OldDecisionTreeModel(rootNode.toOld(1), OldAlgo.Classification)
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 f7d662df2fe5a6c53dce8ebf8bf925af8651251c..5da04d341d7c52b14f6e4d74abbe536a62ac606e 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
@@ -230,10 +230,10 @@ final class RandomForestClassificationModel private[ml] (
* - Average over trees:
* - importance(feature j) = sum (over nodes which split on feature j) of the gain,
* where gain is scaled by the number of instances passing through node
- * - Normalize importances for tree based on total number of training instances used
- * to build tree.
+ * - Normalize importances for tree to sum to 1.
* - Normalize feature importance vector to sum to 1.
*/
+ @Since("1.5.0")
lazy val featureImportances: Vector = RandomForest.featureImportances(trees, numFeatures)
/** (private[ml]) Convert to a model in the old API */
diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala
index 18c94f36387b9edd56cea85ec241c2e5aba00c06..897b23383c0cb4b775f4b3e0c5c068fdd04ad08f 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/regression/DecisionTreeRegressor.scala
@@ -169,6 +169,25 @@ final class DecisionTreeRegressionModel private[ml] (
s"DecisionTreeRegressionModel (uid=$uid) of depth $depth with $numNodes nodes"
}
+ /**
+ * Estimate of the importance of each feature.
+ *
+ * This generalizes the idea of "Gini" importance to other losses,
+ * following the explanation of Gini importance from "Random Forests" documentation
+ * by Leo Breiman and Adele Cutler, and following the implementation from scikit-learn.
+ *
+ * This feature importance is calculated as follows:
+ * - importance(feature j) = sum (over nodes which split on feature j) of the gain,
+ * where gain is scaled by the number of instances passing through node
+ * - Normalize importances for tree to sum to 1.
+ *
+ * Note: Feature importance for single decision trees can have high variance due to
+ * correlated predictor variables. Consider using a [[RandomForestRegressor]]
+ * to determine feature importance instead.
+ */
+ @Since("2.0.0")
+ lazy val featureImportances: Vector = RandomForest.featureImportances(this, numFeatures)
+
/** Convert to a model in the old API */
private[ml] def toOld: OldDecisionTreeModel = {
new OldDecisionTreeModel(rootNode.toOld(1), OldAlgo.Regression)
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 71e40b513ee0ac78a3ec228ed93da469f78582ae..798947b94a2802a9a774d90406a82b1b87b9d2f4 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
@@ -189,10 +189,10 @@ final class RandomForestRegressionModel private[ml] (
* - Average over trees:
* - importance(feature j) = sum (over nodes which split on feature j) of the gain,
* where gain is scaled by the number of instances passing through node
- * - Normalize importances for tree based on total number of training instances used
- * to build tree.
+ * - Normalize importances for tree to sum to 1.
* - Normalize feature importance vector to sum to 1.
*/
+ @Since("1.5.0")
lazy val featureImportances: Vector = RandomForest.featureImportances(trees, numFeatures)
/** (private[ml]) Convert to a model in the old API */
diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala
index ea733d577a5fd71bfd2678a32b7986e7cdbb697c..f994c258b2cadebb7064d85f80dab2d5e951eef0 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala
@@ -1089,12 +1089,9 @@ private[ml] object RandomForest extends Logging {
* - Average over trees:
* - importance(feature j) = sum (over nodes which split on feature j) of the gain,
* where gain is scaled by the number of instances passing through node
- * - Normalize importances for tree based on total number of training instances used
- * to build tree.
+ * - Normalize importances for tree to sum to 1.
* - Normalize feature importance vector to sum to 1.
*
- * Note: This should not be used with Gradient-Boosted Trees. It only makes sense for
- * independently trained trees.
* @param trees Unweighted forest of trees
* @param numFeatures Number of features in model (even if not all are explicitly used by
* the model).
@@ -1128,13 +1125,34 @@ private[ml] object RandomForest extends Logging {
maxFeatureIndex + 1
}
if (d == 0) {
- assert(totalImportances.size == 0, s"Unknown error in computing RandomForest feature" +
- s" importance: No splits in forest, but some non-zero importances.")
+ assert(totalImportances.size == 0, s"Unknown error in computing feature" +
+ s" importance: No splits found, but some non-zero importances.")
}
val (indices, values) = totalImportances.iterator.toSeq.sortBy(_._1).unzip
Vectors.sparse(d, indices.toArray, values.toArray)
}
+ /**
+ * Given a Decision Tree model, compute the importance of each feature.
+ * This generalizes the idea of "Gini" importance to other losses,
+ * following the explanation of Gini importance from "Random Forests" documentation
+ * by Leo Breiman and Adele Cutler, and following the implementation from scikit-learn.
+ *
+ * This feature importance is calculated as follows:
+ * - importance(feature j) = sum (over nodes which split on feature j) of the gain,
+ * where gain is scaled by the number of instances passing through node
+ * - Normalize importances for tree to sum to 1.
+ *
+ * @param tree Decision tree to compute importances for.
+ * @param numFeatures Number of features in model (even if not all are explicitly used by
+ * the model).
+ * If -1, then numFeatures is set based on the max feature index in all trees.
+ * @return Feature importance values, of length numFeatures.
+ */
+ private[ml] def featureImportances(tree: DecisionTreeModel, numFeatures: Int): Vector = {
+ featureImportances(Array(tree), numFeatures)
+ }
+
/**
* Recursive method for computing feature importances for one tree.
* This walks down the tree, adding to the importance of 1 feature at each node.
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala
index 9169bcd390dce9d0e5dffb531a3dd2674b1bc387..6d6836449979c10d01153759c043554f6c543a68 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/DecisionTreeClassifierSuite.scala
@@ -313,6 +313,27 @@ class DecisionTreeClassifierSuite extends SparkFunSuite with MLlibTestSparkConte
}
}
+ test("Feature importance with toy data") {
+ val dt = new DecisionTreeClassifier()
+ .setImpurity("gini")
+ .setMaxDepth(3)
+ .setSeed(123)
+
+ // In this data, feature 1 is very important.
+ val data: RDD[LabeledPoint] = TreeTests.featureImportanceData(sc)
+ val numFeatures = data.first().features.size
+ val categoricalFeatures = (0 to numFeatures).map(i => (i, 2)).toMap
+ val df = TreeTests.setMetadata(data, categoricalFeatures, 2)
+
+ val model = dt.fit(df)
+
+ val importances = model.featureImportances
+ val mostImportantFeature = importances.argmax
+ assert(mostImportantFeature === 1)
+ assert(importances.toArray.sum === 1.0)
+ assert(importances.toArray.forall(_ >= 0.0))
+ }
+
/////////////////////////////////////////////////////////////////////////////
// Tests of model save/load
/////////////////////////////////////////////////////////////////////////////
diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/RandomForestClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/RandomForestClassifierSuite.scala
index deb8ec771cb2796cc986195e73f8698045359f5d..6b810ab9eaa179122fa812b0826913ab3a2ee885 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/classification/RandomForestClassifierSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/classification/RandomForestClassifierSuite.scala
@@ -167,19 +167,15 @@ class RandomForestClassifierSuite extends SparkFunSuite with MLlibTestSparkConte
.setSeed(123)
// In this data, feature 1 is very important.
- val data: RDD[LabeledPoint] = sc.parallelize(Seq(
- new LabeledPoint(0, Vectors.dense(1, 0, 0, 0, 1)),
- new LabeledPoint(1, Vectors.dense(1, 1, 0, 1, 0)),
- new LabeledPoint(1, Vectors.dense(1, 1, 0, 0, 0)),
- new LabeledPoint(0, Vectors.dense(1, 0, 0, 0, 0)),
- new LabeledPoint(1, Vectors.dense(1, 1, 0, 0, 0))
- ))
+ val data: RDD[LabeledPoint] = TreeTests.featureImportanceData(sc)
val categoricalFeatures = Map.empty[Int, Int]
val df: DataFrame = TreeTests.setMetadata(data, categoricalFeatures, numClasses)
val importances = rf.fit(df).featureImportances
val mostImportantFeature = importances.argmax
assert(mostImportantFeature === 1)
+ assert(importances.toArray.sum === 1.0)
+ assert(importances.toArray.forall(_ >= 0.0))
}
/////////////////////////////////////////////////////////////////////////////
diff --git a/mllib/src/test/scala/org/apache/spark/ml/impl/TreeTests.scala b/mllib/src/test/scala/org/apache/spark/ml/impl/TreeTests.scala
index a808177cb9bf0e3a27bf42c24ceab5735239ef21..5561f6f0ef3c422672953fcb174e833fbd22652c 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/impl/TreeTests.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/impl/TreeTests.scala
@@ -19,10 +19,12 @@ package org.apache.spark.ml.impl
import scala.collection.JavaConverters._
+import org.apache.spark.SparkContext
import org.apache.spark.SparkFunSuite
import org.apache.spark.api.java.JavaRDD
import org.apache.spark.ml.attribute.{AttributeGroup, NominalAttribute, NumericAttribute}
import org.apache.spark.ml.tree._
+import org.apache.spark.mllib.linalg.Vectors
import org.apache.spark.mllib.regression.LabeledPoint
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.{DataFrame, SQLContext}
@@ -141,4 +143,15 @@ private[ml] object TreeTests extends SparkFunSuite {
val pred = parentImp.predict
new InternalNode(pred, parentImp.calculate(), gain, left, right, split, parentImp)
}
+
+ /**
+ * Create some toy data for testing feature importances.
+ */
+ def featureImportanceData(sc: SparkContext): RDD[LabeledPoint] = sc.parallelize(Seq(
+ new LabeledPoint(0, Vectors.dense(1, 0, 0, 0, 1)),
+ new LabeledPoint(1, Vectors.dense(1, 1, 0, 1, 0)),
+ new LabeledPoint(1, Vectors.dense(1, 1, 0, 0, 0)),
+ new LabeledPoint(0, Vectors.dense(1, 0, 0, 0, 0)),
+ new LabeledPoint(1, Vectors.dense(1, 1, 0, 0, 0))
+ ))
}
diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/DecisionTreeRegressorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/DecisionTreeRegressorSuite.scala
index 13165f67014c5b4512e2e1751f4bacc46a29f8a8..56b335a33a6b8369ffe8ff33e1a6acb0ceb29627 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/regression/DecisionTreeRegressorSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/regression/DecisionTreeRegressorSuite.scala
@@ -96,6 +96,26 @@ class DecisionTreeRegressorSuite extends SparkFunSuite with MLlibTestSparkContex
}
}
+ test("Feature importance with toy data") {
+ val dt = new DecisionTreeRegressor()
+ .setImpurity("variance")
+ .setMaxDepth(3)
+ .setSeed(123)
+
+ // In this data, feature 1 is very important.
+ val data: RDD[LabeledPoint] = TreeTests.featureImportanceData(sc)
+ val categoricalFeatures = Map.empty[Int, Int]
+ val df: DataFrame = TreeTests.setMetadata(data, categoricalFeatures, 0)
+
+ val model = dt.fit(df)
+
+ val importances = model.featureImportances
+ val mostImportantFeature = importances.argmax
+ assert(mostImportantFeature === 1)
+ assert(importances.toArray.sum === 1.0)
+ assert(importances.toArray.forall(_ >= 0.0))
+ }
+
/////////////////////////////////////////////////////////////////////////////
// Tests of model save/load
/////////////////////////////////////////////////////////////////////////////
diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/RandomForestRegressorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/RandomForestRegressorSuite.scala
index 7e751e4b553b6cf73a64daa2a246b6a45867c914..efb117f8f9b16e9ab0aab55d57223d3a5e91f032 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/regression/RandomForestRegressorSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/regression/RandomForestRegressorSuite.scala
@@ -20,7 +20,6 @@ package org.apache.spark.ml.regression
import org.apache.spark.SparkFunSuite
import org.apache.spark.ml.impl.TreeTests
import org.apache.spark.ml.util.MLTestingUtils
-import org.apache.spark.mllib.linalg.Vectors
import org.apache.spark.mllib.regression.LabeledPoint
import org.apache.spark.mllib.tree.{EnsembleTestHelper, RandomForest => OldRandomForest}
import org.apache.spark.mllib.tree.configuration.{Algo => OldAlgo}
@@ -82,23 +81,17 @@ class RandomForestRegressorSuite extends SparkFunSuite with MLlibTestSparkContex
.setSeed(123)
// In this data, feature 1 is very important.
- val data: RDD[LabeledPoint] = sc.parallelize(Seq(
- new LabeledPoint(0, Vectors.dense(1, 0, 0, 0, 1)),
- new LabeledPoint(1, Vectors.dense(1, 1, 0, 1, 0)),
- new LabeledPoint(1, Vectors.dense(1, 1, 0, 0, 0)),
- new LabeledPoint(0, Vectors.dense(1, 0, 0, 0, 0)),
- new LabeledPoint(1, Vectors.dense(1, 1, 0, 0, 0))
- ))
+ val data: RDD[LabeledPoint] = TreeTests.featureImportanceData(sc)
val categoricalFeatures = Map.empty[Int, Int]
val df: DataFrame = TreeTests.setMetadata(data, categoricalFeatures, 0)
val model = rf.fit(df)
- // copied model must have the same parent.
- MLTestingUtils.checkCopy(model)
val importances = model.featureImportances
val mostImportantFeature = importances.argmax
assert(mostImportantFeature === 1)
+ assert(importances.toArray.sum === 1.0)
+ assert(importances.toArray.forall(_ >= 0.0))
}
/////////////////////////////////////////////////////////////////////////////