Train AutoML Model in Sparkling Water¶

Sparkling Water provides API in Scala and Python for H2O AutoML. The following sections describe how to train an AutoML model in Sparkling Water in both languages.

Scala

First, let’s start Sparkling Shell as

./bin/sparkling-shell

Start H2O cluster inside the Spark environment

import org.apache.spark.h2o._
import java.net.URI
val hc = H2OContext.getOrCreate()

Parse the data using H2O and convert them to Spark Frame

val frame = new H2OFrame(new URI("https://raw.githubusercontent.com/h2oai/sparkling-water/master/examples/smalldata/prostate/prostate.csv"))
val sparkDF = hc.asSparkFrame(frame).withColumn("CAPSULE", $"CAPSULE" cast "string")
val Array(trainingDF, testingDF) = sparkDF.randomSplit(Array(0.8, 0.2))

Create an H2OAutoML instance and configure it according to your use case via provided setters. If feature columns are not specified explicitly, all columns excluding label, fold, weight, and ignored columns are considered as features.

import ai.h2o.sparkling.ml.algos.H2OAutoML
val automl = new H2OAutoML()
automl.setLabelCol("CAPSULE")
automl.setIgnoredCols(Array("ID"))

By default, AutoML goes through a huge space of H2O algorithms and their hyper-parameters which requires some time. If you wish to speed up the training phase, you can exclude some H2O algorithms and limit the number of trained models.

automl.setExcludeAlgos(Array("GLM"))
automl.setMaxModels(10)

Train the AutoML model. The training phase returns the best model according to the sortMetric. If it’s not specified, the sortMetric is chosen automatically.

automl.setSortMetric("AUC")
val model = automl.fit(trainingDF)

You can also get raw model details by calling the getModelDetails() method available on the model as:

model.getModelDetails()

Run Predictions

model.transform(testingDF).show(false)

If you are curious to see information about other models created during the AutoML training process, you can get a model leaderboard represented by Spark DataFrame.

val leaderboard = automl.getLeaderboard()
leaderboard.show(false)

By default, the leaderboard contains the model name (model_id) and various performance metrics like AUC, RMSE, etc. If you want to see more information about models, you can add extra columns to the leaderboard by passing column names to the getLeaderboard() method.

val leaderboard = automl.getLeaderboard("training_time_ms", "predict_time_per_row_ms")
leaderboard.show(false)

Extra columns don’t have to be specified explicitly. You can specify the addition of all possible extra columns as:

val leaderboard = automl.getLeaderboard("ALL")
leaderboard.show(false)

Python

First, let’s start PySparkling Shell as

./bin/pysparkling