跳到主要内容

无需 Pickle 的模型格式

实验性功能

此功能处于实验阶段，未来版本可能会有所更改。

使用 Python 的 pickle 或 cloudpickle 保存模型依赖于 Python 的对象序列化机制，该机制在反序列化期间可能会执行任意代码。MLflow 为多种模型风格提供了更安全的无 pickle (pickle-free) 保存格式。在可能的情况下，请优先使用这些格式。

即将到来的默认行为变更

无 pickle 保存格式将在即将发布的 MLflow 版本中成为默认设置。大多数用户（scikit-learn、LightGBM、LangChain、自定义 Python 模型）将不会受到破坏性影响。 PyTorch 用户应查阅相关要求 — 此后将需要提供 input_example 且不支持 torch.jit.ScriptModule。为避免受此变更影响，请锁定 MLflow 版本或在模型记录代码中指定 serialization_format="pickle"。

Scikit-learn 模型的无 pickle 格式​

保存 scikit-learn 模型时，设置参数 serialization_format="skops" 可使用 skops 格式，从而实现 scikit-learn 模型安全的反序列化。skops 格式不依赖于 Python 的 pickle，并且在加载时避免了任意代码执行。

python
import mlflow
from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier

X, y = load_iris(return_X_y=True)
model = DecisionTreeClassifier().fit(X, y)

with mlflow.start_run():
    mlflow.sklearn.log_model(
        model,
        name="model",
        serialization_format="skops",
    )

对于某些包含自定义或第三方类型的 scikit-learn 模型，您需要将 skops_trusted_types 设置为完全限定类型名称的列表，以便 skops 能够加载它们。例如，包含自定义转换器的管道必须将其转换器类型列为受信任类型。

python
from numpy.random import randint
from sklearn.base import BaseEstimator, TransformerMixin
import pandas as pd
from sklearn.pipeline import Pipeline


class CustomTransformer(BaseEstimator, TransformerMixin):
    def fit(self, X, y=None):
        return self

    def transform(self, X, y=None):
        # Perform arbitrary transformation
        X["random_int"] = randint(0, 10, X.shape[0])
        return X


df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]})

pipeline = Pipeline(steps=[("use_custom_transformer", CustomTransformer())])

# Custom classes must be marked as trusted for skops
with mlflow.start_run():
    mlflow.sklearn.log_model(
        pipeline,
        name="model",
        serialization_format="skops",
        skops_trusted_types=["__main__.CustomTransformer"],
    )

有关完整参数，请参阅 mlflow.sklearn.log_model() 和 mlflow.sklearn.save_model()。

PyTorch 模型的无 pickle 格式​

保存 PyTorch 模型时，使用 serialization_format="pt2" 会调用 torch.export.save，并将模型存储为跟踪图 (traced graph)，而不是使用 pickle 序列化格式。

要求与限制
  • torch >= 2.4
  • 必须提供 input_example（且对于导出的模型，仅支持 Tensor 类型的输入）
  • 不支持 torch.jit.ScriptModule 模型
  • 模型加载后，权重无法在设备之间传输。如果模型保存时权重位于 GPU 0 上，则加载时也必须将其加载到 GPU 0，而不能加载到 CPU 或其他 GPU 上。
  • 为避免这些限制，请通过设置 serialization_format="pickle" 来使用基于 pickle 的序列化。
python
import mlflow
import torch
from torch import nn
from sklearn.datasets import load_diabetes

# Load a real dataset and use a sample as the input example
X, _ = load_diabetes(return_X_y=True)
input_example = torch.tensor(X[:5], dtype=torch.float32).numpy()
sequential_model = nn.Sequential(nn.Linear(10, 3), nn.ReLU(), nn.Linear(3, 1))
with mlflow.start_run():
    mlflow.pytorch.log_model(
        sequential_model,
        name="model",
        serialization_format="pt2",
        input_example=input_example,
    )

详情请参阅 mlflow.pytorch.log_model() 和 mlflow.pytorch.save_model()。

LightGBM 模型的无 pickle 格式​

对于属于 scikit-learn 模型类型的 LightGBM 模型（如 LGBMClassifier、LGBMRegressor），可以使用 skops 格式。这不适用于使用原生格式的 lightgbm.Booster 实例。

python
import mlflow
from lightgbm import LGBMClassifier
from sklearn.datasets import load_iris

X, y = load_iris(return_X_y=True)
model = LGBMClassifier(objective="multiclass", random_state=42).fit(X, y)

with mlflow.start_run():
    mlflow.lightgbm.log_model(
        model,
        name="model",
        serialization_format="skops",
        skops_trusted_types=[
            "collections.OrderedDict",
            "lightgbm.basic.Booster",
            "lightgbm.sklearn.LGBMClassifier",
        ],
    )

详情请参阅 mlflow.lightgbm.log_model() 和 mlflow.lightgbm.save_model()。

LangChain 模型的无 pickle 格式​

LangChain 模型支持保存为代码模型 (Models From Code) 工件，从而完全避免使用 pickle。有关详细信息和示例，请参阅 代码模型 (Models From Code)。

自定义 Python 模型的无 pickle 格式​

自定义 Python 模型支持保存为代码模型 (Models From Code) 工件，从而完全避免使用 pickle。有关详细信息和示例，请参阅 代码模型 (Models From Code)。

强制执行 MLflow 无 pickle 模型加载的全局配置​

您可以通过设置以下环境变量，在全局范围内禁止 MLflow 使用 pickle 或 cloudpickle 加载模型：

bash
export MLFLOW_ALLOW_PICKLE_DESERIALIZATION=false

当设置为 false 时，加载使用 pickle 或 cloudpickle 保存的模型将会报错，除非您在记录模型时使用了无 pickle 的保存选项。为了向后兼容，默认值为 true。

注意

在 Databricks Runtime 中，无论 MLFLOW_ALLOW_PICKLE_DESERIALIZATION 环境变量如何设置，始终允许加载使用 pickle 或 cloudpickle 序列化的模型。这是因为 Databricks 使用受管理的 MLflow 工件存储，可保护模型工件免受未经授权的修改。通过确保工件完整性，该平台降低了在 pickle 或 cloudpickle 反序列化期间执行恶意代码的风险。

总结​

模型风格无 pickle 选项备注成为默认设置时的影响
Scikit-learnserialization_format="skops"必要时使用 skops_trusted_types。无影响 — 对大多数模型为透明升级。
PyTorchserialization_format="pt2"需要 input_example，torch >= 2.4；不适用于 ScriptModule。破坏性变更 — 将需要 input_example；不支持 torch.jit.ScriptModule。
DSPyuse_dspy_model_save=True需要 dspy > 3.1.0。需要升级到 dspy > 3.1.0。
LightGBMserialization_format="skops"仅适用于 sklearn 类型的 LightGBM 模型。无影响 — 对 sklearn 类型的 LightGBM 模型为透明升级。
LangChain来自代码的模型Pickle 保存将发出警告。无影响。
自定义 Python 模型来自代码的模型Pickle 保存将发出警告。无影响。