From 224786f67f59d0a6ddd8c19c8ddfccf38bf157d5 Mon Sep 17 00:00:00 2001 From: Shaochen Shi Date: Fri, 8 Mar 2019 10:36:47 +0800 Subject: [PATCH] [xgboost4j-spark] Allow set the parameter "maxLeaves". (#4226) * Allow set the parameter "maxLeaves". * Add "setMaxLeaves" to XGBoostRegressor. --- .../scala/ml/dmlc/xgboost4j/scala/spark/XGBoostClassifier.scala | 2 ++ .../scala/ml/dmlc/xgboost4j/scala/spark/XGBoostRegressor.scala | 2 ++ 2 files changed, 4 insertions(+) diff --git a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostClassifier.scala b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostClassifier.scala index 04fbfda97..366f871c2 100644 --- a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostClassifier.scala +++ b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostClassifier.scala @@ -113,6 +113,8 @@ class XGBoostClassifier ( def setMaxBins(value: Int): this.type = set(maxBins, value) + def setMaxLeaves(value: Int): this.type = set(maxLeaves, value) + def setSketchEps(value: Double): this.type = set(sketchEps, value) def setScalePosWeight(value: Double): this.type = set(scalePosWeight, value) diff --git a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostRegressor.scala b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostRegressor.scala index 20dd24996..0abad8b9c 100644 --- a/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostRegressor.scala +++ b/jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostRegressor.scala @@ -113,6 +113,8 @@ class XGBoostRegressor ( def setMaxBins(value: Int): this.type = set(maxBins, value) + def setMaxLeaves(value: Int): this.type = set(maxLeaves, value) + def setSketchEps(value: Double): this.type = set(sketchEps, value) def setScalePosWeight(value: Double): this.type = set(scalePosWeight, value)