企业官网建设流程全解析

MLflow XGBoost Scikit-learn 模型自动日志记录实战：以糖尿病数据集回归为例

【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow

本文基于仓库 examples/xgboost/xgboost_sklearn 中的官方示例，讲解如何使用mlflow.xgboost.autolog()对 XGBoost 的 Scikit-learn 风格估计器（如XGBRegressor）进行超参数、评估指标与模型工件的自动记录，并打通模型保存与加载的完整闭环。读完本文，你将掌握在 MLflow 项目中一键启用 XGBoost 自动日志、运行脚本后读取 Run 内数据，以及用原生 XGBoost 与pyfunc两种方式加载已记录模型的完整实战方法。

示例概览：训练什么、记录什么

该示例使用 XGBoost 的 Scikit-learn API 训练一个XGBRegressor回归模型，数据集为 scikit-learn 内置的糖尿病数据集（load_diabetes），训练完成后自动记录：

• 超参数：n_estimators、reg_lambda、gamma、max_depth等模型配置；
• 评估指标：训练过程中基于eval_set产生的逐轮评估结果（如 RMSE）；
• 训练后的模型：以 XGBoost flavor 记录到当前 Run，并附带输入示例与推断出的模型签名（signature）；
• 特征重要性：以 JSON 文件与 PNG 图两种形式记录为工件（artifact）。

示例目录共包含四个文件，结构如下：

examples/xgboost/xgboost_sklearn/ ├── MLproject # MLflow Project 定义，指定入口命令与 python_env ├── README.md # 示例说明 ├── python_env.yaml # 运行环境依赖声明 ├── train.py # 训练 + 自动日志主脚本 └── utils.py # 辅助函数：从 Run 中取回已记录的 params/metrics/tags/artifacts

环境准备：声明依赖

python_env.yaml 声明了运行本示例所需的依赖：

build_dependencies: - pip dependencies: - mlflow - pandas - scikit-learn - xgboost

其中pandas与scikit-learn是为load_diabetes(as_frame=True)、train_test_split以及 Scikit-learn 风格估计器服务的，xgboost提供模型本身，mlflow提供跟踪与模型管理能力。若将本示例作为 MLflow Project 运行，该环境会被自动构建。

核心脚本逐行拆解

train.py 是完整可运行的训练脚本，其执行流程如下。

1. 准备数据集

X, y = load_diabetes(return_X_y=True, as_frame=True) X_train, X_test, y_train, y_test = train_test_split(X, y)

as_frame=True将特征返回为 Pandas DataFrame，便于后续 MLflow 以结构化数据集（dataset）的形式记录训练数据；默认按 75% / 25% 划分训练集与测试集。

2. 启用 XGBoost 自动日志

mlflow.xgboost.autolog()

这是整个示例的关键一行。mlflow.xgboost.autolog()同时作用于 XGBoost 的原生 API（xgboost.train）与Scikit-learn API（xgboost.sklearn下的估计器）。从源码 mlflow/xgboost/init.py 的 docstring 可以看到，它会自动记录：

• xgboost.train中指定的参数；
• 指定了evals时每个迭代轮的评估指标；
• 指定了early_stopping_rounds时的最佳迭代指标（best_iteration、stopped_iteration）；
• 特征重要性（JSON 文件与图）；
• 训练好的模型及其输入示例与推断签名。

3. 训练模型

regressor = xgb.XGBRegressor(n_estimators=20, reg_lambda=1, gamma=0, max_depth=3) regressor.fit(X_train, y_train, eval_set=[(X_test, y_test)])

这里刻意指定了eval_set，目的是让自动日志能在每个迭代轮记录测试集上的评估指标。fit()返回后，模型工件会被记录到 MLflow 自动创建的 Run 中。

4. 预测与取回 Run 信息

y_pred = regressor.predict(X_test) mean_squared_error(y_test, y_pred) run_id = mlflow.last_active_run().info.run_id print(f"Logged data and model in run {run_id}")

由于autolog()会自动管理 Run，训练完成后可通过mlflow.last_active_run()获取当前 Run 的 ID，再交给fetch_logged_data(run_id)把已记录的数据打印出来：

for key, data in fetch_logged_data(run_id).items(): print(f"\n---------- logged {key} ----------") pprint(data)

fetch_logged_data定义在 utils.py 中，内部通过MlflowClient().get_run(run_id).data取回该 Run 的 params、metrics 与 tags（并过滤掉以mlflow.开头的系统标签），再通过client.list_artifacts(run_id)递归列出全部工件路径，最终返回一个包含params、metrics、tags、artifacts四个键的字典，方便直观核对自动日志到底记录了哪些内容。

autolog() 是如何同时覆盖 Scikit-learn 估计器的

从源码层面看，mlflow/xgboost/init.py 中的autolog()使用safe_patch对三类对象打了补丁：

• xgboost.train：patch 后接管原生训练流程，负责参数、逐轮指标、特征重要性与 Booster 模型记录；
• xgboost.sklearn.train：关闭其模型记录（functools.partial(train, False, log_datasets)），因为 Scikit-learn 风格模型要作为完整估计器对象保存，而不是 Booster；
• xgboost.DMatrix.__init__：截获训练数据用于生成输入示例（input example）。

最后一行是关键：mlflow.xgboost.autolog()内部还会调用mlflow.sklearn._autolog(flavor_name="xgboost", ...)，从而让 Scikit-learn 风格的XGBRegressor/XGBClassifier等估计器在fit()返回后，以其原生模型类（如xgboost.sklearn.XGBRegressor）被记录到 Run 中。这正是示例 README 中"Saving / loading models also supports XGBoost scikit-learn models"的底层实现。

仓库测试 tests/xgboost/test_xgboost_autolog.py 验证了这一行为：训练XGBRegressor后断言参数与指标被记录、feature_importance_weight.png与feature_importance_weight.json两个工件存在，并通过mlflow.xgboost.load_model(f"runs:/{run_id}/model")加载模型后与原始模型预测一致。

autolog() 常用参数

autolog()提供若干可调参数，常用配置如下：

例如想要在训练的同时把模型注册进模型仓库，可写成：

mlflow.xgboost.autolog(registered_model_name="diabetes_xgb", importance_types=["weight", "gain"])

模型保存与加载：不只是自动日志

示例 README 特别指出模型保存/加载同样支持 XGBoost Scikit-learn 模型，这对应mlflow.xgboost模块中的几个公开 API：

• save_model(xgb_model, path, ...)：将模型保存到本地路径，内部调用xgb_model.save_model(model_data_path, **kwargs)，并同时生成MLmodel、conda.yaml、requirements.txt与python_env.yaml（见 mlflow/xgboost/init.py）；
• log_model(xgb_model, name, ...)：将模型作为当前 Run 的工件记录，返回包含元数据的ModelInfo（mlflow/xgboost/init.py）；
• load_model(model_uri)：从本地路径或runs:/<run_id>/model等 URI 加载模型，返回的类型取决于保存时的模型类（xgboost.Booster或 Scikit-learn 估计器），见 mlflow/xgboost/init.py 与底层_load_model中的model_class恢复逻辑（mlflow/xgboost/init.py）。

此外，save_model/log_model/autolog均支持model_format参数，默认"ubj"（UBJSON），官方推荐用于获得最佳性能与跨平台兼容性，也可选用"json"或"xgb"。

由于 XGBoost flavor 自动附加了pyfunc兼容层（pyfunc.add_to_model(...)），保存后的模型同样可用于通用的 pyfunc 部署与批量推理。加载示例：

import mlflow # 方式一：按 XGBoost 原生模型加载 model = mlflow.xgboost.load_model("runs:/<run_id>/model") # 方式二：按 pyfunc 通用模型加载，得到可直接 predict(DataFrame) 的包装器 pyfunc_model = mlflow.pyfunc.load_model("runs:/<run_id>/model") predictions = pyfunc_model.predict(X_test)

运行方式

最简单的方式是直接在示例目录下运行：

python train.py

运行结束后控制台会打印类似Logged data and model in run <run_id>的信息，随后依次输出该 Run 中记录的 params、metrics、tags 与 artifacts 列表。

也可以作为 MLflow Project 运行，由 MLproject 定义入口：

mlflow run examples/xgboost/xgboost_sklearn

MLproject 中通过python_env: python_env.yaml指向环境声明，入口命令为python train.py。

查看与对比实验结果

训练产生的 Run 默认写入本地mlruns目录，启动 MLflow UI 即可查看：

mlflow server

打开 UI 后可以浏览本次 Run 的指标曲线、超参数、特征重要性图等工件，并将多次不同超参数配置的 Run 进行对比。配合autolog()自动记录的逐轮指标，还能在图表中直接观察模型在测试集上的收敛过程。

小结

• 通过一行mlflow.xgboost.autolog()，即可让 XGBoost 的 Scikit-learn 风格估计器自动完成超参数、逐轮指标、特征重要性与模型工件的记录，无需手写log_param/log_metric/log_model；
• 自动日志在底层同时 patch 了原生xgboost.train、xgboost.sklearn.train与DMatrix构造，并将 Scikit-learn 风格模型的保存交给mlflow.sklearn._autolog处理，确保模型以原始估计器类被保存；
• 记录后的模型可通过mlflow.xgboost.load_model以原生形式加载，也可通过mlflow.pyfunc.load_model用于通用部署场景。

可继续阅读 examples/xgboost/xgboost_native/README.md 了解 XGBoost 原生 API（xgboost.train）对应的自动日志示例，两者共用mlflow.xgboostflavor 的底层实现。

【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow