使用MLlib從Spark的決策樹中找到重要性值

Question

我們正在使用MLlib為決策樹運行Spark 1.0或1.1。

當我使用示例數據運行示例SCALA代碼時，它沒有任何錯誤，但是我從結果中找不到功能的重要性。

有人知道如何獲取這些值嗎？

Answer 1

在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

Answer 2

當您最終訓練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])