如果您想了解 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)
让我们深入探讨其工作原理，从一些基本概念开始。
模型与模型风格 (Model Flavors)
MLflow 模型是一个目录，包含在不同环境中重现机器学习模型所需的一切。除了存储的模型本身之外，存储的最重要组件是 MLmodel YAML 文件，它指定了模型支持的模型风格 (model flavors)。模型风格是一组规则，规定了 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]),
)
请注意，在保存/记录模型时，您也可以（而且应该）包含一个签名 (signature) 和一个输入示例 (input example)。如果您传入输入示例，签名将自动推断。模型签名是 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 方法——也就是说，如果您的模型没有任何需要特殊处理的工件 (artifacts)——您可以直接保存或记录 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 的模型的访问凭据和端点。
在某些情况下，我们可能希望能够在推理时而不是在初始化模型时传递参数。这可以通过模型推理参数 (model inference params) 实现。要使用推理参数，我们必须传递包含 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 实现的信息吗？您可以访问