Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[jvm-packages] rdds containing boosters should be cleaned once we got boosters to driver (#2183)

Browse Source
This commit is contained in:
Nan Zhu 2017-04-11 06:12:49 -07:00 committed by GitHub
parent f08077606c
commit a837fa9620
1 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
jvm-packages/xgboost4j-spark/src/main/scala/ml/dmlc/xgboost4j/scala/spark/XGBoost.scala
View File

@ -123,10 +123,8 @@ object XGBoost extends Serializable {
}
val partitionItr = fromDenseToSparseLabeledPoints(trainingSamples, missing)
val trainingSet = new DMatrix(new JDMatrix(partitionItr, cacheFileName))
if (xgBoostConfMap.isDefinedAt("groupData")
&& xgBoostConfMap.get("groupData").get != null) {
trainingSet.setGroup(
xgBoostConfMap.get("groupData").get.asInstanceOf[Seq[Seq[Int]]](
if (xgBoostConfMap.contains("groupData") && xgBoostConfMap("groupData") != null) {
trainingSet.setGroup(xgBoostConfMap("groupData").asInstanceOf[Seq[Seq[Int]]](
TaskContext.getPartitionId()).toArray)
}
booster = SXGBoost.train(trainingSet, xgBoostConfMap, round,
@ -309,7 +307,10 @@ object XGBoost extends Serializable {
configMap: Map[String, Any], sparkJobThread: Thread, isClassificationTask: Boolean):
XGBoostModel = {
if (trackerReturnVal == 0) {
convertBoosterToXGBoostModel(distributedBoosters.first(), isClassificationTask)
val xgboostModel = convertBoosterToXGBoostModel(distributedBoosters.first(),
isClassificationTask)
distributedBoosters.unpersist(false)
xgboostModel
} else {
try {
if (sparkJobThread.isAlive) {