记录 MLflow 模型

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

先决条件

MLflow SDK mlflow 包（版本 2.16.x 或更早版本）
Azure Machine Learning MLflow 插件 azureml-mlflow
包 xgboost （版本 2.1.1 或更低版本）
scikit-learn 包

重要

Azure Machine Learning与 MLflow 2.16 及更早版本兼容。 MLflow 2.17 及更高版本引入的工件存储库和 LoggedModels API 变更目前不受 azureml-mlflow 插件支持。锁定 MLflow 和 XGBoost 版本以避免兼容性问题。

pip install "mlflow<=2.16.2" azureml-mlflow "xgboost<=2.1.1" scikit-learn

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

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

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

import mlflow
from xgboost import XGBClassifier
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# Load sample data and split into train/test sets.
X, y = load_breast_cancer(return_X_y=True, as_frame=True)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

mlflow.autolog()

model = XGBClassifier(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)

参考：mlflow.autolog | XGBClassifier | accuracy_score

注释

本文中的其余代码示例重复使用X_train前面的示例数据设置中的变量X_testy_train和y_test变量。

小窍门

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

为什么选用日志模型而不是工件？

记录模型而非原始工件，为模型创建者和使用者之间提供了一个结构化的协议。有关详细比较，请参阅 MLflow 中的 Artifacts 和模型。

模型日志记录具有以下优势：

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

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

使用 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

mlflow.autolog(log_models=False)

with mlflow.start_run():
    model = XGBClassifier(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)

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

    # Log the model manually.
    mlflow.xgboost.log_model(model, 
                             artifact_path="classifier", 
                             pip_requirements=["xgboost>=2.0"],
                             signature=signature,
                             input_example=input_example)

参考：mlflow.xgboost.log_model | infer_signature

注释

对 autolog 的调用使用 log_models=False 配置。此设置关闭自动 MLflow 模型日志记录。稍后使用 log_model 方法手动记录模型。
该方法 infer_signature 尝试直接从输入和输出推断签名。
该 pip_requirements 参数指定模型所需的包。或者，可以将 conda_env 参数与字典或 YAML 文件一起使用。

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

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

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

记录自定义模型

MLflow 支持许多machine learning框架，包括 CatBoost、Keras、LightGBM、ONNX、PyTorch、scikit-learn、Spark MLlib、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, params=None):
        # 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)

with mlflow.start_run():
    model = XGBClassifier(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(artifact_path="classifier", 
                            python_model=ModelWrapper(model),
                            signature=signature)

参考：mlflow.pyfunc.log_model | PythonModel

小窍门

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

有时，模型由需要加载的多个部分组成。有时，无法将模型序列化为 Pickle 文件。在这些情况下，请使用类 PythonModel 。它支持任意工件列表。将每个工件与模型一起打包。

在以下情况下使用此技术：

不能以 Pickle 格式串行化模型，或者有更好的串行化格式可用。
模型包含一个或多个必须引用的构件以加载模型。
你想要保留一些推理配置属性，例如要推荐的项数。
你想要自定义模型加载方式以及 predict 函数的工作原理。

以下代码演示如何记录使用artifacts的自定义模型：

encoder_path = 'encoder.pkl'
joblib.dump(encoder, encoder_path)

model_path = 'xgb.model'
model.save_model(model_path)

mlflow.pyfunc.log_model(artifact_path="classifier", 
                        python_model=ModelWrapper(),
                        artifacts={ 
                            'encoder': encoder_path,
                            'model': model_path 
                        },
                        signature=signature)

注释

模型未以 Pickle 格式保存。相反，代码使用所使用的框架的 save 方法保存模型。
类 ModelWrapper 包装模型，但模型不会作为参数 ModelWrapper 传递给构造函数。
log_model 方法具有字典参数，artifacts。其键包含项目名称。每个值都包含本地文件系统中一个项目的路径。

对应的模型封装器类似于以下代码：

from mlflow.pyfunc import PythonModel, PythonModelContext

class ModelWrapper(PythonModel):
    def load_context(self, context: PythonModelContext):
        import pickle
        from xgboost import XGBClassifier
        from sklearn.preprocessing import OrdinalEncoder
        
        self._encoder = pickle.loads(context.artifacts["encoder"])
        self._model = XGBClassifier(eval_metric="logloss")
        self._model.load_model(context.artifacts["model"])

    def predict(self, context: PythonModelContext, data, params=None):
        return self._model.predict_proba(data)

完整的训练例程类似于以下代码：

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

mlflow.xgboost.autolog(log_models=False)

with mlflow.start_run():
    encoder = OrdinalEncoder(
        handle_unknown='use_encoded_value', unknown_value=np.nan
    )
    X_train['thal'] = encoder.fit_transform(
        X_train['thal'].to_frame()
    )
    X_test['thal'] = encoder.transform(
        X_test['thal'].to_frame()
    )

    model = XGBClassifier(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)

    encoder_path = 'encoder.pkl'
    joblib.dump(encoder, encoder_path)
    model_path = "xgb.model"
    model.save_model(model_path)

    signature = infer_signature(X, y_probs)
    mlflow.pyfunc.log_model(
        artifact_path="classifier", 
        python_model=ModelWrapper(),
        artifacts={ 
            'encoder': encoder_path,
            'model': model_path 
        },
        signature=signature
    )

模型有时具有复杂的逻辑，或者在推理时加载多个源文件。例如，如果模型有 Python 库，则会出现这种情况。在此方案中，将库与模型一起打包，以便它们充当单个部分。

在以下情况下使用此技术：

不能以 Pickle 格式串行化模型，或者有更好的串行化格式可用。
可以将你的模型工件存储在一个存储所有必需工件的文件夹中。
模型源代码很复杂，它需要多个 Python 文件。可能有一个库支持你的模型。
你想要自定义模型加载方式以及函数的运行方式 predict 。

MLflow 支持这些类型的模型。使用 MLflow 时，只要源代码具有 加载程序模块，就可以指定任意源代码来打包模型。可以使用参数在调用log_modelloader_module中指定加载程序模块，该参数指示实现加载程序的 Python 命名空间。还需要该 code_paths 参数。它提供定义加载程序模块的源文件。在此命名空间中，必须实现接收artifacts路径的 _load_pyfunc(data_path: str) 函数，并返回至少实现 predict 方法的对象。

model_path = 'xgb.model'
model.save_model(model_path)

mlflow.pyfunc.log_model(artifact_path="classifier", 
                        data_path=model_path,
                        code_paths=['src'],
                        loader_module='loader_module',
                        signature=signature)

注释

模型未以 Pickle 格式保存。相反，代码使用所使用的框架的 save 方法保存模型。
log_model 方法具有一个 data_path 参数，该参数指向包含模型工件的文件夹。这些工件可以放置在一个文件夹或文件中。工件与模型一起打包。
该方法 log_model 具有指向 code_paths 源代码位置的参数。该值 code_paths 可以是路径或单个文件。源代码与模型一起打包。
函数 _load_pyfunc 存储在 Python 模块中 loader_module 。

src 文件夹包含 loader_module.py 文件。该文件充当加载程序模块，包含以下行：

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

    def predict(self, data):
        return self._model.predict_proba(data)

def _load_pyfunc(data_path: str):
    import os

    model = XGBClassifier(eval_metric='logloss')
    model.load_model(os.path.abspath(data_path))

    return MyModel(model)

注释

MyModel 类没有从 PythonModel 继承，如前所示。但是，它具有一个 predict 函数。
模型源代码在文件中。任何源代码都适合。在这种情况下，名为 src 的文件夹是理想的文件夹。
函数 _load_pyfunc 返回模型的类的实例。