Add new parameter singlePrecisionHistogram to xgboost4j-spark

Expose the existing 'singlePrecisionHistogram' param to the Spark layer.
dmlc · Jun 18, 2020 · 42dd8ba · 42dd8ba
1 parent 38ee514
commit 42dd8ba
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 1 deletion.
diff --git a/...ages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostClassifier.scala b/...ages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostClassifier.scala
@@ -141,6 +141,9 @@ class XGBoostClassifier (
 
   def setCustomEval(value: EvalTrait): this.type = set(customEval, value)
 
+  def setSinglePrecisionHistogram(value: Boolean): this.type =
+    set(singlePrecisionHistogram, value)
+
   // called at the start of fit/train when 'eval_metric' is not defined
   private def setupDefaultEvalMetric(): String = {
     require(isDefined(objective), "Users must set \'objective\' via xgboostParams.")

diff --git a/...kages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostRegressor.scala b/...kages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoostRegressor.scala
@@ -145,6 +145,9 @@ class XGBoostRegressor (
 
   def setCustomEval(value: EvalTrait): this.type = set(customEval, value)
 
+  def setSinglePrecisionHistogram(value: Boolean): this.type =
+    set(singlePrecisionHistogram, value)
+
   // called at the start of fit/train when 'eval_metric' is not defined
   private def setupDefaultEvalMetric(): String = {
     require(isDefined(objective), "Users must set \'objective\' via xgboostParams.")

diff --git a/...s/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/params/BoosterParams.scala b/...s/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/params/BoosterParams.scala
@@ -18,7 +18,7 @@ package ml.dmlc.xgboost4j.scala.spark.params
 
 import scala.collection.immutable.HashSet
 
-import org.apache.spark.ml.param.{DoubleParam, IntParam, Param, Params}
+import org.apache.spark.ml.param.{DoubleParam, IntParam, BooleanParam, Param, Params}
 
 private[spark] trait BoosterParams extends Params {
 
@@ -173,6 +173,14 @@ private[spark] trait BoosterParams extends Params {
 
   final def getMaxBins: Int = $(maxBins)
 
+  /**
+   * whether to build histograms using single precision floating point values
+   */
+  final val singlePrecisionHistogram = new BooleanParam(this, "singlePrecisionHistogram",
+    "whether to use single precision to build histograms")
+
+  final def getSinglePrecisionHistogram: Boolean = $(singlePrecisionHistogram)
+
   /**
    * This is only used for approximate greedy algorithm.
    * This roughly translated into O(1 / sketch_eps) number of bins. Compared to directly select