通过

记录 MLflow 模型

本文介绍如何将训练的机器学习模型或项目记录为 MLflow 模型。 MLflow 是用于管理机器学习工作流的开源框架。 本文探讨用于自定义 MLflow 包和运行模型的方式的各种选项。

先决条件

  • MLflow SDK mlflow 包

为什么记录模型而不是项目？

MLflow 模型是一种工件。 但是，模型具有一个特定的结构，该结构充当创建模型的人员与打算使用它的人员之间的协定。 此合同有助于在项目本身和其含义之间建立桥。

有关日志记录项目或文件与日志记录 MLflow 模型之间的差异，请参阅 MLflow 中的项目和模型。

可以将模型的文件记录为项目，但模型日志记录功能具有以下优势：

  • 可用于 mlflow.<flavor>.load_model 直接加载模型进行推理，并且可以使用函数 predict 。
  • 管道输入可以直接使用模型。
  • 无需指定评分脚本或环境即可部署模型。
  • 在部署的终结点中，Swagger 自动启用。 因此，可以使用 Azure 机器学习工作室中的测试功能测试模型。
  • 可以使用负责任 AI 仪表板。 有关详细信息，请参阅 Azure 机器学习工作室中的“使用负责任的 AI”仪表板。

使用自动日志记录记录模型

可以使用 MLflow autolog 功能自动记录模型。 使用自动日志记录时，MLflow 会捕获框架中的所有相关指标、参数、项目和模型。 记录的数据取决于框架。 默认情况下，如果启用自动日志记录，则会记录大多数模型。 在某些情况下，某些版本不会保存模型。 例如，PySpark 风格不会记录超过特定大小的模型。

使用mlflow.autolog或mlflow.<flavor>.autolog来激活自动日志记录。 以下代码使用 autolog 记录使用 XGBoost 训练的分类器模型：

import mlflow
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score

mlflow.autolog()

model = XGBClassifier(use_label_encoder=False, eval_metric="logloss")
model.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)

y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)

提示

如果使用机器学习管道，例如 Scikit-Learn 管道，请使用该管道风格的 autolog 功能来记录模型。 在管道对象上调用该方法时 fit ，会自动运行模型日志记录。 有关记录模型（包括预处理）和使用管道的笔记本，请参阅使用 MLflow 训练和跟踪 XGBoost 分类器。

使用自定义签名、环境或示例的日志模型

可以使用 MLflow mlflow.<flavor>.log_model 方法手动记录模型。 此工作流提供对模型日志记录的各个方面的控制。

在以下情况下使用此方法：

  • 你想要指示与自动检测到的包或环境不同的 Conda 环境或 pip 包。
  • 你想要包括输入示例。
  • 希望在所需的包中包括特定项目。
  • 该方法 autolog 无法正确推断签名。 使用张量输入时，会出现这种情况，这要求签名具有特定形状。
  • autolog 方法不满足您的所有需求。

以下代码记录 XGBoost 分类器模型：

import mlflow
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score
from mlflow.models import infer_signature
from mlflow.utils.environment import _mlflow_conda_env

mlflow.autolog(log_models=False)

model = XGBClassifier(use_label_encoder=False, eval_metric="logloss")
model.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)
y_pred = model.predict(X_test)

accuracy = accuracy_score(y_test, y_pred)

# Infer the signature.
signature = infer_signature(X_test, y_test)

# Set up a Conda environment.
custom_env =_mlflow_conda_env(
    additional_conda_deps=None,
    additional_pip_deps=["xgboost==1.5.2"],
    additional_conda_channels=None,
)

# Sample the data.
input_example = X_train.sample(n=1)

# Log the model manually.
mlflow.xgboost.log_model(model, 
                         artifact_path="classifier", 
                         conda_env=custom_env,
                         signature=signature,
                         input_example=input_example)

注意

  • 对 autolog 的调用使用 log_models=False 配置。 此设置关闭自动 MLflow 模型日志记录。 稍后使用 log_model 方法手动记录模型。
  • 该方法 infer_signature 用于尝试直接从输入和输出推断签名。
  • mlflow.utils.environment._mlflow_conda_env 方法是 MLflow SDK 中的专用方法。 在此示例中，它简化了代码。 但请谨慎使用此方法，因为它将来可能会更改。 作为替代方案，可以将 YAML 定义手动生成为 Python 字典。

使用修改后的预测行为的日志模型

使用 mlflow.autolog 或 mlflow.<flavor>.log_model 记录模型时，模型风格决定如何执行推理。 口味还决定模型的返回结果。 MLflow 不会强制实施有关生成 predict 结果的特定行为。 在某些情况下，可能需要预处理或后处理您的数据。

在这种情况下，可以实现直接从输入移动到输出的机器学习管道。 虽然这种类型的实现有时可以提高性能，但实现可能非常困难。 在这种情况下，自定义模型处理推理的方式会很有帮助。 有关详细信息，请参阅下一部分， 记录自定义模型。

记录自定义模型

MLflow 支持许多机器学习框架，包括以下风格：

  • CatBoost
  • FastAI
  • H₂O
  • Keras
  • LightGBM
  • MLeap
  • ONNX
  • 先知
  • PyTorch
  • scikit-learn
  • spaCy
  • Spark MLlib
  • statsmodels
  • TensorFlow
  • XGBoost

有关完整列表，请参阅内置模型版本。

但是，可能需要更改模型版本的运作方式，或记录 MLflow 原生不支持的模型。 或者，可能需要记录使用各种框架中的多个元素的模型。 在这些情况下，可以创建自定义模型风格。

为了解决此问题，MLflow 提供了 PyFunc 风格，这是 Python 模型的默认模型接口。 只要该对象满足两个条件，此模式即可将任何对象记录为模型：

  • 您至少要实现 predict 方法。
  • Python 对象继承自 mlflow.pyfunc.PythonModel 类。

提示

实现 scikit-learn API 的可序列化模型可使用 scikit-learn 风格来记录模型，无论该模型是否由 scikit-learn 构建。 如果可以采用 Pickle 格式保存模型，并且该对象至少具有 predict 和 predict_proba 方法，则可以使用 mlflow.sklearn.log_model 在 MLflow 运行中记录模型。

为自定义模型创建风格的最简单方法是围绕现有模型对象创建包装器。 MLflow 为你序列化和打包模型。 当对象可以作为文件（通常为 Pickle 格式）存储在文件系统中时，Python 对象是可串行化的。 在运行时，可以从该文件加载对象。 加载会还原保存时可用的所有值、属性和方法。

在以下情况下使用此方法：

  • 可以使用 Pickle 格式序列化模型。
  • 希望保留模型在刚刚完成训练后的状态。
  • 希望自定义 predict 函数的工作方式。

以下代码包装使用 XGBoost 创建的模型，使其行为与 XGBoost 风格默认实现不同。 它返回概率而不是类。

from mlflow.pyfunc import PythonModel, PythonModelContext

class ModelWrapper(PythonModel):
    def __init__(self, model):
        self._model = model

    def predict(self, context: PythonModelContext, data):
        # The next line uses a prediction function. However, you could also use model.recommend(), model.forecast(), or a similar function instead.
        return self._model.predict_proba(data)

    # You can add extra functions if you need to. Because the model is serialized,
    # all of them are available when you load your model.
    def predict_batch(self, data):
        pass

运行期间使用以下代码记录自定义模型：

import mlflow
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score
from mlflow.models import infer_signature

mlflow.xgboost.autolog(log_models=False)

model = XGBClassifier(use_label_encoder=False, eval_metric="logloss")
model.fit(X_train, y_train, eval_set=[(X_test, y_test)], verbose=False)
y_probs = model.predict_proba(X_test)

accuracy = accuracy_score(y_test, y_probs.argmax(axis=1))
mlflow.log_metric("accuracy", accuracy)

signature = infer_signature(X_test, y_probs)
mlflow.pyfunc.log_model("classifier", 
                        python_model=ModelWrapper(model),
                        signature=signature)

提示

在前面的代码中，infer_signature 方法使用 y_probs 推断签名。 目标列包含目标类，但模型为每个类返回两个概率。

后续步骤