模型持久化

模型持久化是将训练好的机器学习模型保存到磁盘，以便后续加载和重用的过程。在生产环境中，模型训练和预测通常在不同的时间和环境中进行，因此模型持久化是机器学习工作流程中不可或缺的一环。本章将详细介绍如何正确地保存和加载 sklearn 模型。

为什么需要模型持久化？

实际应用场景

离线训练，在线预测：在生产环境中，模型通常在离线环境训练完成后部署到在线服务。训练可能需要数小时甚至数天，而预测需要在毫秒级完成。持久化使得训练和预测解耦成为可能。

模型版本管理：不同时间训练的模型可能效果不同。通过持久化，可以保存多个版本的模型，方便对比和回滚。

共享与协作：团队成员可以共享训练好的模型，而不需要重复训练。

节省计算资源：复杂模型的训练成本高昂，持久化避免了每次使用都重新训练。

持久化的基本要求

一个完善的持久化方案应该满足：

• 完整性：保存模型的所有必要信息，包括学习到的参数、预处理步骤等
• 可复现性：加载后的模型行为应与保存前一致
• 安全性：避免加载恶意构造的模型文件
• 兼容性：考虑版本升级带来的兼容性问题

joblib：推荐的持久化工具

sklearn 官方推荐使用 joblib 进行模型持久化。相比 Python 内置的 pickle，joblib 对大型 numpy 数组的处理更高效。

基本用法

import joblib
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris

# 训练模型
iris = load_iris()
X, y = iris.data, iris.target
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X, y)

# 保存模型
joblib.dump(model, 'random_forest_model.joblib')

# 加载模型
loaded_model = joblib.load('random_forest_model.joblib')

# 验证加载后的模型
print(f"模型类型: {type(loaded_model)}")
print(f"预测结果一致性: {all(model.predict(X) == loaded_model.predict(X))}")

joblib vs pickle

特性	joblib	pickle
大型数组效率	高（使用内存映射）	一般
压缩支持	内置支持	需要额外处理
sklearn 推荐	是	否
通用性	专门针对科学计算	Python 通用序列化

joblib 的优势：sklearn 模型通常包含大量 numpy 数组（如系数矩阵、特征重要性等）。joblib 使用内存映射技术处理这些数组，避免了不必要的数据复制，效率更高。

压缩存储

对于大型模型，可以使用压缩来节省存储空间：

# 使用压缩存储
joblib.dump(model, 'random_forest_model_compressed.joblib', compress=3)

# compress 参数说明：
# 0 或 False：不压缩（默认）
# 1-9：压缩级别，数值越大压缩率越高但速度越慢
# True：相当于 compress=3

保存 Pipeline

在实际项目中，数据预处理和模型训练通常组成 Pipeline。保存 Pipeline 可以确保预处理步骤和模型参数的一致性。

完整 Pipeline 的持久化

from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import load_iris
import joblib

# 创建 Pipeline
pipeline = Pipeline([
    ('scaler', StandardScaler()),
    ('classifier', LogisticRegression(max_iter=200))
])

# 训练
iris = load_iris()
X, y = iris.data, iris.target
pipeline.fit(X, y)

# 保存整个 Pipeline
joblib.dump(pipeline, 'full_pipeline.joblib')

# 加载 Pipeline
loaded_pipeline = joblib.load('full_pipeline.joblib')

# 直接对新数据进行预测（自动包含预处理）
X_new = [[5.1, 3.5, 1.4, 0.2]]
prediction = loaded_pipeline.predict(X_new)
probability = loaded_pipeline.predict_proba(X_new)

print(f"预测类别: {prediction}")
print(f"预测概率: {probability}")

为什么要保存整个 Pipeline？

保存完整的 Pipeline 而不是单独保存模型有以下好处：

• 预处理参数保留：StandardScaler 的均值和标准差会被自动保存
• 一致性保证：预测时的数据处理与训练时完全相同
• 避免数据泄露：确保预处理参数只来自训练数据
• 简化部署：加载一个对象即可进行端到端预测

保存自定义对象

当模型包含自定义类或函数时，持久化需要特别注意。

可序列化的条件

Python 对象可以被 pickle 序列化的条件：

• 在模块顶层定义的类和函数
• 不依赖运行时状态的对象

常见问题及解决

问题：lambda 函数无法序列化

# 错误示例
from sklearn.preprocessing import FunctionTransformer
import joblib

# lambda 函数无法直接序列化
transformer = FunctionTransformer(lambda x: x ** 2)
joblib.dump(transformer, 'transformer.joblib')  # 可能失败

# 正确做法：定义在模块顶层
def square_transform(x):
    return x ** 2

transformer = FunctionTransformer(square_transform)
joblib.dump(transformer, 'transformer.joblib')  # 正确

问题：自定义类需要定义在模块顶层

# 在 .py 文件中定义自定义类
# custom_transformer.py
from sklearn.base import BaseEstimator, TransformerMixin

class CustomScaler(BaseEstimator, TransformerMixin):
    def __init__(self, factor=1.0):
        self.factor = factor
    
    def fit(self, X, y=None):
        self.mean_ = X.mean(axis=0) * self.factor
        return self
    
    def transform(self, X):
        return X - self.mean_

# 在另一个脚本中使用和保存
from custom_transformer import CustomScaler
import joblib

scaler = CustomScaler(factor=0.5)
scaler.fit(X_train)
joblib.dump(scaler, 'custom_scaler.joblib')

安全性考虑

pickle 的安全风险

pickle 可以序列化任意 Python 对象，包括可执行代码。加载恶意构造的 pickle 文件可能导致任意代码执行。

# 警告：永远不要加载不信任来源的模型文件
# 恶意的 pickle 文件可以执行任意代码

# 不安全示例
import pickle

# 假设这是从不可信来源获取的文件
with open('untrusted_model.pkl', 'rb') as f:
    model = pickle.load(f)  # 危险！可能执行恶意代码

安全最佳实践

使用签名验证：对模型文件进行数字签名，加载前验证签名。

import hmac
import hashlib
import joblib

def save_model_with_signature(model, filepath, secret_key):
    """保存模型并附加签名"""
    joblib.dump(model, filepath)
    
    # 计算文件哈希并签名
    with open(filepath, 'rb') as f:
        file_hash = hashlib.sha256(f.read()).hexdigest()
    
    signature = hmac.new(
        secret_key.encode(), 
        file_hash.encode(), 
        hashlib.sha256
    ).hexdigest()
    
    # 保存签名（实际应用中应单独存储）
    with open(filepath + '.sig', 'w') as f:
        f.write(signature)

def load_model_with_signature(filepath, secret_key):
    """验证签名后加载模型"""
    with open(filepath, 'rb') as f:
        file_hash = hashlib.sha256(f.read()).hexdigest()
    
    with open(filepath + '.sig', 'r') as f:
        expected_signature = f.read().strip()
    
    actual_signature = hmac.new(
        secret_key.encode(), 
        file_hash.encode(), 
        hashlib.sha256
    ).hexdigest()
    
    if not hmac.compare_digest(expected_signature, actual_signature):
        raise ValueError("签名验证失败！文件可能被篡改。")
    
    return joblib.load(filepath)

只加载可信来源的模型：在生产环境中，确保模型文件来源可信，并妥善保管存储位置。

版本兼容性

sklearn 版本兼容性

不同版本的 sklearn 可能不兼容：

• 新版本训练的模型可能无法在旧版本中加载
• 模型内部结构可能随版本变化

# 保存时记录版本信息
import sklearn
import json
import joblib

def save_model_with_metadata(model, filepath):
    """保存模型并记录元数据"""
    joblib.dump(model, filepath)
    
    metadata = {
        'sklearn_version': sklearn.__version__,
        'python_version': platform.python_version(),
        'model_type': type(model).__name__,
        'timestamp': datetime.now().isoformat()
    }
    
    with open(filepath + '.meta', 'w') as f:
        json.dump(metadata, f)

# 加载时检查版本
def load_model_with_metadata(filepath):
    """加载模型并验证版本"""
    with open(filepath + '.meta', 'r') as f:
        metadata = json.load(f)
    
    if metadata['sklearn_version'] != sklearn.__version__:
        print(f"警告：模型训练版本 {metadata['sklearn_version']} "
              f"与当前版本 {sklearn.__version__} 不同")
    
    return joblib.load(filepath)

向前兼容的最佳实践

• 记录环境信息：保存模型时记录 sklearn、Python、numpy 等依赖版本
• 使用虚拟环境：为每个项目维护独立的虚拟环境
• 容器化部署：使用 Docker 确保运行环境一致
• 定期重新训练：对于长期运行的系统，定期用新版本重新训练模型

大型模型的持久化

分块存储

对于非常大的模型（如大型随机森林），可以考虑分块存储：

import joblib
from sklearn.ensemble import RandomForestClassifier

# 训练大型模型
rf = RandomForestClassifier(n_estimators=1000, random_state=42)
rf.fit(X_train, y_train)

# 分块存储（将每个树单独保存）
import os
save_dir = 'model_chunks'
os.makedirs(save_dir, exist_ok=True)

# 保存模型元信息
metadata = {
    'n_estimators': rf.n_estimators,
    'n_classes_': rf.n_classes_,
    'n_features_in_': rf.n_features_in_,
    'classes_': rf.classes_.tolist()
}

import json
with open(os.path.join(save_dir, 'metadata.json'), 'w') as f:
    json.dump(metadata, f)

# 单独保存每棵树
for i, estimator in enumerate(rf.estimators_):
    joblib.dump(estimator, os.path.join(save_dir, f'tree_{i}.joblib'))

# 加载时重建模型
def load_chunked_model(save_dir):
    with open(os.path.join(save_dir, 'metadata.json'), 'r') as f:
        metadata = json.load(f)
    
    rf = RandomForestClassifier(n_estimators=metadata['n_estimators'])
    rf.n_classes_ = metadata['n_classes_']
    rf.n_features_in_ = metadata['n_features_in_']
    rf.classes_ = np.array(metadata['classes_'])
    
    rf.estimators_ = []
    for i in range(metadata['n_estimators']):
        tree = joblib.load(os.path.join(save_dir, f'tree_{i}.joblib'))
        rf.estimators_.append(tree)
    
    return rf

使用 HDF5 格式

对于包含大型 numpy 数组的模型，HDF5 格式可能更高效：

import h5py
import numpy as np

def save_to_hdf5(model, filepath):
    """将模型参数保存到 HDF5 文件"""
    with h5py.File(filepath, 'w') as f:
        for key, value in model.__dict__.items():
            if isinstance(value, np.ndarray):
                f.create_dataset(key, data=value)
            elif isinstance(value, (int, float, str)):
                f.attrs[key] = value

完整示例：生产级模型持久化

下面是一个完整的示例，展示生产环境中模型持久化的最佳实践：

import joblib
import json
import sklearn
import platform
from datetime import datetime
from pathlib import Path

class ModelPersistence:
    """生产级模型持久化管理器"""
    
    def __init__(self, base_dir='models'):
        self.base_dir = Path(base_dir)
        self.base_dir.mkdir(exist_ok=True)
    
    def save(self, model, name, version=None, metadata=None):
        """
        保存模型及其元数据
        
        Parameters:
        -----------
        model : sklearn estimator 或 Pipeline
            要保存的模型
        name : str
            模型名称
        version : str, optional
            版本号，默认使用时间戳
        metadata : dict, optional
            额外的元数据
        """
        version = version or datetime.now().strftime('%Y%m%d_%H%M%S')
        model_dir = self.base_dir / name / version
        model_dir.mkdir(parents=True, exist_ok=True)
        
        # 保存模型
        model_path = model_dir / 'model.joblib'
        joblib.dump(model, model_path, compress=3)
        
        # 准备元数据
        full_metadata = {
            'name': name,
            'version': version,
            'created_at': datetime.now().isoformat(),
            'sklearn_version': sklearn.__version__,
            'python_version': platform.python_version(),
            'model_type': type(model).__name__,
            'model_params': model.get_params() if hasattr(model, 'get_params') else {}
        }
        
        if metadata:
            full_metadata.update(metadata)
        
        # 保存元数据
        metadata_path = model_dir / 'metadata.json'
        with open(metadata_path, 'w') as f:
            json.dump(full_metadata, f, indent=2, default=str)
        
        print(f"模型已保存到: {model_dir}")
        return model_dir
    
    def load(self, name, version='latest'):
        """
        加载模型
        
        Parameters:
        -----------
        name : str
            模型名称
        version : str
            版本号，'latest' 加载最新版本
        
        Returns:
        --------
        model : 加载的模型
        metadata : 模型元数据
        """
        model_base = self.base_dir / name
        
        if version == 'latest':
            versions = sorted([d.name for d in model_base.iterdir() if d.is_dir()])
            if not versions:
                raise ValueError(f"未找到模型: {name}")
            version = versions[-1]
        
        model_dir = model_base / version
        model_path = model_dir / 'model.joblib'
        metadata_path = model_dir / 'metadata.json'
        
        if not model_path.exists():
            raise ValueError(f"模型文件不存在: {model_path}")
        
        # 加载模型
        model = joblib.load(model_path)
        
        # 加载元数据
        with open(metadata_path, 'r') as f:
            metadata = json.load(f)
        
        # 版本警告
        if metadata.get('sklearn_version') != sklearn.__version__:
            print(f"警告：模型版本 {metadata.get('sklearn_version')} "
                  f"与当前版本 {sklearn.__version__} 不同")
        
        return model, metadata
    
    def list_models(self, name=None):
        """列出所有可用的模型及其版本"""
        if name:
            model_dir = self.base_dir / name
            if model_dir.exists():
                return sorted([d.name for d in model_dir.iterdir() if d.is_dir()])
            return []
        
        return {d.name: sorted([v.name for v in d.iterdir() if v.is_dir()])
                for d in self.base_dir.iterdir() if d.is_dir()}

# 使用示例
if __name__ == '__main__':
    from sklearn.datasets import load_iris
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.model_selection import train_test_split
    
    # 训练模型
    iris = load_iris()
    X_train, X_test, y_train, y_test = train_test_split(
        iris.data, iris.target, test_size=0.2, random_state=42
    )
    
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)
    
    # 创建持久化管理器
    persistence = ModelPersistence('saved_models')
    
    # 保存模型
    persistence.save(
        model, 
        name='iris_classifier',
        metadata={
            'accuracy': model.score(X_test, y_test),
            'description': '鸢尾花分类器'
        }
    )
    
    # 加载模型
    loaded_model, metadata = persistence.load('iris_classifier')
    
    print(f"加载的模型: {metadata['model_type']}")
    print(f"训练时间: {metadata['created_at']}")
    print(f"测试集准确率: {metadata['accuracy']:.3f}")
    
    # 列出所有模型
    print(f"\n可用模型: {persistence.list_models()}")

小结

模型持久化是机器学习工作流程的重要环节：

1. 使用 joblib：sklearn 官方推荐使用 joblib 进行模型持久化，对大型数组处理更高效
2. 保存完整 Pipeline：确保预处理步骤和模型参数一起保存，保证预测一致性
3. 注意安全性：只加载可信来源的模型，考虑使用签名验证
4. 记录版本信息：保存 sklearn、Python 等版本信息，便于排查兼容性问题
5. 生产级实践：建立完整的模型管理体系，包括版本控制、元数据记录等

正确实现模型持久化是模型部署的基础。下一章我们将学习常见陷阱与最佳实践，帮助避免机器学习项目中的常见错误。