How to access individual trees in a model created by RandomForestClassifier (spark.ml version)?

Question

How to access individual trees in a model created by Spark ML RandomForestClassifier ? I am using the Scala version of RandomForestClassifier.

Answer 1

In fact, it has the trees attribute:

 import org.apache.spark.ml.attribute.NominalAttribute import org.apache.spark.ml.classification.{ RandomForestClassificationModel, RandomForestClassifier, DecisionTreeClassificationModel } val meta = NominalAttribute .defaultAttr .withName("label") .withValues("0.0", "1.0") .toMetadata val data = sqlContext.read.format("libsvm") .load("data/mllib/sample_libsvm_data.txt") .withColumn("label", $"label".as("label", meta)) val rf: RandomForestClassifier = new RandomForestClassifier() .setLabelCol("label") .setFeaturesCol("features") val trees: Array[DecisionTreeClassificationModel] = rf.fit(data).trees.collect { case t: DecisionTreeClassificationModel => t } 

As you can see, the only problem is getting the types correctly so that we can actually use them:

 trees.head.transform(data).show(3) // +-----+--------------------+-------------+-----------+----------+ // |label| features|rawPrediction|probability|prediction| // +-----+--------------------+-------------+-----------+----------+ // | 0.0|(692,[127,128,129...| [33.0,0.0]| [1.0,0.0]| 0.0| // | 1.0|(692,[158,159,160...| [0.0,59.0]| [0.0,1.0]| 1.0| // | 1.0|(692,[124,125,126...| [0.0,59.0]| [0.0,1.0]| 1.0| // +-----+--------------------+-------------+-----------+----------+ // only showing top 3 rows 

Note

If you work with pipelines, you can also extract individual trees:

 import org.apache.spark.ml.Pipeline val model = new Pipeline().setStages(Array(rf)).fit(data) // There is only one stage and know its type // but lets be thorough val rfModelOption = model.stages.headOption match { case Some(m: RandomForestClassificationModel) => Some(m) case _ => None } val trees = rfModelOption.map { _.trees // ... as before }.getOrElse(Array())