PyCaret训练的Pickle模型本地保存报错，求问题排查方案

问题描述 #

我有一个通过PyCaret训练得到的Pickle格式模型，尝试调用save_model(model, model_name)保存到本地时触发错误：NameError: name 'prep_pipe' is not defined，报错发生在代码中标注的行。

相关代码：

from asyncore import read
import logging
import pickle
import mlflow
from numpy import save
from setuptools import setup
import wrapper
import os
import sqlite3
import sqlalchemy
import sys
from mlflow.models.signature import infer_signature
from pycaret.classification import save_model

class MlflowModelService:

   def saveModel(self,model,variant,readable_model_id,preprocess_file_path=None):
    print('inside storeModel of model service.......')
    readable_model_id = readable_model_id.replace("/","__$__")
    model_name = "Original-Model"
    with mlflow.start_run() as active_run:#mlflow work starts
        active_run = mlflow.active_run()

        #mlflow.keras.save_model(model,model_name) #Save a scikit-learn model to a path on the local file system
        #print(model , model_name)
        #exp_clf101 = setup(data = dataset, target = 'result', use_gpu=False, silent=True)
        save_model(model , model_name)#ERROR ON THIS LINE
        pyfunc_model_uri = self.logModel(readable_model_id,model_name,preprocess_file_path)
        self.registerModel(pyfunc_model_uri,readable_model_id)

原因分析 #

PyCaret的save_model函数依赖训练过程中生成的预处理管道（prep_pipe），这个管道是在PyCaret的setup()函数执行时创建并绑定到当前会话上下文的。如果直接传入一个孤立的Pickle模型对象（未关联PyCaret训练会话的上下文或预处理管道），save_model找不到对应的prep_pipe就会触发该错误。

你的代码中注释掉了setup()的执行行，这是核心问题——save_model必须在setup()初始化的会话环境中才能正常工作。

解决方法 #

方案1：复用训练时的PyCaret会话 #

如果模型是在当前代码会话中训练得到的，取消注释setup()行，确保在调用save_model前完成PyCaret会话初始化：

# 恢复setup执行，需替换为你的数据集和目标列
exp_clf101 = setup(data = dataset, target = 'result', use_gpu=False, silent=True)
save_model(model , model_name)

方案2：保存孤立的Pickle模型（无需PyCaret上下文） #

如果模型是从外部加载的Pickle文件，直接用Python原生pickle保存即可，不需要依赖PyCaret的save_model：

with open(f"{model_name}.pkl", "wb") as f:
    pickle.dump(model, f)

方案3：重新关联PyCaret上下文（针对外部加载的模型） #

如果必须用save_model保存（需要保留PyCaret完整流程），需重新初始化PyCaret会话并关联模型：

from pycaret.classification import setup, load_model

# 用和训练时一致的数据集结构初始化会话
setup(data=your_dataset, target=your_target_column, silent=True, use_gpu=False)
# 加载外部Pickle模型
loaded_model = pickle.load(open("your_model_file.pkl", "rb"))
# 现在可以正常调用save_model
save_model(loaded_model, model_name)

注意事项 #

• PyCaret的save_model保存的是包含预处理管道+模型的完整流程，不是单独的模型对象，因此不能直接传入孤立的Pickle模型。
• 如果要结合MLflow使用，建议在PyCaret训练阶段就开启MLflow跟踪（setup()中设置log_experiment=True），后续保存和注册会更顺畅。

内容的提问来源于stack exchange，提问作者Ashu mishra