超参数调优：使用Hyperopt的SynapseML

%pip install hyperopt mlflow

from synapse.ml.core.platform import *

if running_on_synapse_internal():
    experiment_name = "hyperopt-synapseml"
elif running_on_synapse():
    experiment_name = "hyperopt-synapseml"
else:
    experiment_name = "/Shared/hyperopt-synapseml"

import mlflow

mlflow.__version__

# Set pyspark autologging logModelAllowlist to include SynapseML models
spark.sparkContext._conf.set(
    "spark.mlflow.pysparkml.autolog.logModelAllowlistFile",
    "https://mmlspark.blob.core.windows.net/publicwasb/log_model_allowlist.txt",
)

# enable autologging
mlflow.pyspark.ml.autolog()

# Set MLflow experiment.

if running_on_synapse():
    from notebookutils.mssparkutils import azureML

    linked_service = "AzureMLService1"  # use your linked service name
    ws = azureML.getWorkspace(linked_service)
    mlflow.set_tracking_uri(ws.get_mlflow_tracking_uri())
mlflow.set_experiment(experiment_name)

import numpy as np
import pandas as pd
from sklearn.datasets import fetch_california_housing
import time

try:
    california = fetch_california_housing()
except EOFError:
    print("Encountered EOFError while downloading, retrying once...")
    time.sleep(5)
    california = fetch_california_housing()

feature_cols = ["f" + str(i) for i in range(california.data.shape[1])]
header = ["target"] + feature_cols
df = spark.createDataFrame(
    pd.DataFrame(
        data=np.column_stack((california.target, california.data)), columns=header
    )
).repartition(1)

print("Dataframe has {} rows".format(df.count()))
display(df)

display(df.summary().toPandas())

from pyspark.ml.feature import VectorAssembler

# Convert features into a single vector column
featurizer = VectorAssembler(inputCols=feature_cols, outputCol="features")
data = featurizer.transform(df)["target", "features"]

train_data, test_data = data.randomSplit([0.75, 0.25], seed=42)
train_data, validation_data = train_data.randomSplit([0.85, 0.15], seed=42)

display(train_data)

# Using one partition since the training dataset is very small
repartitioned_data = train_data.repartition(1).cache()

from synapse.ml.lightgbm import LightGBMRegressor
from synapse.ml.train import ComputeModelStatistics


def train_tree(alpha, learningRate, numLeaves, numIterations):
    """
    This train() function:
     - takes hyperparameters as inputs (for tuning later)
     - returns the F1 score on the validation dataset

    Wrapping code as a function makes it easier to reuse the code later with Hyperopt.
    """
    # Use MLflow to track training.
    # Specify "nested=True" since this single model will be logged as a child run of Hyperopt's run.
    with mlflow.start_run(nested=True):

        lgr = LightGBMRegressor(
            objective="quantile",
            alpha=alpha,
            learningRate=learningRate,
            numLeaves=numLeaves,
            labelCol="target",
            numIterations=numIterations,
        )

        model = lgr.fit(repartitioned_data)

        cms = ComputeModelStatistics(
            evaluationMetric="regression", labelCol="target", scoresCol="prediction"
        )

        # Define an evaluation metric and evaluate the model on the test dataset.
        predictions = model.transform(test_data)
        metrics = cms.transform(predictions).collect()[0].asDict()

        # log metrics with mlflow
        mlflow.log_metric("MSE", metrics["mean_squared_error"])
        mlflow.log_metric("RMSE", metrics["root_mean_squared_error"])
        mlflow.log_metric("R^2", metrics["R^2"])
        mlflow.log_metric("MAE", metrics["mean_absolute_error"])

    return model, metrics["R^2"]

initial_model, val_metric = train_tree(
    alpha=0.2, learningRate=0.3, numLeaves=31, numIterations=50
)
print(
    f"The trained decision tree achieved a R^2 of {val_metric} on the validation data"
)

from hyperopt import fmin, tpe, hp, Trials, STATUS_OK


def train_with_hyperopt(params):
    """
    An example train method that calls into MLlib.
    This method is passed to hyperopt.fmin().

    :param params: hyperparameters as a dict. Its structure is consistent with how search space is defined. See below.
    :return: dict with fields 'loss' (scalar loss) and 'status' (success/failure status of run)
    """
    # For integer parameters, make sure to convert them to int type if Hyperopt is searching over a continuous range of values.
    alpha = params["alpha"]
    learningRate = params["learningRate"]
    numLeaves = int(params["numLeaves"])
    numIterations = int(params["numIterations"])

    model, r_squared = train_tree(alpha, learningRate, numLeaves, numIterations)

    # Hyperopt expects you to return a loss (for which lower is better), so take the negative of the R^2 (for which higher is better).
    loss = -r_squared

    return {"loss": loss, "status": STATUS_OK}

space = {
    "alpha": hp.uniform("alpha", 0, 1),
    "learningRate": hp.uniform("learningRate", 0, 1),
    "numLeaves": hp.uniformint("numLeaves", 30, 50),
    "numIterations": hp.uniformint("numIterations", 20, 100),
}

algo = tpe.suggest

with mlflow.start_run():
    best_params = fmin(fn=train_with_hyperopt, space=space, algo=algo, max_evals=8)

# Print out the parameters that produced the best model
best_params

best_alpha = best_params["alpha"]
best_learningRate = best_params["learningRate"]
best_numIterations = int(best_params["numIterations"])
best_numLeaves = int(best_params["numLeaves"])

final_model, val_r_squared = train_tree(
    best_alpha, best_learningRate, best_numIterations, best_numLeaves
)

# Define an evaluation metric and evaluate the model on the test dataset.
cms = ComputeModelStatistics(
    evaluationMetric="regression", labelCol="target", scoresCol="prediction"
)

initial_model_predictions = initial_model.transform(test_data)
initial_model_test_metric = (
    cms.transform(initial_model_predictions).collect()[0].asDict()["R^2"]
)

final_model_predictions = final_model.transform(test_data)
final_model_test_metric = (
    cms.transform(final_model_predictions).collect()[0].asDict()["R^2"]
)

print(
    f"On the test data, the initial (untuned) model achieved R^2 {initial_model_test_metric}, and the final (tuned) model achieved {final_model_test_metric}."
)