spark-reviews mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From mengxr <...@git.apache.org>
Subject [GitHub] spark pull request: [SPARK-2851] [mllib] DecisionTree Python consi...
Date Wed, 06 Aug 2014 19:47:49 GMT
Github user mengxr commented on a diff in the pull request:

    https://github.com/apache/spark/pull/1798#discussion_r15898657
  
    --- Diff: mllib/src/main/scala/org/apache/spark/mllib/tree/DecisionTree.scala ---
    @@ -300,6 +309,93 @@ object DecisionTree extends Serializable with Logging {
         new DecisionTree(strategy).train(input)
       }
     
    +  /**
    +   * Method to train a decision tree model for binary or multiclass classification.
    +   *
    +   * @param input Training dataset: RDD of [[org.apache.spark.mllib.regression.LabeledPoint]].
    +   *              Labels should take values {0, 1, ..., numClasses-1}.
    +   * @param numClassesForClassification number of classes for classification.
    +   * @param categoricalFeaturesInfo Map storing arity of categorical features.
    +   *                                E.g., an entry (n -> k) indicates that feature
n is categorical
    +   *                                with k categories indexed from 0: {0, 1, ..., k-1}.
    +   * @param impurity Criterion used for information gain calculation.
    +   *                 Supported values: "gini" (recommended) or "entropy".
    +   * @param maxDepth Maximum depth of the tree.
    +   *                 E.g., depth 0 means 1 leaf node; depth 1 means 1 internal node +
2 leaf nodes.
    +   *                  (suggested value: 4)
    +   * @param maxBins maximum number of bins used for splitting features
    +   *                 (suggested value: 100)
    +   * @return DecisionTreeModel that can be used for prediction
    +   */
    +  def trainClassifier(
    +      input: RDD[LabeledPoint],
    +      numClassesForClassification: Int,
    +      categoricalFeaturesInfo: Map[Int, Int],
    +      impurity: String,
    +      maxDepth: Int,
    +      maxBins: Int): DecisionTreeModel = {
    +    val impurityType = Impurities.fromString(impurity)
    +    train(input, Classification, impurityType, maxDepth, numClassesForClassification,
maxBins, Sort,
    +      categoricalFeaturesInfo)
    +  }
    +
    +  /**
    +   * Java-friendly API for [[org.apache.spark.mllib.tree.DecisionTree$#trainClassifier]]
    +   */
    +  def trainClassifier(
    +      input: RDD[LabeledPoint],
    --- End diff --
    
    `RDD` -> `JavaRDD`. Sorry I missed this in the first pass. 


---
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 infrastructure@apache.org or file a JIRA ticket
with INFRA.
---

---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Mime
View raw message