使用 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 文件。模型口味（A model flavor）是一组规则，指定 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 实现的信息吗？您可以访问