使用MLflow无法记录Pipeline对象，如何保存完整GridSearchCV实例？

你当前只能拿到最优Pipeline的核心原因是：你调用mlflow.sklearn.log_model时传入的是best_grid.best_estimator_，仅保存了网格搜索筛选出的最优模型，而非完整的GridSearch实例；同时你使用mlflow.pyfunc.load_model加载，该方法返回的是封装后的预测接口，会丢失原sklearn类的所有非预测属性。

方案1：直接保存完整GridSearch实例（兼容性最好） #

GridSearchCV本身是符合sklearn规范的估计器，MLflow的sklearn日志能力支持直接保存完整实例：

1. 修改日志代码，直接传入best_grid而非best_estimator_：

# 替换原有的mlflow.sklearn.log_model(fitted_model, "tester")
mlflow.sklearn.log_model(best_grid, "full_grid_search")

2. 加载时使用mlflow.sklearn.load_model而非pyfunc加载器，即可还原完整GridSearch对象：

import mlflow.sklearn
# <run_id>替换为你的训练运行ID
loaded_grid = mlflow.sklearn.load_model("runs:/<run_id>/full_grid_search")
# 可正常访问所有GridSearch属性
print(loaded_grid.cv_results_) # 全量交叉验证结果
print(loaded_grid.best_params_) # 最优参数配置
print(loaded_grid.best_estimator_) # 最优Pipeline

方案2：序列化保存自定义对象（适合保存HyperclassifierSearch实例） #

如果需要保存HyperclassifierSearch封装类的完整实例，或者担心第三方封装类和MLflow sklearn日志的兼容性，可以用通用artifact能力配合pickle序列化：

保存代码 #

import pickle
with mlflow.start_run():
    # 原有训练逻辑不变
    search = HyperclassifierSearch(models, params)
    best_grid = search.train_model(X_train, y_train, cv=3, scoring='accuracy')
    results = search.evaluate_model()
    # 序列化对象到本地文件
    with open("full_search_obj.pkl", "wb") as f:
        # 可根据需求选择保存search（HyperclassifierSearch实例）或best_grid（GridSearch实例）
        pickle.dump(search, f)
    # 上传文件到MLflow artifact存储
    mlflow.log_artifact("full_search_obj.pkl")

加载代码 #

import pickle
import mlflow
# 下载对应的artifact文件
artifact_path = mlflow.artifacts.download_artifacts("runs:/<run_id>/full_search_obj.pkl")
# 反序列化得到完整对象
with open(artifact_path, "rb") as f:
    loaded_search_obj = pickle.load(f)

注意事项 #

• 用pickle序列化的对象，加载环境的依赖版本（包括HyperclassifierSearch、sklearn、xgboost等）必须和训练环境完全一致，否则会出现反序列化失败问题
• MLflow的sklearn autolog只会自动记录GridSearch的参数、指标、最优模型，不会自动保存完整GridSearch实例，需要手动显式执行保存操作

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