From ab68a0ccc7dd16336e914c5f41c35c1dde363aaf Mon Sep 17 00:00:00 2001 From: CodingCat Date: Fri, 11 Mar 2016 13:57:03 -0500 Subject: [PATCH] fix examples --- .../example/flink/DistTrainWithFlink.scala | 2 +- .../example/spark/DistTrainWithSpark.scala | 17 +++++++++++------ .../ml/dmlc/xgboost4j/scala/flink/XGBoost.scala | 9 +++++---- .../xgboost4j/scala/flink/XGBoostModel.scala | 11 ++++++++++- .../scala/ml/dmlc/xgboost4j/scala/DMatrix.scala | 4 ++-- 5 files changed, 29 insertions(+), 14 deletions(-) diff --git a/jvm-packages/xgboost4j-example/src/main/scala/ml/dmlc/xgboost4j/scala/example/flink/DistTrainWithFlink.scala b/jvm-packages/xgboost4j-example/src/main/scala/ml/dmlc/xgboost4j/scala/example/flink/DistTrainWithFlink.scala index 200ca3ba1..74b24ac35 100644 --- a/jvm-packages/xgboost4j-example/src/main/scala/ml/dmlc/xgboost4j/scala/example/flink/DistTrainWithFlink.scala +++ b/jvm-packages/xgboost4j-example/src/main/scala/ml/dmlc/xgboost4j/scala/example/flink/DistTrainWithFlink.scala @@ -34,7 +34,7 @@ object DistTrainWithFlink { // number of iterations val round = 2 // train the model - val model = XGBoost.train(paramMap, trainData, round) + val model = XGBoost.train(trainData, paramMap, round) val predTest = model.predict(testData.map{x => x.vector}) model.saveModelAsHadoopFile("file:///path/to/xgboost.model") } diff --git a/jvm-packages/xgboost4j-example/src/main/scala/ml/dmlc/xgboost4j/scala/example/spark/DistTrainWithSpark.scala b/jvm-packages/xgboost4j-example/src/main/scala/ml/dmlc/xgboost4j/scala/example/spark/DistTrainWithSpark.scala index 0c65ce59d..82e6e626b 100644 --- a/jvm-packages/xgboost4j-example/src/main/scala/ml/dmlc/xgboost4j/scala/example/spark/DistTrainWithSpark.scala +++ b/jvm-packages/xgboost4j-example/src/main/scala/ml/dmlc/xgboost4j/scala/example/spark/DistTrainWithSpark.scala @@ -16,29 +16,34 @@ package ml.dmlc.xgboost4j.scala.example.spark -import ml.dmlc.xgboost4j.scala.spark.XGBoost +import ml.dmlc.xgboost4j.scala.DMatrix +import ml.dmlc.xgboost4j.scala.spark.{DataUtils, XGBoost} import org.apache.spark.SparkContext import org.apache.spark.mllib.util.MLUtils object DistTrainWithSpark { def main(args: Array[String]): Unit = { - if (args.length != 4) { + if (args.length != 5) { println( - "usage: program num_of_rounds num_workers training_path model_path") + "usage: program num_of_rounds num_workers training_path test_path model_path") sys.exit(1) } val sc = new SparkContext() val inputTrainPath = args(2) - val outputModelPath = args(3) + val inputTestPath = args(3) + val outputModelPath = args(4) // number of iterations val numRound = args(0).toInt - val trainRDD = MLUtils.loadLibSVMFile(sc, inputTrainPath).repartition(args(1).toInt) + import DataUtils._ + val trainRDD = MLUtils.loadLibSVMFile(sc, inputTrainPath) + val testSet = MLUtils.loadLibSVMFile(sc, inputTestPath).collect().iterator // training parameters val paramMap = List( "eta" -> 0.1f, "max_depth" -> 2, "objective" -> "binary:logistic").toMap - val xgboostModel = XGBoost.train(trainRDD, paramMap, numRound) + val xgboostModel = XGBoost.train(trainRDD, paramMap, numRound, nWorkers = args(1).toInt) + xgboostModel.predict(new DMatrix(testSet)) // save model to HDFS path xgboostModel.saveModelAsHadoopFile(outputModelPath) } diff --git a/jvm-packages/xgboost4j-flink/src/main/scala/ml/dmlc/xgboost4j/scala/flink/XGBoost.scala b/jvm-packages/xgboost4j-flink/src/main/scala/ml/dmlc/xgboost4j/scala/flink/XGBoost.scala index 94b36be91..8d00ec9c1 100644 --- a/jvm-packages/xgboost4j-flink/src/main/scala/ml/dmlc/xgboost4j/scala/flink/XGBoost.scala +++ b/jvm-packages/xgboost4j-flink/src/main/scala/ml/dmlc/xgboost4j/scala/flink/XGBoost.scala @@ -81,13 +81,14 @@ object XGBoost { /** * Train a xgboost model with link. * - * @param params The parameters to XGBoost. * @param dtrain The training data. + * @param params The parameters to XGBoost. * @param round Number of rounds to train. */ - def train(params: Map[String, Any], - dtrain: DataSet[LabeledVector], - round: Int): XGBoostModel = { + def train( + dtrain: DataSet[LabeledVector], + params: Map[String, Any], + round: Int): XGBoostModel = { val tracker = new RabitTracker(dtrain.getExecutionEnvironment.getParallelism) if (tracker.start()) { dtrain diff --git a/jvm-packages/xgboost4j-flink/src/main/scala/ml/dmlc/xgboost4j/scala/flink/XGBoostModel.scala b/jvm-packages/xgboost4j-flink/src/main/scala/ml/dmlc/xgboost4j/scala/flink/XGBoostModel.scala index 54bcdb27b..b66439c38 100644 --- a/jvm-packages/xgboost4j-flink/src/main/scala/ml/dmlc/xgboost4j/scala/flink/XGBoostModel.scala +++ b/jvm-packages/xgboost4j-flink/src/main/scala/ml/dmlc/xgboost4j/scala/flink/XGBoostModel.scala @@ -37,6 +37,15 @@ class XGBoostModel (booster: Booster) extends Serializable { .create(new Path(modelPath))) } + /** + * predict with the given DMatrix + * @param testSet the local test set represented as DMatrix + * @return prediction result + */ + def predict(testSet: DMatrix): Array[Array[Float]] = { + booster.predict(testSet, true, 0) + } + /** * Predict given vector dataset. * @@ -44,7 +53,7 @@ class XGBoostModel (booster: Booster) extends Serializable { * @return The prediction result. */ def predict(data: DataSet[Vector]) : DataSet[Array[Float]] = { - val predictMap: Iterator[Vector] => TraversableOnce[Array[Float]] = + val predictMap: Iterator[Vector] => Traversable[Array[Float]] = (it: Iterator[Vector]) => { val mapper = (x: Vector) => { val (index, value) = x.toSeq.unzip diff --git a/jvm-packages/xgboost4j/src/main/scala/ml/dmlc/xgboost4j/scala/DMatrix.scala b/jvm-packages/xgboost4j/src/main/scala/ml/dmlc/xgboost4j/scala/DMatrix.scala index cdf3a9844..27bcc618c 100644 --- a/jvm-packages/xgboost4j/src/main/scala/ml/dmlc/xgboost4j/scala/DMatrix.scala +++ b/jvm-packages/xgboost4j/src/main/scala/ml/dmlc/xgboost4j/scala/DMatrix.scala @@ -35,10 +35,10 @@ class DMatrix private[scala](private[scala] val jDMatrix: JDMatrix) { * init DMatrix from Iterator of LabeledPoint * * @param dataIter An iterator of LabeledPoint - * @param cacheInfo Cache path information, used for external memory setting, can be null. + * @param cacheInfo Cache path information, used for external memory setting, null by default. * @throws XGBoostError native error */ - def this(dataIter: Iterator[LabeledPoint], cacheInfo: String) { + def this(dataIter: Iterator[LabeledPoint], cacheInfo: String = null) { this(new JDMatrix(dataIter.asJava, cacheInfo)) }