Github user mengxr commented on a diff in the pull request:
https://github.com/apache/spark/pull/5626#discussion_r28987650
--- Diff: mllib/src/main/scala/org/apache/spark/ml/tree/treeModels.scala ---
@@ -58,3 +58,43 @@ trait DecisionTreeModel {
header + rootNode.subtreeToString(2)
}
}
+
+/**
+ * :: AlphaComponent ::
+ *
+ * Abstraction for models which are ensembles of decision trees
+ *
+ * TODO: Add support for predicting probabilities and raw predictions
+ */
+@AlphaComponent
+trait TreeEnsembleModel {
+
+ // Note: We use getTrees since subclasses of TreeEnsembleModel will
store subclasses of
+ // DecisionTreeModel.
+
+ /** Trees in this ensemble. Warning: These have null parent Estimators.
*/
+ def getTrees: Array[DecisionTreeModel]
--- End diff --
If we don't allow users to set trees, it is a property and maybe calling it
`trees` would be sufficient. Same for `getTreeWeights`.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at [email protected] or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
