add enhancement for xgboost (intel-analytics#2807)

dding3 · Sep 2, 2020 · 188d430 · 188d430
1 parent 086efd4
commit 188d430
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/dllib/nnframes/NNClassifier.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/dllib/nnframes/NNClassifier.scala
@@ -371,6 +371,7 @@ class XGBRegressor () {
 
   private val model = new XGBoostRegressor()
   model.setNthread(EngineRef.getCoreNumber())
+  model.setMaxBins(256)
 
   def setLabelCol(labelColName : String) : this.type = {
     model.setLabelCol(labelColName)
@@ -383,6 +384,7 @@ class XGBRegressor () {
   }
 
   def fit(df: DataFrame): XGBRegressorModel = {
+    df.repartition(EngineRef.getNodeNumber())
     val xgbModel = model.fit(df)
     new XGBRegressorModel(xgbModel)
   }