使用 mlflow.pyfunc 的自定义 MLflow 模型
如果您想了解 MLflow 自定义模型所提供的所有灵活性和自定义选项，这篇博客将帮助您更好地理解如何利用这种强大且高度可定制的模型存储格式。
MLflow 提供了内置方法，用于记录和处理来自许多流行的机器学习和生成式 AI 框架和模型提供商的模型，例如 scikit-learn、PyTorch、HuggingFace Transformers 和 LangChain。例如，mlflow.sklearn.log_model 将 scikit-learn 模型记录为 MLflow 工件，而无需您定义自定义预测方法或处理工件。
然而，在某些情况下，您可能使用的框架 MLflow 没有内置方法，或者您可能需要与模型默认预测输出不同的东西。在这些情况下，MLflow 允许您创建自定义模型，以支持几乎任何框架，并将自定义逻辑集成到现有受支持的框架中。
最简单的情况下，只需定义一个自定义的 `predict` 方法并记录模型。以下示例定义了一个简单的 `pyfunc` 模型，它只返回其输入的平方
import mlflow
# Define a custom model
class MyModel(mlflow.pyfunc.PythonModel):
def predict(self, context, model_input):
# Directly return the square of the input
return model_input**2
# Save the model
with mlflow.start_run():
model_info = mlflow.pyfunc.log_model(
artifact_path="model",
python_model=MyModel()
)
# Load the model
loaded_model = mlflow.pyfunc.load_model(
model_uri=model_info.model_uri
)
# Predict
loaded_model.predict(2)
让我们深入了解其工作原理，从一些基本概念开始。
模型和模型风格
MLflow 模型是一个目录，包含重现机器学习模型在不同环境中所需的一切。除了存储的模型本身，最重要的组件是指定模型支持的模型风格的 `MLmodel` YAML 文件。模型风格是一组规则，规定 MLflow 如何与模型交互（即保存、加载和从中获取预测）。
当您在 MLflow 中创建自定义模型时，它具有 `python_function` 或 `pyfunc` 模型风格，这是一种跨 MLflow 格式的“通用转换器”。当您使用内置模型风格（例如通过 mlflow.sklearn.log_model）在 MLflow 中保存模型时，该模型除了其特定于框架的风格外，还具有 `pyfunc` 模型风格。同时拥有特定于框架和 `pyfunc` 模型风格允许您通过框架的本机 API（例如 `scikit-learn`）或通过 `pyfunc` 风格的与框架无关的推理 API 使用模型。
带有 `pyfunc` 风格的模型被加载为 mlflow.pyfunc.PyFuncModel 类的实例，该类公开了一个标准化的 `predict` 方法。这使得通过单个函数调用即可进行直接推理，而无需关心底层模型的实现细节。
定义自定义 MLflow Pyfunc 模型
将任何受支持的机器学习框架中的模型保存为 MLflow 模型，都会创建一个 `pyfunc` 模型风格，它为模型的管理和使用提供了与框架无关的接口。但是，如果您使用的框架没有 MLflow 集成，或者您试图从模型中获取一些自定义行为呢？自定义 `pyfunc` 模型允许您使用几乎任何框架，并集成自定义逻辑。
要实现一个自定义 `pyfunc` 模型，需要定义一个继承自 `PythonModel` 类的新 Python 类，并实现必要的方法。最少需要实现一个自定义的 `predict` 方法。接下来，创建模型实例并记录或保存模型。加载已保存或已记录的模型后，即可将其用于预测。
让我们看几个例子，每个例子都会增加一些复杂性并突出定义自定义 `pyfunc` 模型的不同方面。我们将介绍在 `pyfunc` 模型中实现自定义行为的四种主要技术
- 实现一个自定义的 `predict` 方法
- 实现一个自定义的 `__init__` 方法
- 实现一个自定义的 `load_context` 方法
- 实现用户定义的自定义方法
定义自定义的 `predict` 方法
最低限度，一个 `pyfunc` 模型应该指定一个自定义的 `predict` 方法，它定义了当我们调用 `model.predict` 时会发生什么。下面是一个自定义模型的例子，它对模型输入进行简单的线性变换，将每个输入乘以二并加上三
import pandas as pd
import mlflow
from mlflow.pyfunc import PythonModel
# Custom PythonModel class
class SimpleLinearModel(PythonModel):
def predict(self, context, model_input):
"""
Applies a simple linear transformation
to the input data. For example, y = 2x + 3.
"""
# Assuming model_input is a pandas DataFrame with one column
return pd.DataFrame(2 * model_input + 3)
with mlflow.start_run():
model_info = mlflow.pyfunc.log_model(
artifact_path="linear_model",
python_model=SimpleLinearModel(),
input_example=pd.DataFrame([10, 20, 30]),
)
请注意，您也可以（并且应该）在保存/记录模型时包含签名和输入示例。如果您传入输入示例，签名将自动推断。模型签名提供了一种方式，让 MLflow 强制正确使用您的模型。
一旦我们定义了模型路径并保存了模型实例，我们就可以加载已保存的模型并用它来生成预测
# Now the model can be loaded and used for predictions
loaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
model_input = pd.DataFrame([1, 2, 3]) # Example input data
print(loaded_model.predict(model_input)) # Outputs transformed data
这将返回
: 0
: 0 5
: 1 7
: 2 9
请注意，如果您只需要一个自定义的 `predict` 方法——也就是说，如果您的模型没有任何需要特殊处理的工件——您可以直接保存或记录 `predict` 方法，而无需将其封装在一个 Python 类中
import mlflow
import pandas as pd
def predict(model_input):
"""
Applies a simple linear transformation
to the input data. For example, y = 2x + 3.
"""
# Assuming model_input is a pandas DataFrame with one column
return pd.DataFrame(2 * model_input + 3)
# Pass predict method as python_model argument to save/log model
with mlflow.start_run():
model_info = mlflow.pyfunc.log_model(
artifact_path="simple_function",
python_model=predict,
input_example=pd.DataFrame([10, 20, 30]),
)
请注意，使用这种方法，我们**必须**在自定义 `predict` 方法中包含一个输入示例。我们还必须修改 `predict` 方法，使其只接受一个输入（即，没有 `self` 或 `context`）。运行此示例，然后使用与前面代码块相同的代码进行加载，将保留与使用类定义示例相同的输出。
参数化自定义模型
现在假设我们希望对自定义线性函数模型进行参数化，以便它可以使用不同的斜率和截距。我们可以定义 `__init__` 方法来设置自定义参数，如下例所示。请注意，自定义模型类的 `__init__` 方法不应该用于加载外部资源，例如数据文件或预训练模型；这些在 `load_context` 方法中处理，我们稍后将讨论。
import pandas as pd
import mlflow
from mlflow.pyfunc import PythonModel
# Custom PythonModel class
class ParameterizedLinearModel(PythonModel):
def __init__(self, slope, intercept):
"""
Initialize the parameters of the model. Note that we are not loading
any external resources here, just setting up the parameters.
"""
self.slope = slope
self.intercept = intercept
def predict(self, context, model_input):
"""
Applies a simple linear transformation
to the input data. For example, y = 2x + 3.
"""
# Assuming model_input is a pandas DataFrame with one column
return pd.DataFrame(self.slope * model_input + self.intercept)
linear_model = ParameterizedLinearModel(10, 20)
# Saving the model with mlflow
with mlflow.start_run():
model_info = mlflow.pyfunc.log_model(
artifact_path="parameter_model",
python_model=linear_model,
input_example=pd.DataFrame([10, 20, 30]),
)
同样，我们可以加载这个模型并进行一些预测
loaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
model_input = pd.DataFrame([1, 2, 3]) # Example input data
print(loaded_model.predict(model_input)) # Outputs transformed data
: 0
: 0 30
: 1 40
: 2 50
在许多情况下，我们可能希望以这种方式对模型进行参数化。我们可以在 `__init__` 方法中定义变量来
- 设置模型超参数。
- 使用不同参数集进行模型 A/B 测试。
- 设置用户特定自定义项。
- 切换功能。
- 设置（例如）访问凭据和访问外部 API 模型的端点。
在某些情况下，我们可能希望在推理时而不是在初始化模型时传递参数。这可以通过模型推理参数实现。要使用推理参数，我们必须传递一个包含 `params` 的有效模型签名。以下是如何修改上述示例以使用推理参数
import pandas as pd
import mlflow
from mlflow.models import infer_signature
from mlflow.pyfunc import PythonModel
# Custom PythonModel class
class LinearFunctionInferenceParams(PythonModel):
def predict(self, context, model_input, params):
"""
Applies a simple linear transformation
to the input data. For example, y = 2x + 3.
"""
slope = params["slope"]
# Assuming model_input is a pandas DataFrame with one column
return pd.DataFrame(slope * model_input + params["intercept"])
# Set default params
params = {"slope": 2, "intercept": 3}
# Define model signature
signature = infer_signature(model_input=pd.DataFrame([10, 20, 30]), params=params)
# Save the model with mlflow
with mlflow.start_run():
model_info = mlflow.pyfunc.log_model(
artifact_path="model_with_params",
python_model=LinearFunctionInferenceParams(),
signature=signature,
)
像之前一样加载模型后，您现在可以将 `params` 参数传递给 `predict` 方法，从而能够将相同的已加载模型用于不同的参数组合
loaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
parameterized_predictions = loaded_model.predict(
pd.DataFrame([10, 20, 30]), params={"slope": 2, "intercept": 10}
)
print(parameterized_predictions)
: 0
: 0 30
: 1 50
: 2 70
使用 `load_context` 加载外部资源
自定义模型通常需要外部文件（例如模型权重）才能执行推理。这些文件（或工件）必须小心处理，以避免不必要地将文件加载到内存中或在模型序列化期间出错。在 MLflow 中构建自定义 `pyfunc` 模型时，您可以使用 `load_context` 方法正确处理模型工件。
`load_context` 方法接收一个包含模型在推理期间可以使用的工件的 `context` 对象。您可以在保存或记录模型时使用 `artifacts` 参数指定这些工件，使它们可以通过 `context.artifacts` 字典访问 `load_context` 方法。
实际上，`load_context` 方法通常通过处理模型工件的加载来初始化 `predict` 方法调用的模型。
这提出了一个重要问题：为什么我们要在 `load_context` 方法中加载工件并定义模型，而不是在 `__init__` 或直接在 `predict` 中？正确使用 `load_context` 对于 MLflow `pyfunc` 模型的维护性、效率、可扩展性和可移植性至关重要。这是因为
- `load_context` 方法在通过 `mlflow.pyfunc.load_model` 加载模型时执行一次。此设置确保此方法中定义的资源密集型过程（例如加载大型模型文件）不会不必要地重复。如果在 `predict` 方法中进行工件加载，则每次进行预测时都会发生。这对于资源密集型模型来说效率极低。
- 保存或记录 MLflow `pyfunc` 模型涉及序列化 Python 模型类（您创建的 `mlflow.pyfunc.PythonModel` 的子类）及其属性。复杂的 ML 模型并不总是与用于序列化 Python 对象的方法兼容，如果它们作为 Python 对象的属性创建，可能会导致错误。
例如，假设我们要加载 `gguf` 模型格式（一种用于存储推理模型的文件格式）的大型语言模型 (LLM)，并使用 ctransformers 库运行它。在撰写本文时，没有内置的模型风格允许我们使用 `gguf` 模型进行推理，因此我们将创建一个自定义 `pyfunc` 模型，在 `load_context` 方法中加载所需的库和模型文件。具体来说，我们将加载 Mistral 7B 模型 AWQ 版本的量化版本。
首先，我们将使用 huggingface hub cli 下载模型快照
huggingface-cli download TheBloke/Mistral-7B-v0.1-GGUF \
mistral-7b-v0.1.Q4_K_M.gguf \
--local-dir /path/to/mistralfiles/ \
--local-dir-use-symlinks False
然后我们将定义我们的自定义 `pyfunc` 模型。请注意 `load_context` 方法的添加
import ctransformers
from mlflow.pyfunc import PythonModel
class CTransformersModel(PythonModel):
def __init__(self, gpu_layers):
"""
Initialize with GPU layer configuration.
"""
self.gpu_layers = gpu_layers
self.model = None
def load_context(self, context):
"""
Load the model from the specified artifacts directory.
"""
model_file_path = context.artifacts["model_file"]
# Load the model
self.model = ctransformers.AutoModelForCausalLM.from_pretrained(
model_path_or_repo_id=model_file_path,
gpu_layers=self.gpu_layers,
)
def predict(self, context, model_input):
"""
Perform prediction using the loaded model.
"""
if self.model is None:
raise ValueError(
"The model has not been loaded. "
"Ensure that 'load_context' is properly executed."
)
return self.model(model_input)
这里发生了很多事情，让我们来分解一下。以下是关键点
- 和以前一样，我们使用 `__init__` 方法对模型进行参数化（在本例中，用于设置模型的 `gpu_layers` 参数）。
- `load_context` 方法的目的是加载在 `predict` 方法中使用所需的工件。在这种情况下，我们需要加载模型及其权重。
- 您会注意到我们引用了 `context.artifacts["model_file"]`。这来自 `mlflow.pyfunc.save_model` 或 `mlflow.pyfunc.load_model` 的 `artifacts` 参数，如以下代码片段所示。这是使用 `pyfunc` 模型的重要组成部分。`predict` 和 `load_context` 方法可以通过 `context.artifacts` 对象访问在 `save_model` 或 `log_model` 方法的 `artifacts` 参数中定义的工件。`load_context` 在通过 `load_model` 加载模型时执行；如前所述，这提供了一种方法，以确保模型的潜在耗时初始化不会在每次使用模型进行预测时发生。
现在我们可以初始化并保存模型实例。请注意 `save_model` 函数的 `artifacts` 参数
# Create an instance of the model
mistral_model = CTransformersModel(gpu_layers=50)
# Log the model using mlflow with the model file as an artifact
with mlflow.start_run():
model_info = mlflow.pyfunc.log_model(
artifact_path="mistral_model",
python_model=mistral_model,
artifacts={"model_file": model_file_path},
pip_requirements=[
"ctransformers==0.2.27",
],
)
# Load the saved model
loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
# Make a prediction with the model
loaded_model.predict("Question: What is the MLflow Pyfunc model flavor?")
总结：正确使用 `load_context` 方法有助于确保高效处理模型工件，并防止因尝试将工件定义为模型属性而可能导致的序列化错误。
定义自定义方法
您可以在自定义 `pyfunc` 模型中定义自己的方法来处理任务，例如预处理输入或后处理输出。然后，这些自定义方法可以由 `predict` 方法调用。请记住，这些自定义方法，就像 `__init__` 和 `predict` 一样，**不应**用于加载工件。加载工件是 `load_context` 方法的专属职责。
例如，我们可以修改 `CTransformersModel` 以包含一些提示格式，如下所示
import ctransformers
from mlflow.pyfunc import PythonModel
class CTransformersModel(PythonModel):
def __init__(self, gpu_layers):
"""
Initialize with GPU layer configuration.
"""
self.gpu_layers = gpu_layers
self.model = None
def load_context(self, context):
"""
Load the model from the specified artifacts directory.
"""
model_file_path = context.artifacts["model_file"]
# Load the model
self.model = ctransformers.AutoModelForCausalLM.from_pretrained(
model_path_or_repo_id=model_file_path,
gpu_layers=self.gpu_layers,
)
@staticmethod
def _format_prompt(prompt):
"""
Formats the user's prompt
"""
formatted_prompt = (
"Question: What is an MLflow Model?\n\n"
"Answer: An MLflow Model is a directory that includes "
"everything that is needed to reproduce a machine "
"learning model across different environments. "
"It is essentially a container holding the trained model "
"files, dependencies, environment details, input examples, "
"and additional metadata. The directory also includes an "
"MLmodel YAML file, which describes the different "
f"flavors of the model.\n\nQuestion: {prompt}\nAnswer: "
)
return formatted_prompt
def predict(self, context, model_input):
"""
Perform prediction using the loaded model.
"""
if self.model is None:
raise ValueError(
"Model was not loaded. Ensure that 'load_context' "
"is properly executed."
)
return self.model(self._format_prompt(model_input))
现在，`predict` 方法可以访问私有的 `_format_prompt` 静态方法，以对提示应用自定义格式。
依赖项和源代码
上面定义的自定义 `CTransformersModel` 使用了 `ctransformers` 库。有几种不同的方法可以确保这个库（以及任何其他源代码，包括来自您本地设备的源代码）与您的模型正确加载。正确指定依赖项对于确保自定义模型在不同环境中按预期工作至关重要。
有三种主要方法可以指定依赖项
- 使用 `save_model` 或 `log_model` 的 `pip_requirements` 参数显式定义 pip 需求。
- 使用 `save_model` 或 `log_model` 的 `extra_pip_requirements` 参数将额外的 pip 需求添加到自动生成的需求集中。
- 使用 `save_model` 或 `log_model` 的 `conda_env` 参数定义 Conda 环境。
早些时候，我们使用了第一种方法来指定需要 `ctransformers` 库
# Log the model using mlflow with the model file as an artifact
with mlflow.start_run():
model_info = mlflow.pyfunc.save_model(
artifact_path="mistralmodel",
python_model=mistral_model,
artifacts={"model_file": "path/to/mistral/model/on/local/filesystem"},
pip_requirements=[
"ctransformers==0.2.27",
],
)
如果您未明确指定依赖项，MLflow 将尝试推断正确的需求集和环境详细信息。为了提高准确性，**强烈建议**在保存或记录模型时包含一个 `input_example`，因为内部会执行一个示例推理步骤，该步骤将捕获与推理执行相关联的任何已加载库引用，从而大大提高记录正确依赖项的可能性。
您还可以使用 `code_path` 参数处理自己文件系统上的自定义代码。`code_path` 接受一个 Python 文件依赖项路径列表，并在模型加载之前将它们添加到系统路径的开头，这样自定义 `pyfunc` 模型就可以从这些模块中导入。
有关 `pip`、`Conda` 和本地代码要求的接受格式的更多详细信息，请参阅 log_model 和 save_model 函数的文档。
总结：MLflow 中的自定义 Pyfunc 模型
MLflow 提供了用于处理来自许多流行机器学习框架（例如 scikit-learn、PyTorch 和 Transformers）的模型的内置方法。当您想要处理尚无内置模型风格的模型，或者想要为具有内置模型风格的模型实现自定义 `predict` 方法时，您可以定义自己的自定义 `mlflow.pyfunc` 模型。
有几种方法可以自定义 `pyfunc` 模型以获得所需行为。最少，您可以实现一个自定义的 `predict` 方法。如果您的模型需要保存或加载工件，您还应该实现一个 `load_context` 方法。为了进一步自定义，您可以使用 `__init__` 方法设置自定义属性，并定义自己的自定义方法进行预处理和后处理。结合这些方法，您可以灵活地为您的机器学习模型定义自定义逻辑。
进一步学习
有兴趣了解更多关于自定义 `pyfunc` 实现的信息吗？您可以访问