MLflow + Scikit-learn 入门

下载此笔记本

在本指南中，我们将向您展示如何使用 scikit-learn 训练模型，并使用 MLflow 记录您的训练。

我们将使用 Databricks 免费试用版，它内置了对 MLflow 的支持。Databricks 免费试用版提供了免费使用 Databricks 平台的机会。如果您尚未注册，请通过 此链接 注册一个帐户。

您可以从基于云的笔记本（如 Databricks 笔记本或 Google Colab）运行本指南中的代码，或者在本地计算机上运行它。

安装依赖项

让我们安装 mlflow 包。

%pip install mlflow

然后我们导入包

from sklearn.datasets import load_iris
from sklearn.linear_model import ElasticNet
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split

import mlflow

加载并准备数据集

我们将使用 鸢尾花数据集 训练一个简单的鸢尾花多分类模型。

让我们使用 load_iris() 将数据集加载到 pandas Dataframe 中，并查看数据。

iris_df = load_iris(as_frame=True).frame
iris_df

	萼片长度 (cm)	萼片宽度 (cm)	花瓣长度 (cm)	花瓣宽度 (cm)	目标
0	5.1	3.5	1.4	0.2	0
1	4.9	3.0	1.4	0.2	0
2	4.7	3.2	1.3	0.2	0
3	4.6	3.1	1.5	0.2	0
4	5.0	3.6	1.4	0.2	0
...	...	...	...	...	...
145	6.7	3.0	5.2	2.3	2
146	6.3	2.5	5.0	1.9	2
147	6.5	3.0	5.2	2.0	2
148	6.2	3.4	5.4	2.3	2
149	5.9	3.0	5.1	1.8	2

150 行 × 5 列

现在我们将把数据集分成训练集和测试集

# Split into 80% training and 20% testing
train_df, test_df = train_test_split(iris_df, test_size=0.2, random_state=42)
train_df.shape, test_df.shape

((120, 5), (30, 5))

# Separate the target column for the training set
train_dataset = mlflow.data.from_pandas(train_df, name="train")
train_x = train_dataset.df.drop(["target"], axis=1)
train_y = train_dataset.df[["target"]]

train_x.shape, train_y.shape

((120, 4), (120, 1))

# Separate the target column for the testing set
test_dataset = mlflow.data.from_pandas(test_df, name="test")
test_x = test_dataset.df.drop(["target"], axis=1)
test_y = test_dataset.df[["target"]]

test_x.shape, test_y.shape

((30, 4), (30, 1))

定义模型

对于此示例，我们将使用具有一些预定义超参数的 ElasticNet 模型。我们还将定义一个辅助函数来计算一些指标以评估我们模型的性能。

lr = ElasticNet(alpha=0.5, l1_ratio=0.5, random_state=42)

def compute_metrics(actual, predicted):
  rmse = mean_squared_error(actual, predicted)
  mae = mean_absolute_error(actual, predicted)
  r2 = r2_score(actual, predicted)

  return rmse, mae, r2

连接到 MLflow 跟踪服务器

在训练之前，我们需要配置 MLflow 跟踪服务器，因为我们将数据记录到 MLflow 中。在本教程中，我们将使用 Databricks 免费试用版作为 MLflow 跟踪服务器。有关其他选项，例如使用本地 MLflow 服务器，请阅读 跟踪服务器概述。

如果您尚未设置 Databricks 免费试用版的帐户和访问令牌，请按照 本指南 进行操作。注册时间应不超过 5 分钟。对于本指南，我们需要 ML 实验仪表板来跟踪我们的训练进度。

在 Databricks 免费试用版上成功注册帐户后，让我们将 MLflow 连接到 Databricks Workspace。您将需要输入以下信息

• Databricks 主机：https://<您的工作区主机>.cloud.databricks.com
• 令牌：您的个人访问令牌

mlflow.login()

现在此笔记本已连接到托管跟踪服务器。让我们配置一些 MLflow 元数据。需要设置两件事

• mlflow.set_tracking_uri：始终使用“databricks”。
• mlflow.set_experiment：选择您喜欢的名称，以 /Users/<您的电子邮件>/ 开头。

mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Users/<your email>/mlflow-sklearn-quickstart")

使用 MLflow 进行日志记录

MLflow 具有强大的跟踪 API，可让我们记录运行和模型及其关联的元数据，例如参数和指标。让我们首先开始训练运行以训练我们的模型。

# Start a training run
with mlflow.start_run() as training_run:
  # Log the parameters for our model
  mlflow.log_param("alpha", 0.5)
  mlflow.log_param("l1_ratio", 0.5)

  # Train and log our model, which inherits the parameters
  lr.fit(train_x, train_y)
  model_info = mlflow.sklearn.log_model(sk_model=lr, name="elasticnet", input_example=train_x)

  # Evaluate the model on the training dataset and log metrics
  # These metrics will be linked to both the model and run
  predictions = lr.predict(train_x)
  (rmse, mae, r2) = compute_metrics(train_y, predictions)
  mlflow.log_metrics(
      metrics={
          "rmse": rmse,
          "r2": r2,
          "mae": mae,
      },
      dataset=train_dataset,
  )

现在让我们在测试数据集上评估我们的模型

# Start an evaluation run
with mlflow.start_run() as evaluation_run:
  # Load our previous model
  logged_model = mlflow.sklearn.load_model(f"models:/{model_info.model_id}")

  # Evaluate the model on the training dataset and log metrics
  predictions = logged_model.predict(test_x)
  (rmse, mae, r2) = compute_metrics(test_y, predictions)
  mlflow.log_metrics(
      metrics={
          "rmse": rmse,
          "r2": r2,
          "mae": mae,
      },
      dataset=test_dataset,
      model_id=model_info.model_id,
  )

查看结果

让我们看看我们的训练和测试结果。登录您的 Databricks 工作区，然后单击左侧菜单中的 Experiments 选项卡。初始页面显示运行列表，我们可以在其中看到我们的训练和评估运行。

现在让我们前往模型选项卡，在那里我们可以看到我们记录的模型

单击模型名称将带您进入模型详细信息页面，其中包含其参数、两次运行的指标以及其他元数据的信息。

我们还可以使用 API 检查我们的模型

logged_model = mlflow.get_logged_model(model_info.model_id)

logged_model, logged_model.metrics, logged_model.params