注册一个模型

在本教程中,我们将利用一个本地跟踪服务器和模型注册中心以简化操作。然而,对于生产用例,我们建议使用 远程跟踪服务器

步骤 0: 安装依赖项

pip install --upgrade mlflow

步骤 1:注册模型

要使用 MLflow 模型注册表,您需要将您的 MLflow 模型添加到其中。这可以通过以下命令之一注册给定的模型来完成:

  • mlflow.<model_flavor>.log_model(registered_model_name=<model_name>): 在记录模型到跟踪服务器的同时**注册**模型。

  • mlflow.register_model(<model_uri>, <model_name>): 在将模型记录到跟踪服务器**之后**注册模型。请注意,在运行此命令之前,您必须先记录模型以获取模型URI。

MLflow 有很多模型风格。在下面的例子中,我们将利用 scikit-learn 的 RandomForestRegressor 来展示注册模型的最简单方法,但请注意,你可以利用任何 支持的模型风格。在下面的代码片段中,我们开始一个 mlflow 运行并训练一个随机森林模型。然后我们记录一些相关的超参数、模型的均方误差(MSE),最后记录并注册模型本身。

from sklearn.datasets import make_regression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split

import mlflow
import mlflow.sklearn

with mlflow.start_run() as run:
    X, y = make_regression(n_features=4, n_informative=2, random_state=0, shuffle=False)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )

    params = {"max_depth": 2, "random_state": 42}
    model = RandomForestRegressor(**params)
    model.fit(X_train, y_train)

    # Log parameters and metrics using the MLflow APIs
    mlflow.log_params(params)

    y_pred = model.predict(X_test)
    mlflow.log_metrics({"mse": mean_squared_error(y_test, y_pred)})

    # Log the sklearn model and register as version 1
    mlflow.sklearn.log_model(
        sk_model=model,
        artifact_path="sklearn-model",
        input_example=X_train,
        registered_model_name="sk-learn-random-forest-reg-model",
    )
示例输出
Successfully registered model 'sk-learn-random-forest-reg-model'.
Created version '1' of model 'sk-learn-random-forest-reg-model'.

太好了!我们已经注册了一个模型。

在继续之前,让我们强调一些重要的实现注意事项。

  • 要注册一个模型,你可以在 mlflow.sklearn.log_model() 中使用 registered_model_name 参数,或者在记录模型后调用 mlflow.register_model()。通常,我们建议前者,因为它更简洁。

  • 模型签名 为我们的模型输入和输出提供验证。在 log_model() 中的 input_example 会自动推断并记录一个签名。再次建议使用此实现,因为它简洁。