使用MLlib从Spark的决策树中查找重要性值

时间:2014-12-04 07:55:00

标签: scala apache-spark decision-tree apache-spark-mllib

我们使用MLlib为决策树运行Spark 1.0或1.1。

当我使用示例数据运行示例SCALA代码时,它没有错误,但我无法从结果中找到功能重要性。

任何人都有关于如何获取价值的信息?

2 个答案:

答案 0 :(得分:3)

在Spark 2+中,您可以执行以下操作:

val vectorAssembler = new VectorAssembler().setInputCols(featureArray)
val decisionTreeModel = decisionTree.fit(trainingDataset)
val featureImportances = decisionTreeModel.featureImportances // Sparse or Dense Vector

featureArray.zip(featureImportances.toArray).sortBy(_._2).reverse

答案 1 :(得分:1)

当您在最后训练DecisionTreeModel时,您将拥有此类

class DecisionTreeModel(val topNode: Node, val algo: Algo) {
   ...
}

您可以从顶部开始遍历节点,您可以从中获得所需的一切(预测+ InformationGainStats)

class Node (
    val id: Int,
    val predict: Double,
    val isLeaf: Boolean,
    val split: Option[Split],
    var leftNode: Option[Node],
    var rightNode: Option[Node],
    val stats: Option[InformationGainStats])