scikit-learn模型持久化技术详解-CSDN博客

本文链接：https://blog.csdn.net/gitblog_00511/article/details/148323893

scikit-learn模型持久化技术详解

scikit-learn 一个基于 Python 的机器学习框架项目，适合对机器学习算法和应用感兴趣的人士学习和实践，内容包括分类、回归、聚类等多个领域。特点是功能丰富，算法齐全，易于理解和应用。项目地址: https://gitcode.com/gh_mirrors/sc/scikit-learn

概述

在机器学习项目中，模型持久化是一个关键环节。scikit-learn提供了多种模型持久化方法，每种方法都有其适用场景和优缺点。本文将全面解析这些持久化技术，帮助开发者根据实际需求选择最合适的方案。

持久化方法对比

| 持久化方法 | 优点 | 风险/缺点 | |------------------|----------------------------------------------------------------------|---------------------------------------------------------------------------| | ONNX | 无需Python环境即可服务模型；训练和服务环境独立；安全性最高 | 不是所有scikit-learn模型都支持；自定义估计器需要额外工作；无法重建原始Python对象 | | skops.io | 比pickle更安全；加载前可部分验证内容 | 不如pickle快；支持的类型少于pickle；需要与训练环境相同的环境 | | pickle | Python原生；可序列化大多数Python对象；protocol=5内存效率高 | 加载时可执行任意代码；需要与训练环境相同的环境 | | joblib | 内存效率高；支持内存映射；压缩解压便捷 | 基于pickle；加载时可执行任意代码；需要与训练环境相同的环境 | | cloudpickle | 可序列化非打包的自定义Python代码；加载效率与protocol=5的pickle相当 | 基于pickle；加载时可执行任意代码；无向前兼容保证；需要与训练环境相同的环境 |

选择持久化方法的关键问题

是否需要保留Python对象
如果仅需预测服务而不需要Python对象本身，ONNX是最佳选择。但需注意并非所有模型都支持ONNX。
安全性考虑
如果对模型来源有安全顾虑，应选择skops.io。它比pickle更安全，不会自动执行任意代码，但需要手动验证文件内容。
加载性能需求
如果需要高性能加载或进程间共享内存映射对象，joblib是理想选择。否则可以使用标准pickle模块。
特殊序列化需求
当pickle或joblib无法序列化模型（如包含用户定义函数）时，可考虑使用cloudpickle。

典型工作流程

1. 训练并持久化模型

from sklearn import ensemble, datasets

# 训练模型
clf = ensemble.HistGradientBoostingClassifier()
X, y = datasets.load_iris(return_X_y=True)
clf.fit(X, y)

2. 持久化路径选择

ONNX路径

需要ONNX运行时环境
服务环境可以非常精简，甚至不需要Python
内存占用通常比Python更少

Python序列化路径(skops/pickle/joblib/cloudpickle)

需要与训练环境相同的Python环境
依赖包及其版本必须完全一致
通常可跨硬件平台运行

各持久化方法详解

ONNX持久化

ONNX(Open Neural Network Exchange)格式最适合需要跨平台部署且不依赖Python环境的场景。

特点：

二进制序列化格式
提高模型在不同框架间的互操作性
增强在不同计算架构上的可移植性

使用示例：

from skl2onnx import to_onnx
import numpy

# 转换为ONNX格式
onx = to_onnx(clf, X[:1].astype(numpy.float32), target_opset=12)
with open("model.onnx", "wb") as f:
    f.write(onx.SerializeToString())

# 加载并使用
from onnxruntime import InferenceSession
with open("model.onnx", "rb") as f:
    onx = f.read()
sess = InferenceSession(onx, providers=["CPUExecutionProvider"])
pred = sess.run(None, {"X": X_test.astype(numpy.float32)})[0]

局限性：

不支持所有scikit-learn模型
第三方估计器支持有限
自定义转换器文档较少

skops.io持久化

skops.io提供了比pickle更安全的替代方案，它不会自动加载任意代码。

特点：

仅加载用户明确信任的类型和函数
加载前可检查文件内容
API设计与pickle类似

使用示例：

import skops.io as sio

# 保存模型
sio.dump(clf, "model.skops")

# 加载模型
unknown_types = sio.get_untrusted_types(file="model.skops")
# 检查unknown_types内容后
clf = sio.load("model.skops", trusted=unknown_types)

pickle/joblib/cloudpickle持久化

这三种方法都基于pickle协议，但有细微差别：

pickle：Python标准库，可序列化大多数Python对象
joblib：处理大型模型和numpy数组更高效
cloudpickle：可序列化pickle和joblib无法处理的特殊对象

使用示例：

# 使用pickle保存
import pickle
with open("model.pkl", "wb") as f:
    pickle.dump(clf, f, protocol=5)

# 使用joblib保存
import joblib
joblib.dump(clf, "model.joblib")

# 使用cloudpickle保存
import cloudpickle
with open("model.pkl", "wb") as f:
    cloudpickle.dump(clf, f)

安全性与维护性考量

安全风险：
- pickle系列方法加载时可执行任意代码
- ONNX计算也可能存在安全隐患，建议在沙箱环境中运行
版本兼容性：
- 不支持跨scikit-learn版本加载模型
- 建议保存训练数据、源代码和依赖版本等元数据
环境一致性：
- 生产环境应与训练环境保持完全一致
- 可使用容器技术(如Docker)冻结环境

版本不一致警告：

from sklearn.exceptions import InconsistentVersionWarning
warnings.simplefilter("error", InconsistentVersionWarning)

try:
    with open("model.pkl", "rb") as f:
        est = pickle.load(f)
except InconsistentVersionWarning as w:
    print(w.original_sklearn_version)