Expose setCustomObj & setCustomEval for XGBoostClassifier & XGBoostRegressor. (#3486)

This commit is contained in:
Yanbo Liang 2018-07-17 21:16:51 -07:00 committed by Nan Zhu
parent b6dcbf0e07
commit c004cea788
2 changed files with 10 additions and 0 deletions

View File

@ -22,6 +22,7 @@ import scala.collection.mutable
import ml.dmlc.xgboost4j.java.Rabit import ml.dmlc.xgboost4j.java.Rabit
import ml.dmlc.xgboost4j.scala.{Booster, DMatrix, XGBoost => SXGBoost} import ml.dmlc.xgboost4j.scala.{Booster, DMatrix, XGBoost => SXGBoost}
import ml.dmlc.xgboost4j.scala.{EvalTrait, ObjectiveTrait}
import ml.dmlc.xgboost4j.scala.spark.params._ import ml.dmlc.xgboost4j.scala.spark.params._
import ml.dmlc.xgboost4j.{LabeledPoint => XGBLabeledPoint} import ml.dmlc.xgboost4j.{LabeledPoint => XGBLabeledPoint}
@ -134,6 +135,10 @@ class XGBoostClassifier (
def setNumEarlyStoppingRounds(value: Int): this.type = set(numEarlyStoppingRounds, value) def setNumEarlyStoppingRounds(value: Int): this.type = set(numEarlyStoppingRounds, value)
def setCustomObj(value: ObjectiveTrait): this.type = set(customObj, value)
def setCustomEval(value: EvalTrait): this.type = set(customEval, value)
// called at the start of fit/train when 'eval_metric' is not defined // called at the start of fit/train when 'eval_metric' is not defined
private def setupDefaultEvalMetric(): String = { private def setupDefaultEvalMetric(): String = {
require(isDefined(objective), "Users must set \'objective\' via xgboostParams.") require(isDefined(objective), "Users must set \'objective\' via xgboostParams.")

View File

@ -23,6 +23,7 @@ import ml.dmlc.xgboost4j.java.Rabit
import ml.dmlc.xgboost4j.{LabeledPoint => XGBLabeledPoint} import ml.dmlc.xgboost4j.{LabeledPoint => XGBLabeledPoint}
import ml.dmlc.xgboost4j.scala.spark.params.{DefaultXGBoostParamsReader, _} import ml.dmlc.xgboost4j.scala.spark.params.{DefaultXGBoostParamsReader, _}
import ml.dmlc.xgboost4j.scala.{Booster, DMatrix, XGBoost => SXGBoost} import ml.dmlc.xgboost4j.scala.{Booster, DMatrix, XGBoost => SXGBoost}
import ml.dmlc.xgboost4j.scala.{EvalTrait, ObjectiveTrait}
import org.apache.hadoop.fs.Path import org.apache.hadoop.fs.Path
import org.apache.spark.TaskContext import org.apache.spark.TaskContext
@ -136,6 +137,10 @@ class XGBoostRegressor (
def setNumEarlyStoppingRounds(value: Int): this.type = set(numEarlyStoppingRounds, value) def setNumEarlyStoppingRounds(value: Int): this.type = set(numEarlyStoppingRounds, value)
def setCustomObj(value: ObjectiveTrait): this.type = set(customObj, value)
def setCustomEval(value: EvalTrait): this.type = set(customEval, value)
// called at the start of fit/train when 'eval_metric' is not defined // called at the start of fit/train when 'eval_metric' is not defined
private def setupDefaultEvalMetric(): String = { private def setupDefaultEvalMetric(): String = {
require(isDefined(objective), "Users must set \'objective\' via xgboostParams.") require(isDefined(objective), "Users must set \'objective\' via xgboostParams.")