第 86 课·工程化与部署进阶

模型持久化：Pickle、Joblib 与 ONNX

deploymentserialization

第86课 - 模型持久化：Pickle、Joblib 与 ONNX

学习目标

完成本课后，你将能够：

  1. 理解模型持久化的重要性及其在机器学习工作流中的作用。
  2. 掌握使用Pickle和Joblib保存与加载Scikit-learn模型的标准流程。
  3. 了解ONNX格式的优势，并能将一个训练好的模型转换为ONNX格式。
  4. 根据项目需求选择合适的持久化方案，并规避常见的陷阱。

核心概念

想象一下，你花了一整天时间精心训练了一个机器学习模型，它表现优异。现在，你需要把这份“智能”分享给同事，或者部署到服务器上，让它处理真实世界的请求。但你总不能每次都从头开始训练吧？这就需要用到模型持久化——简单来说，就是把训练好的模型“保存”成一个文件，下次需要时再“加载”回来，即装即用。

在Python的机器学习生态中，我们有三种主流的选择：

  1. Pickle：Python的内置序列化模块。它就像一个万能文件夹，可以把几乎任何Python对象（包括我们的模型）打包存起来。优点是简单、无需额外安装；缺点是对大型NumPy数组的序列化效率不是最高，且跨不同Python版本或库版本时可能有兼容性问题。

  2. Joblib：专门为高效序列化包含大量NumPy数组的对象（如Scikit-learn模型）而优化的库。它通常比Pickle更快，生成的文件也更小，是保存Scikit-learn模型的推荐选择。可以把它看作是一个“优化版文件夹”，专门为科学计算数据打包。

  3. ONNX (Open Neural Network Exchange)：这是一个开放的模型表示标准。它就像一个“通用翻译器”，无论你的模型是用PyTorch、TensorFlow还是Scikit-learn训练的，都可以转换成ONNX格式。转换后的模型可以在任何支持ONNX的运行时（如ONNX Runtime）中运行，实现跨框架、跨平台部署，是生产环境部署的利器。

代码示例

下面，我们将以一个经典的鸢尾花分类模型为例，演示三种持久化方法。

准备工作：训练一个模型

首先，我们训练一个简单的模型作为后续操作的对象。

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
import numpy as np

# 加载数据
iris = load_iris()
X, y = iris.data, iris.target
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 训练一个随机森林分类器
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# 评估一下
y_pred = model.predict(X_test)
print(f"模型在测试集上的准确率: {accuracy_score(y_test, y_pred):.2f}")

方法一：使用Pickle保存与加载

import pickle

# 1. 保存模型到文件
with open('iris_model.pkl', 'wb') as f:  # 'wb' 表示以二进制写模式打开
    pickle.dump(model, f)
print("模型已使用Pickle保存至 iris_model.pkl")

# 2. 从文件加载模型
with open('iris_model.pkl', 'rb') as f:  # 'rb' 表示以二进制读模式打开
    loaded_model = pickle.load(f)

# 3. 验证加载的模型
loaded_pred = loaded_model.predict(X_test)
print(f"加载后的模型准确率: {accuracy_score(y_test, loaded_pred):.2f}")

方法二：使用Joblib保存与加载（推荐用于Scikit-learn）

import joblib

# 1. 保存模型到文件
joblib.dump(model, 'iris_model.joblib')
print("模型已使用Joblib保存至 iris_model.joblib")

# 2. 从文件加载模型
loaded_model_joblib = joblib.load('iris_model.joblib')

# 3. 验证加载的模型
loaded_pred_joblib = loaded_model_joblib.predict(X_test)
print(f"使用Joblib加载的模型准确率: {accuracy_score(y_test, loaded_pred_joblib):.2f}")

# 可以顺便比较一下文件大小
import os
pkl_size = os.path.getsize('iris_model.pkl')
joblib_size = os.path.getsize('iris_model.joblib')
print(f"Pickle文件大小: {pkl_size} bytes, Joblib文件大小: {joblib_size} bytes")

方法三：使用ONNX转换与推理

这需要先安装 sklearn-onnx 和 onnxruntime。

pip install sklearn-onnx onnxruntime

from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
import onnxruntime as rt

# 1. 将Scikit-learn模型转换为ONNX
# 需要指定输入数据的类型和形状（样本数可以是None）
initial_type = [('float_input', FloatTensorType([None, 4]))]
onnx_model = convert_sklearn(model, initial_types=initial_type)

# 2. 将ONNX模型保存为文件
with open("iris_model.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())
print("模型已转换并保存为ONNX格式: iris_model.onnx")

# 3. 使用ONNX Runtime进行推理
# 创建推理会话
sess = rt.InferenceSession("iris_model.onnx")
# 获取输入名称
input_name = sess.get_inputs()[0].name
label_name = sess.get_outputs()[0].name

# 准备输入数据（需要是NumPy数组，且类型为float32）
X_test_onnx = X_test.astype(np.float32)

# 进行推理
pred_onnx = sess.run([label_name], {input_name: X_test_onnx})[0]
print(f"使用ONNX Runtime的预测结果形状: {pred_onnx.shape}")
print(f"前5个样本的预测: {pred_onnx[:5]}")

实践练习

练习1：基础操作

  1. 加载 sklearn.datasets 中的 digits 手写数字数据集。
  2. 使用 SVM 分类器（sklearn.svm.SVC）进行训练。
  3. 分别使用 Pickle 和 Joblib 将训练好的模型保存为 digits_svm.pkl 和 digits_svm.joblib。
  4. 编写代码加载这两个模型，并在测试集上验证它们的准确率是否一致。

练习2：封装函数

创建一个函数 save_load_model(model, filename, method='joblib')，它可以接收一个模型对象、文件名和方法（'pickle' 或 'joblib'），并完成保存和加载的全过程，最后返回加载后的模型。测试该函数。

练习3：探索ONNX（进阶）

尝试将你在练习1中训练的SVM模型转换为ONNX格式。

  1. 注意，SVM模型在转换时可能需要指定其他参数（如 options）。
  2. 使用 onnxruntime 加载转换后的模型，并在测试集上计算准确率，与原始Scikit-learn模型的结果对比。
  3. （可选）比较三种格式（.pkl, .joblib, .onnx）文件的大小。

常见错误

  1. 忘记保存预处理器：模型通常需要和预处理步骤（如StandardScaler, OneHotEncoder）一起工作。只保存模型而不保存预处理器，会导致在新数据上推理失败。解决方案是使用Pipeline，并将整个Pipeline一起保存。

    from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
pipe = Pipeline([('scaler', StandardScaler()), ('model', RandomForestClassifier())])
pipe.fit(X_train, y_train)
joblib.dump(pipe, 'full_pipeline.joblib') # 保存整个Pipeline

  2. 版本不兼容：使用Pickle或Joblib加载模型时，如果Pandas、NumPy、Scikit-learn等库的版本与保存时不一致，可能引发错误。在团队协作或部署时，务必记录并保持环境一致性，或使用Docker容器化。

  3. ONNX操作符不支持：不是所有模型或操作都能顺利转换为ONNX。一些自定义的转换器或特定的模型参数可能不受支持，转换时会报错。需要查阅sklearn-onnx的文档，看是否有对应的转换器。

  4. 文件路径错误：保存和加载时使用了不同的路径，或者文件名拼写错误。建议使用绝对路径或明确的工作目录管理。

  5. 使用Pickle加载不受信任的文件：Pickle文件可以包含恶意代码，切勿加载来源不明的.pkl文件。Joblib和ONNX相对安全，但在生产环境中，任何文件都应来自可信源。

小结

本课我们学习了将训练好的机器学习模型“固化”下来的关键技术：

  • Pickle：Python原生，简单通用，适合快速原型和简单的模型序列化。
  • Joblib：针对科学计算优化，是保存Scikit-learn模型的首选，效率更高。
  • ONNX：一种开放的中间表示格式，核心优势是跨平台、跨框架，是模型部署到生产环境（特别是移动端、嵌入式设备或不同技术栈的服务器）的强大工具。

选择建议：

  • 对于实验、开发和简单的模型共享，使用 Joblib。
  • 当你需要将模型部署到非Python环境、或追求极致的推理性能和跨语言支持时，选择 ONNX。
  • 永远记住保存与模型配套的预处理流程，并保持运行环境的一致性。

掌握模型持久化，你才能将机器学习的成果从“实验室”带到“现实世界”。下一课，我们将学习如何用Flask将保存好的模型包装成API，实现真正的服务化部署。

练习编辑器

rust
Loading...

继续学习

完成本课后，建议继续学习下一课「Flask 构建模型推理 API」 以巩固所学知识。