Skip to content

[SPARK-22126][ML][WIP] Fix model-specific optimization support for ML tuning #19350

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
wants to merge 3 commits into from
Closed

[SPARK-22126][ML][WIP] Fix model-specific optimization support for ML tuning #19350

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
d5625a6
init pr
WeichenXu123 Sep 26, 2017
5614390
update & resolve conflicts
WeichenXu123 Nov 15, 2017
1b4a384
add java api
WeichenXu123 Nov 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
48 changes: 48 additions & 0 deletions mllib/src/main/scala/org/apache/spark/ml/Estimator.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,14 @@
package org.apache.spark.ml

import scala.annotation.varargs
import scala.concurrent.{ExecutionContext, Future}
import scala.concurrent.duration.Duration

import org.apache.spark.annotation.{DeveloperApi, Since}
import org.apache.spark.api.java.function.VoidFunction2
import org.apache.spark.ml.param.{ParamMap, ParamPair}
import org.apache.spark.sql.Dataset
import org.apache.spark.util.ThreadUtils

/**
* :: DeveloperApi ::
Expand Down Expand Up @@ -82,5 +86,49 @@ abstract class Estimator[M <: Model[M]] extends PipelineStage {
paramMaps.map(fit(dataset, _))
}

/**
* (Java-specific)
*/
@Since("2.3.0")
def fit(dataset: Dataset[_], paramMaps: Array[ParamMap],
unpersistDatasetAfterFitting: Boolean, executionContext: ExecutionContext,
modelCallback: VoidFunction2[Model[_], Int]): Unit = {
// Fit models in a Future for training in parallel
val modelFutures = paramMaps.map { paramMap =>
Future[Model[_]] {
fit(dataset, paramMap).asInstanceOf[Model[_]]
Copy link
Contributor

@MLnick MLnick Dec 13, 2017

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How will this work in a pipeline?

If the Estimator in CV is a Pipeline, then here it will call fit(dataset, paramMap) on the Pipeline which will in turn fit on each stage with that paramMap. This is what the current parallel CV is doing.

But if we have a stage with model-specific optimization (let's say for arguments sake a LinearRegression that can internally optimize maxIter) then its fit will be called with only a single paramMap arg.

So that pushing the parallel fit into Estimator nullifies any benefit from model-specific optimizations?

Copy link
Contributor Author

@WeichenXu123 WeichenXu123 Dec 13, 2017

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@MLnick Oh, the design is still under discussion on JIRA and will be changed I think. I should mark this WIP. thanks!

Copy link
Contributor Author

@WeichenXu123 WeichenXu123 Dec 15, 2017

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@MLnick I dicussed with @jkbradley @MrBago offline and here is the newest proposal
https://docs.google.com/document/d/1xw5M4sp1e0eQie75yIt-r6-GTuD5vpFf_I6v-AFBM3M/edit?usp=sharing
Thanks!

} (executionContext)
}

if (unpersistDatasetAfterFitting) {
// Unpersist training data only when all models have trained
Future.sequence[Model[_], Iterable](modelFutures)(implicitly, executionContext)
.onComplete { _ => dataset.unpersist() } (executionContext)
}

val modelCallbackFutures = modelFutures.zipWithIndex.map {
case (modelFuture, paramMapIndex) =>
modelFuture.map { model =>
modelCallback.call(model, paramMapIndex)
} (executionContext)
}
modelCallbackFutures.map(ThreadUtils.awaitResult(_, Duration.Inf))
}

/**
* (Scala-specific)
*/
@Since("2.3.0")
def fit(dataset: Dataset[_], paramMaps: Array[ParamMap],
unpersistDatasetAfterFitting: Boolean, executionContext: ExecutionContext,
modelCallback: (Model[_], Int) => Unit): Unit = {
fit(dataset, paramMaps, unpersistDatasetAfterFitting, executionContext,
new VoidFunction2[Model[_], Int] {
override def call(model: Model[_], paramMapIndex: Int): Unit = {
modelCallback(model, paramMapIndex)
}
})
}

override def copy(extra: ParamMap): Estimator[M]
}
30 changes: 8 additions & 22 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 @@ -146,34 +146,20 @@ 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.")

// Fit models in a Future for training in parallel
val modelFutures = epm.zipWithIndex.map { case (paramMap, paramIndex) =>
Future[Model[_]] {
val model = est.fit(trainingDataset, paramMap).asInstanceOf[Model[_]]

val foldMetrics = new Array[Double](epm.length)
est.fit(trainingDataset, epm, true, executionContext,
(model: Model[_], paramMapIndex: Int) => {
val paramMap = epm(paramMapIndex)
if (collectSubModelsParam) {
subModels.get(splitIndex)(paramIndex) = model
subModels.get(splitIndex)(paramMapIndex) = model
}
model
} (executionContext)
}

// Unpersist training data only when all models have trained
Future.sequence[Model[_], Iterable](modelFutures)(implicitly, executionContext)
.onComplete { _ => trainingDataset.unpersist() } (executionContext)

// Evaluate models in a Future that will calulate a metric and allow model to be cleaned up
val foldMetricFutures = modelFutures.zip(epm).map { case (modelFuture, paramMap) =>
modelFuture.map { model =>
// TODO: duplicate evaluator to take extra params from input
val metric = eval.evaluate(model.transform(validationDataset, paramMap))
logDebug(s"Got metric $metric for model trained with $paramMap.")
metric
} (executionContext)
}
foldMetrics(paramMapIndex) = metric
}
)

// Wait for metrics to be calculated before unpersisting validation dataset
val foldMetrics = foldMetricFutures.map(ThreadUtils.awaitResult(_, Duration.Inf))
validationDataset.unpersist()
foldMetrics
}.transpose.map(_.sum / $(numFolds)) // Calculate average metric over all splits
Expand Down
29 changes: 8 additions & 21 deletions mllib/src/main/scala/org/apache/spark/ml/tuning/TrainValidationSplit.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,33 +143,20 @@ class TrainValidationSplit @Since("1.5.0") (@Since("1.5.0") override val uid: St

// Fit models in a Future for training in parallel
logDebug(s"Train split with multiple sets of parameters.")
val modelFutures = epm.zipWithIndex.map { case (paramMap, paramIndex) =>
Future[Model[_]] {
val model = est.fit(trainingDataset, paramMap).asInstanceOf[Model[_]]

val metrics = new Array[Double](epm.length)
est.fit(trainingDataset, epm, true, executionContext,
(model: Model[_], paramMapIndex: Int) => {
val paramMap = epm(paramMapIndex)
if (collectSubModelsParam) {
subModels.get(paramIndex) = model
subModels.get(paramMapIndex) = model
}
model
} (executionContext)
}

// Unpersist training data only when all models have trained
Future.sequence[Model[_], Iterable](modelFutures)(implicitly, executionContext)
.onComplete { _ => trainingDataset.unpersist() } (executionContext)

// Evaluate models in a Future that will calulate a metric and allow model to be cleaned up
val metricFutures = modelFutures.zip(epm).map { case (modelFuture, paramMap) =>
modelFuture.map { model =>
// TODO: duplicate evaluator to take extra params from input
val metric = eval.evaluate(model.transform(validationDataset, paramMap))
logDebug(s"Got metric $metric for model trained with $paramMap.")
metric
} (executionContext)
}

// Wait for all metrics to be calculated
val metrics = metricFutures.map(ThreadUtils.awaitResult(_, Duration.Inf))
metrics(paramMapIndex) = metric
}
)

// Unpersist validation set once all metrics have been produced
validationDataset.unpersist()
Expand Down