在 MLOps 中，一个常见的挑战是在各种算法或框架之间迁移的麻烦。这篇新手易懂的文章将帮助你利用 mlflow.pyfunc 构建与算法无关的模型，从而应对这一挑战。

为什么要进行算法无关模型构建？

考虑一下这种情况：我们目前在生产中部署了一个 sklearn 模型，用于特定的用例。后来，我们发现深度学习模型的性能更好。如果 sklearn 模型是以其原始格式部署的，那么过渡到深度学习模型可能会很麻烦，因为这两个模型的工件截然不同。

为了应对这一挑战，mlflow.pyfunc 模型风味提供了一种通用的方法，用于在 Python 中构建和部署机器学习模型。

1. 通用模型构建： pyfunc 模型风味提供了一种构建模型的统一方法，而不管构建时使用的是什么框架或库。

2. ML Pipeline 的封装：pyfunc 允许我们在模型消费过程中将模型封装为前处理和后处理步骤或其他所需的自定义逻辑。

3. 统一模型表示法： 我们可以使用 pyfunc 部署模型、机器学习管道或任何 python 函数，而不必担心模型的底层格式。这种统一的表示方法简化了模型部署、重新部署和下游评分。

本文将帮助你开始使用 mlflow.pyfunc。

• 首先，让我们通过一个简单的玩具示例来创建 mlflow.pyfunc 类。
• 然后，我们将定义一个 mlflow.pyfunc 类，该类封装了一个机器学习管道（例如，一个估算器加上一些预处理逻辑）。我们还将训练、记录和加载这一机器学习管道，以便进行推理。
• 最后，让我们深入了解封装的 mlflow.pyfunc 对象，探索 mlflow 自动跟踪的丰富元数据和工件，更好地掌握 mlflow.pyfunc 提供的全部功能。

{pyfunc} 简单的玩具模型

首先，让我们创建一个简单的玩具 mlflow.pyfunc 模型，然后将它与 mlflow 工作流一起使用。

• 第 1 步：创建模型
  
• 第 2 步：记录模型
• 第 3 步：加载记录的模型以执行推理

# Step 1: Create a mlflow.pyfunc model
class ToyModel(mlflow.pyfunc.PythonModel):
    """
    ToyModel is a simple example implementation of an MLflow Python model.
    """
    
    def predict(self, context, model_input):
        """
        A basic predict function that takes a model_input list and returns a new list 
        where each element is increased by one.
        Parameters:
        - context (Any): An optional context parameter provided by MLflow.
        - model_input (list of int or float): A list of numerical values that the model will use for prediction.
        Returns:
        - list of int or float: A list with each element in model_input is increased by one.
        """
        return [x + 1 for x in model_input]

从上面的示例中可以看出，你可以创建一个 mlflow.pyfunc 模型来实现任何你认为适合你的 ML 解决方案的自定义 Python 函数，而不一定非得是现成的机器学习算法。

然后，你可以记录这个模型，稍后加载它来执行推理。

# Step 2: log this model as an mlflow run
with mlflow.start_run():
    mlflow.pyfunc.log_model(
        artifact_path = "model", 
        python_model=ToyModel()
    )
    run_id = mlflow.active_run().info.run_id

# Step 3: load the logged model to perform inference
model = mlflow.pyfunc.load_model(f"runs:/{run_id}/model")
# dummy new data
x_new = [1,2,3]
# model inference for the new data
print(model.predict(x_new))

[2, 3, 4]

{pyfunc} 封装的 XGBoost 管道

现在，让我们创建一个 ML 管道，将估算器与额外的自定义逻辑封装在一起。

在下面的示例中，XGB_PIPELINE 类是一个封装器，它将估算器与预处理步骤整合在一起，这对于某些 MLOps 实现来说是可取的。利用 mlflow.pyfunc，该封装器与估计器无关，并提供统一的模型表示法。具体如下：

• fit()： 该类不使用 XGBoost 的本地 API（xgboost.train()），而是使用 .fit()，这符合 sklearn 惯例，可直接集成到 sklearn 管道中，并确保不同估计器之间的一致性。
• DMatrix()： DMatrix 是 XGBoost 的核心数据结构，用于优化训练和预测数据。在该类中，将 pandas DataFrame 转换为 DMatrix 的步骤被封装在该类中，从而实现了与所有其他 sklearn 估算器一样与 pandas DataFrame 的无缝集成。
• predict()： 这是 mlflow.pyfunc 模型的通用推理 API。对于本 ML 管道、上述玩具模型以及我们封装在 mlflow.pyfunc 模型中的任何机器学习算法或自定义逻辑，它都是一致的。

import json
import xgboost as xgb
import mlflow.pyfunc
from typing import Any, Dict, Union
import pandas as pd

class XGB_PIPELINE(mlflow.pyfunc.PythonModel):
    """
    XGBWithPreprocess is an example implementation of an MLflow Python model with XGBoost.
    """
    
    def __init__(self, params: Dict[str, Union[str, int, float]]):
        """
        Initialize the model with given parameters.
        Parameters:
        - params (Dict[str, Union[str, int, float]]): Parameters for the XGBoost model.
        """
        self.params = params
        self.xgb_model = None
        self.config = None      
    def preprocess_input(self, model_input: pd.DataFrame) -> pd.DataFrame:
        """
        Preprocess the input data.
        Parameters:
        - model_input (pd.DataFrame): The input data to preprocess.
        Returns:
        - pd.DataFrame: The preprocessed input data.
        """
        processed_input = model_input.copy()
        # put any desired preprocessing logic here
        processed_input.drop(processed_input.columns[0], axis=1, inplace=True)
        return processed_input
    def fit(self, X_train: pd.DataFrame, y_train: pd.Series):
        """
        Train the XGBoost model.
        Parameters:
        - X_train (pd.DataFrame): The training input data.
        - y_train (pd.Series): The target values.
        """
        processed_model_input = self.preprocess_input(X_train.copy())
        dtrain = xgb.DMatrix(processed_model_input, label=y_train)
        self.xgb_model = xgb.train(self.params, dtrain)
    def predict(self, context: Any, model_input: pd.DataFrame) -> Any:
        """
        Predict using the trained XGBoost model.
        Parameters:
        - context (Any): An optional context parameter provided by MLflow.
        - model_input (pd.DataFrame): The input data for making predictions.
        Returns:
        - Any: The prediction results.
        """
        processed_model_input = self.preprocess_input(model_input.copy())
        dmatrix = xgb.DMatrix(processed_model_input)
        return self.xgb_model.predict(dmatrix)

现在，让我们来训练和记录这个模型。

from sklearn.datasets import make_regression
from sklearn.model_selection import train_test_split
import pandas as pd
# Generate synthetic datasets for demo
X, y = make_regression(n_samples=1000, n_features=10, noise=0.1, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# train and log the model
with mlflow.start_run(run_name = 'xgb_demo') as run:
    # Create an instance of XGB_PIPELINE
    params = {
        'objective': 'reg:squarederror',  
        'max_depth': 3,  
        'learning_rate': 0.1,
    }
    model = XGB_PIPELINE(params)
    # Fit the model
    model.fit(X_train=pd.DataFrame(X_train), y_train=y_train)
    # Log the model
    model_info = mlflow.pyfunc.log_model(
        artifact_path = 'model',
        python_model = model,
    )
    run_id = mlflow.active_run().info.run_id

模型已成功记录。现在，让我们加载模型进行推理。

loaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
loaded_model.predict(pd.DataFrame(X_test))

array([ 4.11692047e+00,  7.30551958e+00, -2.36042137e+01, -1.31888123e+02,
       ...

深入研究 Mlflow.pyfunc 对象

上述过程非常流畅，不是吗？这代表了 mlflow.pyfunc 对象的基本功能。现在，让我们深入探索 mlflow.pyfunc 的全部功能。

1. 模型信息

在上面的示例中，mlflow.pyfunc.log_model() 返回的 model_info 对象是 mlflow.models.model.ModelInfo 类的一个实例。它包含有关记录模型的元数据和信息。例如：

请运行 dir(model_info) 进一步探索，或查看源代码中定义的所有属性。我用得最多的属性是 model_uri，它表示在 mlflow 跟踪系统中可以在哪里找到记录的模型。

2. 加载模型

值得说明的是，loaded_model 并不是 XGB_PIPELINE 类的实例，而是 mlflow.pyfunc 为算法识别推理提供的封装对象。如下所示，如果尝试从 loaded_model 中获取 XGB_PIPELINE 类的属性，将返回错误信息。

print(loaded_model.params)

AttributeError: 'PyFuncModel' object has no attribute 'params'

3. 解包模型

好吧，你可能会问，那么 XGB_PIPELINE 的训练实例在哪里？它也被记录下来并可通过 mlflow 检索吗？

别担心，如下图所示，它被安全地保存起来，你可以轻松地将其拆封。

unwrapped_model = loaded_model.unwrap_python_model()
print(unwrapped_model.params)

{'objective': 'reg:squarederror', 'max_depth': 3, 'learning_rate': 0.1}

就是这样做的。有了 unwrapped_model，你就可以像这样访问自定义 ML 管道的任何属性或方法！我有时会在自定义管道中添加一些方便的方法，如 explain_model 或 post_processing，或者加入一些有用的属性来跟踪模型训练过程并提供诊断......好了，我还是就此打住吧，这些内容留待后续文章介绍。我只想说，你可以根据自己的用例自由定制 ML 管道，并且知道：

1. 你可以访问所有这些量身定制的方法和属性，以供下游使用，并且
2. 这种量身定制的自定义模型将封装在统一的 mlflow.pyfunc 推断 API 中，因此在必要时可以顺利迁移到其他估计器。

4. 上下文

你可能已经注意到，上面定义的 mlflow.pyfunc 类中的预测方法都有一个上下文参数。但有趣的是，当我们使用加载的模型进行预测时，并不需要这个参数。为什么？

loaded_model = mlflow.pyfunc.load_model(model_uri)
# the context parameter is not needed when calling `predict`
loaded_model.predict(model_input)

这是因为上面的 loaded_model 是 mlflow 提供的封装对象。如果我们使用未封装的模型，则需要明确定义上下文，如下所示；否则，代码将返回错误。

unwrapped_model = loaded_model.unwrap_python_model()
# need to provide context mannually
unwrapped_model.predict(context=None, model_input)

那么，这个背景是什么？它在预测方法中扮演什么角色？

上下文是一个 PythonModelContext 对象，它包含了pyfunc 模型在执行推理时可以使用的人工制品。它由 log_method() 方法自动隐式创建。

导航到项目 repo 中的 mlruns 子文件夹，该文件夹会在记录 mlflow 模型时由 mlflow 自动创建。找到以模型 run_id 命名的文件夹。在该文件夹中，你将看到为你自动记录的模型工件，如下所示。

# get run_id of a loaded model
print(loaded_model.metadata.run_id)

38a617d0f30645e8ae95eea4642a03c2

很不错吧？下面是 FYR 文件夹中需求和 MLmodel 文件的截图。

下面的需求说明了重新创建模型运行环境所需的依赖版本。

下面的 MLmodel 文档定义了以 YAML 格式加载和提供模型所需的元数据和配置。

结论

这就是 mlflow.pyfunc 模型构建方法。信息量很大，让我们来回顾一下

1. mlflow.pyfunc 提供了统一的模型表示法，不受用于构建模型的底层框架或库的影响。
   
2. 我们甚至可以将丰富的自定义逻辑封装到 mlflow.pyfunc 模型中，以便为每个用例量身定制，同时保持推理 API 的一致性和统一性。
3. 底层模型可以从加载的 mlflow.pyfunc 模型中解包出来，这样我们就可以利用更多针对每个用例定制的自定义方法/属性。
4. mlflow.pyfunc 模型对象记录了丰富的元数据和工件，这些数据和工件会被 mlflow 自动跟踪。
5. 这种统一的 mlflow.pyfunc 模型表示法可简化在不同算法之间进行实验和迁移的过程，以实现最佳性能。