使用 MLflow 跟踪进行自动日志记录

自动记录是一个强大的功能，它允许你记录指标、参数和模型，而无需显式的日志记录语句。你只需要在训练代码之前调用 mlflow.autolog() 。

import mlflow

mlflow.autolog()

with mlflow.start_run():
    # your training code goes here
    ...

这将使 MLflow 能够自动记录有关您的运行的各种信息，包括：

  • 指标 - MLflow 根据您使用的模型和库预先选择一组要记录的指标

  • 参数 - 为训练指定的超参数，如果未明确设置，则由库提供默认值

  • 模型签名 - 记录 模型签名 实例，该实例描述了模型的输入和输出模式

  • Artifacts - 例如模型检查点

  • 数据集 - 用于训练的数据集对象（如果适用），例如 tensorflow.data.Dataset

如何开始

步骤 1 - 获取 MLflow

MLflow 可以在 PyPI 上获取。如果你还没有在系统上安装它，你可以通过以下命令安装：

pip install mlflow

步骤 2 - 在您的代码中插入 mlflow.autolog

例如，以下代码片段展示了如何为 scikit-learn 模型启用自动日志记录：

import mlflow

from sklearn.model_selection import train_test_split
from sklearn.datasets import load_diabetes
from sklearn.ensemble import RandomForestRegressor

mlflow.autolog()

db = load_diabetes()
X_train, X_test, y_train, y_test = train_test_split(db.data, db.target)

rf = RandomForestRegressor(n_estimators=100, max_depth=6, max_features=3)
# MLflow triggers logging automatically upon model fitting
rf.fit(X_train, y_train)

步骤 3 - 执行你的代码 

python YOUR_ML_CODE.py

步骤 4 - 在 MLflow UI 中查看您的结果

一旦你的训练任务完成，你可以运行以下命令来启动 MLflow UI：

mlflow ui --port 8080

然后，在浏览器中导航到 http://localhost:8080 以查看结果。

自定义自动记录行为

您还可以通过向 mlflow.autolog() 函数传递参数来控制自动日志记录的行为。例如，您可以禁用模型检查点的日志记录，并将标签与您的运行关联，如下所示：

import mlflow

mlflow.autolog(
    log_model_signatures=False,
    extra_tags={"YOUR_TAG": "VALUE"},
)

请参阅 mlflow.autolog() 以获取您可以使用的完整参数集。

启用/禁用特定库的自动日志记录

一个常见的用例是启用/禁用特定库的自动日志记录。例如，如果你在 PyTorch 上训练模型，但使用 scikit-learn 进行数据预处理，你可能希望禁用 scikit-learn 的自动日志记录，同时保持对 PyTorch 的启用。你可以通过以下两种方式实现这一点：(1) 仅使用 PyTorch 风格为 PyTorch 启用自动日志记录 (2) 使用 disable=True 通过其风格禁用 scikit-learn 的自动日志记录。

import mlflow

# Option 1: Enable autologging only for PyTorch
mlflow.pytorch.autolog()

# Option 2: Disable autologging for scikit-learn, but enable it for other libraries
mlflow.sklearn.autolog(disable=True)
mlflow.autolog()

支持的库

备注

通用自动记录函数 mlflow.autolog() 在你导入它时，会为每个已安装的支持库启用自动记录。或者，你可以使用特定于库的自动记录调用，例如 mlflow.pytorch.autolog()，来显式启用（或禁用）特定库的自动记录。

以下库支持自动日志记录：

对于自动将模型保存为工件的版本，会记录用于依赖管理的 附加文件。

Fastai

在训练代码之前调用通用的 autolog 函数 mlflow.fastai.autolog() 以启用指标和参数的自动记录。查看与 Fastai 的示例用法。

Autologging 捕获以下信息：

框架

指标

参数

标签

工件

fastai

用户指定的指标

将优化器数据记录为参数。例如，epochs、lr、opt_func 等；记录 EarlyStoppingCallback 和 OneCycleScheduler 回调的参数

–

模型检查点被记录到‘models’目录；训练结束时的 `MLflow Model`_（fastai Learner模型）；模型摘要文本被记录

Gluon

在训练代码之前调用通用的 autolog 函数 mlflow.gluon.autolog() 以启用指标和参数的自动记录。参见 Gluon 的示例用法。

Autologging 捕获以下信息：

框架

指标

参数

标签

工件

Gluon

训练损失；验证损失；用户指定的指标

层数；优化器名称；学习率；epsilon

–

MLflow 模型 (Gluon 模型); 在训练结束时

Keras/TensorFlow

在训练代码之前调用通用的 autolog 函数或 mlflow.tensorflow.autolog() 以启用指标和参数的自动记录。例如，尝试运行 Keras/Tensorflow 示例。

请注意，仅支持 tensorflow>=2.3 版本。与 tf.estimator 和 EarlyStopping 相关的指标会自动记录。例如，尝试运行 Keras/TensorFlow 示例。

Autologging 捕获以下信息：

框架/模块

指标

参数

标签

工件

tf.keras

训练损失；验证损失；用户指定的指标

fit() 参数；优化器名称；学习率；epsilon

–

训练开始时的模型摘要；`MLflow 模型 <https://mlflow.org/docs/latest/models.html>`_（Keras 模型）；训练结束时的 TensorBoard 日志

tf.keras.callbacks.EarlyStopping

来自 EarlyStopping 回调的指标。例如， stopped_epoch 、 restored_epoch 、 restore_best_weight 等

EarlyStopping 中的 fit() 参数。例如，min_delta、patience、baseline、restore_best_weights 等

–

–

如果在 autolog() 捕获数据时没有活动的运行，MLflow 将自动创建一个运行来记录信息。此外，MLflow 将在通过调用 tf.keras.fit() 完成训练后自动结束运行。

如果在 autolog() 捕获数据时已经存在一个运行，MLflow 将记录到该运行，但在训练后不会自动结束该运行。如果你希望为记录到一个新的运行开始一个新的运行上下文，你将需要手动停止该运行。

LangChain

在训练代码之前调用通用 autolog 函数 mlflow.langchain.autolog() 以启用自动记录跟踪。更多详情请参见 LangChain 自动记录。

Autologging 捕获以下信息：

框架

指标

参数

标签

工件

LangChain

–

–

–

  • 跟踪

  • `MLflow 模型`_（LangChain 模型）在训练端带有模型签名

  • 输入示例

LlamaIndex

在你的训练代码之前调用通用的 autolog 函数 mlflow.llama_index.autolog() 以启用自动记录跟踪。

Autologging 捕获以下信息：

框架

指标

参数

标签

工件

LlamaIndex

–

–

–

  • 跟踪

LightGBM

在你的训练代码之前调用通用的 autolog 函数 mlflow.lightgbm.autolog() 以启用指标和参数的自动记录。

Autologging 捕获以下信息：

框架

指标

参数

标签

工件

LightGBM

用户指定的指标

lightgbm.train 参数

–

`MLflow 模型`_（LightGBM 模型）在训练结束时带有模型签名；特征重要性；输入示例

如果启用了早停，最佳迭代中的指标将被记录为额外的步骤/迭代。

OpenAI

在训练代码之前调用通用自动记录函数 mlflow.openai.autolog() 以启用工件的自动记录。查看与 OpenAI 的示例用法。

Autologging 捕获以下信息：

框架

指标

参数

标签

工件

OpenAI

–

–

–

  • `MLflow 模型`_（OpenAI 模型）在训练端带有模型签名

  • 输入示例

Paddle

在你的训练代码之前调用通用的 autolog 函数 mlflow.paddle.autolog() 以启用指标和参数的自动记录。

Autologging 捕获以下信息：

框架

指标

参数

标签

工件

Paddle

用户指定的指标

paddle.Model.fit 参数

–

MLflow 模型 (Paddle 模型) 在训练端带有模型签名

PySpark

在训练代码之前调用 mlflow.pyspark.ml.autolog() 以启用指标、参数和模型的自动记录。请参阅 PySpark 的示例用法。

pyspark ml 估计器的自动日志记录捕获以下信息：

指标

参数

标签

工件

通过 Evaluator.evaluate 获得的训练后指标

通过 Estimator.fit 获得的参数

  • 类名

  • 完全限定类名

  • 包含拟合估计器的 MLflow 模型

  • metric_info.json 用于训练后指标

PyTorch

在你的 PyTorch Lightning 训练代码之前调用通用的 autolog 函数 mlflow.pytorch.autolog() 以启用自动记录指标、参数和模型。查看示例用法 这里。请注意，目前 PyTorch autologging 仅支持使用 PyTorch Lightning 训练的模型。

自动日志记录在调用 pytorch_lightning.trainer.Trainer.fit 时触发，并捕获以下信息：

框架/模块

指标

参数

标签

工件

pytorch_lightning.trainer.Trainer

训练损失；验证损失；平均测试准确率；用户定义的指标。

fit() 参数；优化器名称；学习率；epsilon。

–

训练开始时的模型摘要，`MLflow 模型 <https://mlflow.org/docs/latest/models.html>`_（PyTorch 模型）在训练结束时；

pytorch_lightning.callbacks.earlystopping

训练损失；验证损失；平均测试准确率；用户定义的指标。来自 EarlyStopping 回调的指标。例如，stopped_epoch、restored_epoch、restore_best_weight 等。

fit() 参数；优化器名称；学习率；来自 EarlyStopping 回调的参数。例如，min_delta、patience、baseline、restore_best_weights 等

–

训练开始时的模型总结；训练结束时的 `MLflow 模型 <https://mlflow.org/docs/latest/models.html>`_（PyTorch 模型）；如果训练由于早期停止回调而停止，则保存最佳的 PyTorch 模型检查点。

如果在 autolog() 捕获数据时没有活动的运行，MLflow 将自动创建一个运行来记录信息，并在 pytorch_lightning.trainer.Trainer.fit() 调用完成后结束运行。

如果在 autolog() 捕获数据时已经存在运行，MLflow 将记录到该运行，但在训练后不会自动结束该运行。

备注

  • 在使用 pytorch_lightning.trainer.Trainer.fit() 时，用户未明确传递的参数（使用默认值的参数）目前不会自动记录

  • 在多优化器场景下（例如使用自编码器），仅记录第一个优化器的参数

Scikit-learn

在你的训练代码之前调用 mlflow.sklearn.autolog() 以启用sklearn指标、参数和模型的自动记录。查看示例用法 这里。

估计器（例如 LinearRegression）和元估计器（例如 Pipeline）的自动记录会创建一个运行并记录：

指标

参数

标签

工件

通过 estimator.score 获得的训练分数

通过 estimator.get_params 获得的参数

  • 类名

  • 完全限定类名

拟合估计器

参数搜索估计器的自动记录（例如 GridSearchCV）会创建一个父运行和嵌套的子运行

- Parent run
  - Child run 1
  - Child run 2
  - ...

包含以下数据：

运行类型

指标

参数

标签

工件

父级

训练得分

  • 参数搜索 估计器的参数

  • 最佳参数组合

  • 类名

  • 完全限定类名

  • 拟合参数搜索估计器

  • 拟合最佳估计器

  • 搜索结果 csv 文件

子

每个参数组合的CV测试分数

每个参数组合

  • 类名

  • 完全限定类名

–

Spark

使用附加了 mlflow-spark JAR 的 SparkSession 进行初始化（例如 SparkSession.builder.config("spark.jars.packages", "org.mlflow.mlflow-spark")），然后调用通用 autolog 函数 mlflow.spark.autolog() 以在读取时自动记录 Spark 数据源信息，无需显式日志语句。请注意，Spark ML（MLlib）模型的自动日志记录尚未支持。

Autologging 捕获以下信息：

框架

指标

参数

标签

工件

Spark

–

–

包含源路径、版本、格式的单个标签。该标签每行包含一个数据源。

–

备注

  • 此外，Spark 数据源自动记录是异步进行的 - 因此，在启动短生命周期的 MLflow 运行时，可能会（尽管不太可能）出现竞争条件，导致数据源信息未被记录。

重要

使用 Pyspark 3.2.0 或更高版本时，Spark 数据源自动记录需要将 PYSPARK_PIN_THREAD 环境变量设置为 false。

Statsmodels

在你的训练代码之前调用通用的 autolog 函数 mlflow.statsmodels.autolog() 以启用指标和参数的自动记录。

Autologging 捕获以下信息：

框架

指标

参数

标签

工件

Statsmodels

用户指定的指标

statsmodels.base.model.Model.fit 参数

–

MLflow 模型 (statsmodels.base.wrapper.ResultsWrapper) 在训练结束时

备注

  • 每个重写 fit 的模型子类都期望并记录其自己的参数。

XGBoost

在训练代码之前调用通用 autolog 函数 mlflow.xgboost.autolog() 以启用指标和参数的自动记录。

Autologging 捕获以下信息：

框架

指标

参数

标签

工件

XGBoost

用户指定的指标

xgboost.train 参数

–

`MLflow 模型`_（XGBoost 模型）在训练端带有模型签名；特征重要性；输入示例

如果启用了早停，最佳迭代中的指标将被记录为额外的步骤/迭代。