aboutsummaryrefslogtreecommitdiff
path: root/mllib
diff options
context:
space:
mode:
authorHolden Karau <holden@pigscanfly.ca>2015-06-22 22:40:19 -0700
committerJoseph K. Bradley <joseph@databricks.com>2015-06-22 22:40:19 -0700
commit164fe2aa44993da6c77af6de5efdae47a8b3958c (patch)
tree938944e023b53542ee306edcecec99a24d66429f /mllib
parent44fa7df64daa55bd6eb1f2c219a9701b34e1c2a3 (diff)
downloadspark-164fe2aa44993da6c77af6de5efdae47a8b3958c.tar.gz
spark-164fe2aa44993da6c77af6de5efdae47a8b3958c.tar.bz2
spark-164fe2aa44993da6c77af6de5efdae47a8b3958c.zip
[SPARK-7781] [MLLIB] gradient boosted trees.train regressor missing max bins
Author: Holden Karau <holden@pigscanfly.ca> Closes #6331 from holdenk/SPARK-7781-GradientBoostedTrees.trainRegressor-missing-max-bins and squashes the following commits: 2894695 [Holden Karau] remove extra blank line 2573e8d [Holden Karau] Update the scala side of the pythonmllibapi and make the test a bit nicer too 3a09170 [Holden Karau] add maxBins to to the train method as well af7f274 [Holden Karau] Add maxBins to GradientBoostedTrees.trainRegressor and correctly mention the default of 32 in other places where it mentioned 100
Diffstat (limited to 'mllib')
-rw-r--r--mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala4
1 files changed, 3 insertions, 1 deletions
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala b/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala
index 634d56d08d..f9a271f47e 100644
--- a/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala
+++ b/mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala
@@ -696,12 +696,14 @@ private[python] class PythonMLLibAPI extends Serializable {
lossStr: String,
numIterations: Int,
learningRate: Double,
- maxDepth: Int): GradientBoostedTreesModel = {
+ maxDepth: Int,
+ maxBins: Int): GradientBoostedTreesModel = {
val boostingStrategy = BoostingStrategy.defaultParams(algoStr)
boostingStrategy.setLoss(Losses.fromString(lossStr))
boostingStrategy.setNumIterations(numIterations)
boostingStrategy.setLearningRate(learningRate)
boostingStrategy.treeStrategy.setMaxDepth(maxDepth)
+ boostingStrategy.treeStrategy.setMaxBins(maxBins)
boostingStrategy.treeStrategy.categoricalFeaturesInfo = categoricalFeaturesInfo.asScala.toMap
val cached = data.rdd.persist(StorageLevel.MEMORY_AND_DISK)