记录 MLflow 模型

本文描述了如何将已定型的模型(或项目)记录为 MLflow 模型。 它探索自定义 MLflow 如何打包模型以及如何运行这些模型的不同方法。

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

MLflow 中从项目到模型描述了与记录 MLflow 模型相比,记录项目或文件之间的差异。

MLflow 模型也是项目。 然而,该模型有特定的结构,作为模型创造者和打算使用模型之人之间的合同。 此合同有助于在项目本身和其含义之间建立桥。

模型记录具有以下优点:

  • 可以使用 mlflow.<flavor>.load_model 直接加载模型进行推理,并且可以使用 predict 函数
  • 管道输入可以直接使用模型
  • 可以在不指示评分脚本或环境的情况下部署模型
  • Swagger 在部署的终结点中自动启用,Azure 机器学习工作室可以使用测试功能
  • 可以使用负责任 AI 仪表板

本节介绍如何在 Azure 机器学习中使用 MLflow 的模型概念:

使用 autolog 记录模型

可以使用 MLflow 自动记录功能。 自动记录允许 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 方法可以手动记录模型。 此工作流可以控制模型记录的不同方面。

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

  • 希望指示不同于自动检测到事项的 pip 包或 conda 环境
  • 希望包括输入示例
  • 希望在所需的包中包括特定项目
  • 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)

# Signature
signature = infer_signature(X_test, y_test)

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

# Sample
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 模型记录。 MLflow 模型的自动记录稍后会作为手动进程进行
  • 使用 infer_signature 方法尝试直接从输入和输出推断签名
  • mlflow.utils.environment._mlflow_conda_env 方法是 MLflow SDK 中的专用方法。 在本例中,它使代码更简单,但使用时请谨慎。 未来可能会发生变化。 作为替代方案,可以将 YAML 定义手动生成为 Python 字典。

使用 predict 方法记录具有不同行为的模型

当使用 mlflow.autologmlflow.<flavor>.log_model 记录模型时,模型风格决定了如何执行推理以及模型返回的内容。 MLflow 没有强制执行任何关于生成 predict 结果的特定行为。 在某些情况下,可能会希望在模型执行前后进行一些预处理或后处理。

在这种情况下,请实现直接从输入移动到输出的机器学习管道。 尽管这种实现是可能的,有时也被鼓励来提高性能,但实现起来可能会很有挑战性。 在这些情况下,它可以帮助自定义模型处理推理的方式,如下一节所述。

记录自定义模型

MLflow 支持许多机器学习框架,包括

  • CatBoost
  • FastAI
  • h2o
  • Keras
  • LightGBM
  • MLeap
  • MXNet Gluon
  • ONNX
  • Prophet
  • 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() 将模型记录在 MLflowv 运行中。

如果在现有模型对象周围创建包装器,那么为自定义模型创建风格就变得十分简单了。 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):
        # You don't have to keep the semantic meaning of `predict`. You can use here model.recommend(), model.forecast(), etc
        return self._model.predict_proba(data)

    # You can even add extra functions if you need to. Since the model is serialized,
    # all of them will be available when you load your model back.
    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 来推断签名。 目标列具有目标类,但模型现在返回每个类的两个概率。

后续步骤