Skip to content

[SPARK-22707][ML] Optimize CrossValidator memory occupation by models in fitting #19904

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
WeichenXu123 wants to merge 4 commits into from
Closed

[SPARK-22707][ML] Optimize CrossValidator memory occupation by models in fitting #19904

Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ec50dad
init pr
WeichenXu123 Dec 5, 2017
2cc7c28
improve code
WeichenXu123 Dec 8, 2017
ccd2689
choose approach 3 update code
WeichenXu123 Dec 19, 2017
cad2104
nit update
WeichenXu123 Dec 20, 2017
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
choose approach 3 update code
  • Loading branch information
@WeichenXu123
WeichenXu123 committed Dec 19, 2017
commit ccd26895e749cce1acbe65a5507cba6b0e6c42d3
6 changes: 1 addition & 5 deletions mllib/src/main/scala/org/apache/spark/ml/tuning/CrossValidator.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -147,15 +147,10 @@ class CrossValidator @Since("1.2.0") (@Since("1.4.0") override val uid: String)
val validationDataset = sparkSession.createDataFrame(validation, schema).cache()
logDebug(s"Train split $splitIndex with multiple sets of parameters.")

val completeFitCount = new AtomicInteger(0)
// Fit models in a Future for training in parallel
val foldMetricFutures = epm.zipWithIndex.map { case (paramMap, paramIndex) =>
Future[Double] {
val model = est.fit(trainingDataset, paramMap).asInstanceOf[Model[_]]
if (completeFitCount.incrementAndGet() == epm.length) {
trainingDataset.unpersist()
}

if (collectSubModelsParam) {
subModels.get(splitIndex)(paramIndex) = model
}
Expand All @@ -168,6 +163,7 @@ class CrossValidator @Since("1.2.0") (@Since("1.4.0") override val uid: String)

// Wait for metrics to be calculated before unpersisting validation dataset
val foldMetrics = foldMetricFutures.map(ThreadUtils.awaitResult(_, Duration.Inf))
trainingDataset.unpersist()
validationDataset.unpersist()
foldMetrics
}.transpose.map(_.sum / $(numFolds)) // Calculate average metric over all splits
Expand Down