快速入门:比较运行,选择模型,并将其部署到REST API

在这个快速入门中,您将:

  • 在训练脚本上运行超参数扫描

  • 在MLflow UI中比较运行的结果

  • 选择最佳运行并将其注册为模型

  • 将模型部署到REST API

  • 构建一个适合部署到云平台的容器镜像

作为一名机器学习工程师或MLOps专业人员,您可以使用MLflow来比较、共享和部署团队生成的最佳模型。在本快速入门中,您将使用MLflow跟踪UI来比较超参数扫描的结果,选择最佳运行,并将其注册为模型。然后,您将模型部署到REST API。最后,您将创建一个适合部署到云平台的Docker容器镜像。

显示数据科学和MLOps工作流程的图表,使用MLflow

设置

要获取关于如何设置一个MLflow环境以配置MLflow跟踪功能的全面指南,您可以 阅读这里的指南

运行超参数扫描

这个示例尝试优化 Keras 深度学习模型在葡萄酒质量数据集上的 RMSE 指标。它有两个超参数需要优化:learning_ratemomentum。我们将使用 Hyperopt 库在 learning_ratemomentum 的不同值上运行超参数扫描,并将结果记录在 MLflow 中。

在运行超参数扫描之前,让我们将 MLFLOW_TRACKING_URI 环境变量设置为我们的 MLflow 跟踪服务器的 URI:

export MLFLOW_TRACKING_URI=http://localhost:5000

备注

如果您想探索其他跟踪服务器部署的可能性,包括使用 Databricks Community Edition 的全托管免费解决方案,请参阅 此页面

导入以下包

import keras
import numpy as np
import pandas as pd
from hyperopt import STATUS_OK, Trials, fmin, hp, tpe
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split

import mlflow
from mlflow.models import infer_signature

现在加载数据集并将其分割为训练集、验证集和测试集。

# Load dataset
data = pd.read_csv(
    "https://raw.githubusercontent.com/mlflow/mlflow/master/tests/datasets/winequality-white.csv",
    sep=";",
)

# Split the data into training, validation, and test sets
train, test = train_test_split(data, test_size=0.25, random_state=42)
train_x = train.drop(["quality"], axis=1).values
train_y = train[["quality"]].values.ravel()
test_x = test.drop(["quality"], axis=1).values
test_y = test[["quality"]].values.ravel()
train_x, valid_x, train_y, valid_y = train_test_split(
    train_x, train_y, test_size=0.2, random_state=42
)
signature = infer_signature(train_x, train_y)

然后让我们定义模型架构并训练模型。train_model 函数使用 MLflow 来跟踪每个试验的参数、结果和模型本身作为子运行。

def train_model(params, epochs, train_x, train_y, valid_x, valid_y, test_x, test_y):
    # Define model architecture
    mean = np.mean(train_x, axis=0)
    var = np.var(train_x, axis=0)
    model = keras.Sequential(
        [
            keras.Input([train_x.shape[1]]),
            keras.layers.Normalization(mean=mean, variance=var),
            keras.layers.Dense(64, activation="relu"),
            keras.layers.Dense(1),
        ]
    )

    # Compile model
    model.compile(
        optimizer=keras.optimizers.SGD(
            learning_rate=params["lr"], momentum=params["momentum"]
        ),
        loss="mean_squared_error",
        metrics=[keras.metrics.RootMeanSquaredError()],
    )

    # Train model with MLflow tracking
    with mlflow.start_run(nested=True):
        model.fit(
            train_x,
            train_y,
            validation_data=(valid_x, valid_y),
            epochs=epochs,
            batch_size=64,
        )
        # Evaluate the model
        eval_result = model.evaluate(valid_x, valid_y, batch_size=64)
        eval_rmse = eval_result[1]

        # Log parameters and results
        mlflow.log_params(params)
        mlflow.log_metric("eval_rmse", eval_rmse)

        # Log model
        mlflow.tensorflow.log_model(model, "model", signature=signature)

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

objective 函数接收超参数并返回该组超参数下 train_model 函数的结果。

def objective(params):
    # MLflow will track the parameters and results for each run
    result = train_model(
        params,
        epochs=3,
        train_x=train_x,
        train_y=train_y,
        valid_x=valid_x,
        valid_y=valid_y,
        test_x=test_x,
        test_y=test_y,
    )
    return result

接下来,我们将为 Hyperopt 定义搜索空间。在这种情况下,我们希望尝试不同的 learning-ratemomentum 值。Hyperopt 通过选择一组初始的超参数开始其优化过程,通常是随机选择的或基于指定的域空间。此域空间定义了每个超参数的可能值的范围和分布。在评估初始设置后,Hyperopt 使用结果来更新其概率模型,以更明智的方式指导后续超参数集的选择,旨在收敛到最优解。

space = {
    "lr": hp.loguniform("lr", np.log(1e-5), np.log(1e-1)),
    "momentum": hp.uniform("momentum", 0.0, 1.0),
}

最后,我们将使用 Hyperopt 运行超参数扫描,传入 objective 函数和搜索空间。Hyperopt 将尝试不同的超参数组合并返回最佳组合的结果。我们将把最佳参数、模型和评估指标存储在 MLflow 中。

mlflow.set_experiment("/wine-quality")
with mlflow.start_run():
    # Conduct the hyperparameter search using Hyperopt
    trials = Trials()
    best = fmin(
        fn=objective,
        space=space,
        algo=tpe.suggest,
        max_evals=8,
        trials=trials,
    )

    # Fetch the details of the best run
    best_run = sorted(trials.results, key=lambda x: x["loss"])[0]

    # Log the best parameters, loss, and model
    mlflow.log_params(best)
    mlflow.log_metric("eval_rmse", best_run["loss"])
    mlflow.tensorflow.log_model(best_run["model"], "model", signature=signature)

    # Print out the best parameters and corresponding loss
    print(f"Best parameters: {best}")
    print(f"Best eval rmse: {best_run['loss']}")

比较结果

在浏览器中打开 MLFLOW_TRACKING_URI 的 MLflow UI。你应该会看到一个嵌套的运行列表。在默认的 表格视图 中,选择 按钮并添加 指标 | eval_rmse 列以及 参数 | lr参数 | momentum 列。要按 RMSE 升序排序,请点击 eval_rmse 列标题。最佳运行通常在 测试 数据集上的 RMSE 约为 0.70。你可以在 参数 列中看到最佳运行的参数。

MLflow 跟踪 UI 表格视图显示运行的截图

选择 图表视图。选择 平行坐标图 并配置它以显示 lrmomentum 坐标以及 eval_rmse 指标。图中的每条线代表一次运行,并将每次超参数评估运行的参数与该运行的评估误差指标关联起来。

Screenshot of MLflow tracking UI parallel coordinates graph showing runs

这个图表上的红色图形是表现不佳的运行。最低的一个是基线运行,其中 lrmomentum 都设置为 0.0。该基线运行的 RMSE 约为 0.89。其他红色线条显示,高 momentum 也可能导致此问题和架构下的结果不佳。

颜色偏向蓝色的图表表示运行效果更好。将鼠标悬停在单个运行上以查看其详细信息。

注册你的最佳模型

选择最佳运行并将其注册为模型。在 表格视图 中,选择最佳运行。在 运行详情 页面,打开 工件 部分并选择 注册模型 按钮。在 注册模型 对话框中,输入模型的名称,例如 wine-quality,然后点击 注册

现在,您的模型已准备好部署。您可以在 MLflow UI 的 模型 页面中查看它。打开您刚刚注册的模型页面。

您可以为模型添加描述,添加标签,并轻松导航回生成此模型的源运行。您还可以将模型转换到不同的阶段。例如,您可以将模型转换到 Staging 以表明它已准备好进行测试。您可以将其转换到 Production 以表明它已准备好部署。

通过选择 Stage 下拉菜单,将模型转换到 Staging 阶段:

MLflow 跟踪 UI 模型页面截图显示已注册的模型

在本地提供模型服务

MLflow 允许你轻松地服务由任何运行或模型版本生成的模型。你可以通过运行以下命令来服务你刚刚注册的模型:

mlflow models serve -m "models:/wine-quality/1" --port 5002

(请注意,如果您在默认端口 5000 上在同一台机器上运行跟踪服务器,则如上指定端口将是必要的。)

你也可以使用 runs:/<run_id> URI 来提供模型,或者使用 Artifact Store 是一个用于跟踪和管理项目中产生的各种工件(如模型、数据集等)的存储系统。 中描述的任何支持的 URI。

请注意,在生产环境中,我们不建议将模型与跟踪服务器部署在同一虚拟机中,因为资源有限,在本指南中,我们只是为了简单起见,从同一台机器上运行所有内容。

要测试模型,您可以使用 curl 命令向 REST API 发送请求:

curl -d '{"dataframe_split": {
"columns": ["fixed acidity","volatile acidity","citric acid","residual sugar","chlorides","free sulfur dioxide","total sulfur dioxide","density","pH","sulphates","alcohol"],
"data": [[7,0.27,0.36,20.7,0.045,45,170,1.001,3,0.45,8.8]]}}' \
-H 'Content-Type: application/json' -X POST localhost:5002/invocations

推理是通过向 localhost 上指定端口的 invocations 路径发送 JSON POST 请求完成的。columns 键指定了输入数据中列的名称。data 值是一个列表的列表,其中每个内部列表是一行数据。为简洁起见,上述仅请求了一次葡萄酒质量的预测(评分范围为3-8)。响应是一个 JSON 对象,包含一个 predictions 键,该键包含一个预测列表,每个数据行对应一个预测。在这种情况下,响应为:

{"predictions": [{"0": 5.310967445373535}]}

输入和输出的模式可以在 MLflow UI 的 Artifacts | Model 描述中找到。该模式之所以可用,是因为 train.py 脚本使用了 mlflow.infer_signature 方法,并将结果传递给了 mlflow.log_model 方法。强烈建议将签名传递给 log_model 方法,因为它可以在输入请求格式错误时提供清晰的错误信息。

为您的模型构建一个容器镜像

大多数部署路径将使用容器来打包您的模型、其依赖项以及运行时环境的相关部分。您可以使用 MLflow 为您的模型构建 Docker 镜像。

mlflow models build-docker --model-uri "models:/wine-quality/1" --name "qs_mlops"

此命令构建一个名为 qs_mlops 的 Docker 镜像,该镜像包含您的模型及其依赖项。在这种情况下,model-uri 指定的是版本号(/1)而不是生命周期阶段(/staging),但您可以使用最适合您工作流的任何一种。构建镜像需要几分钟时间。一旦完成,您可以在本地、本地部署、定制的互联网服务器或云平台上运行该镜像以提供实时推理服务。您可以在本地运行它,命令如下:

docker run -p 5002:8080 qs_mlops

这个 Docker run 命令 运行你刚刚构建的镜像,并将你本地机器上的端口 5002 映射到容器中的端口 8080。你现在可以使用与之前相同的 curl 命令向模型发送请求:

curl -d '{"dataframe_split": {"columns": ["fixed acidity","volatile acidity","citric acid","residual sugar","chlorides","free sulfur dioxide","total sulfur dioxide","density","pH","sulphates","alcohol"], "data": [[7,0.27,0.36,20.7,0.045,45,170,1.001,3,0.45,8.8]]}}' -H 'Content-Type: application/json' -X POST localhost:5002/invocations

部署到云平台

几乎所有的云平台都允许你部署一个 Docker 镜像。这个过程差异很大,所以你需要查阅你的云服务提供商的文档以获取详细信息。

此外,一些云服务提供商已经内置了对 MLflow 的支持。例如:

所有支持 MLflow。云平台通常支持多种部署工作流程:命令行、基于 SDK 的和基于 Web 的。您可以在这些工作流程中的任何一个中使用 MLflow,尽管细节会因平台和版本而异。再次提醒,您需要查阅云提供商的文档以获取详细信息。