跳到主要内容

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

自动日志记录是一项强大的功能,可让您在无需显式日志语句的情况下记录指标、参数和模型。您所要做的就是在训练代码之前调用 mlflow.autolog()

import mlflow

mlflow.autolog()

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

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

  • 指标 - MLflow 会根据您使用的模型和库预先选择一组要记录的指标
  • 参数 - 训练指定的超参数,以及在未显式设置时库提供的默认值
  • 模型签名 - 记录 模型签名 实例,该实例描述了模型的输入和输出架构
  • 制品 - 例如模型检查点
  • 数据集 - 用于训练的数据集对象(如果适用),例如 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

然后,在浏览器中导航到 https://: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) 使用 scikit-learn 的插件并设置 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() 来显式启用(或禁用)特定库的自动日志记录。

以下列表包含 MLflow 中支持自动日志记录的最流行库:

注意

还有许多其他集成支持自动日志记录,并且支持的库列表不断增长。请参阅专门的页面,以获取有关给定库是否支持自动日志记录的更多指导。

对于自动将模型保存为制品的插件,还会记录用于依赖关系管理的附加文件

Keras/TensorFlow

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

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

自动日志记录捕获以下信息:

框架指标参数标签工件
tf.keras训练损失;验证损失;用户指定的指标fit() 参数;优化器名称;学习率;epsilon--训练开始时的模型摘要;MLflow 模型(Keras 模型);训练结束时的 TensorBoard 日志
tf.keras.callbacks.EarlyStopping来自 EarlyStopping 回调的指标。例如,stopped_epochrestored_epochrestore_best_weight来自 EarlyStoppingfit() 参数。例如,min_deltapatiencebaselinerestore_best_weights----

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

如果在 autolog() 捕获数据时已存在运行,MLflow 将记录到该运行,但训练后不会自动结束该运行。如果您希望开始新的运行上下文来记录到新运行,则必须手动停止运行。

LightGBM

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

自动日志记录捕获以下信息:

框架指标参数标签工件
LightGBM用户指定的指标lightgbm.train 参数--训练结束时的MLflow 模型(LightGBM 模型),带有模型签名;特征重要性;输入示例;

如果启用了提前停止,最佳迭代时的指标将作为额外的步骤/迭代进行记录。

Paddle

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

自动日志记录捕获以下信息:

框架指标参数标签工件
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 训练代码之前调用通用的自动日志记录函数 mlflow.pytorch.autolog() 以启用指标、参数和模型的自动日志记录。请参阅此处的示例用法。请注意,目前 PyTorch 自动日志记录仅支持使用 PyTorch Lightning 训练的模型。

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

框架/模块指标参数标签工件
pytorch_lightning.trainer.Trainer训练损失;验证损失;平均测试准确率;用户定义的指标fit() 参数;优化器名称;学习率;epsilon。--训练开始时的模型摘要,训练结束时的MLflow 模型(PyTorch 模型);
pytorch_lightning.callbacks.earlystopping训练损失;验证损失;平均测试准确率;用户定义的指标。来自 EarlyStopping 回调的指标。例如,spotted_epochrestored_epochrestore_best_weight 等。fit() 参数;优化器名称;学习率;epsilon。来自 EarlyStopping 回调的参数。例如,min_deltapatiencebaselinerestore_best_weights 等。--训练开始时的模型摘要;训练结束时的MLflow 模型(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")),然后调用通用的自动日志记录函数 mlflow.spark.autolog() 以在读取时启用 Spark 数据源信息的自动日志记录,无需显式日志语句。请注意,Spark ML(MLlib)模型的自动日志记录尚不支持。

自动日志记录捕获以下信息:

框架指标参数标签工件
Spark----包含源路径、版本、格式的单个标签。该标签包含每个数据源的一行--
注意
  • 此外,Spark 数据源自动日志记录是异步进行的——因此,在启动短暂的 MLflow 运行时,可能会(尽管不太可能)出现竞争条件,导致数据源信息未被记录。
重要提示

对于 Pyspark 3.2.0 或更高版本,Spark 数据源自动日志记录需要将 PYSPARK_PIN_THREAD 环境变量设置为 false

Statsmodels

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

自动日志记录捕获以下信息:

框架指标参数标签工件
Statsmodels用户指定的指标statsmodels.base.model.Model.fit 参数--训练结束时的MLflow 模型(statsmodels.base.wrapper.ResultsWrapper)
注意
  • 每个重写 fit 方法的模型子类都将期望并记录其自己的参数。

XGBoost

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

自动日志记录捕获以下信息:

框架指标参数标签工件
XGBoost用户指定的指标xgboost.train 参数--训练结束时的MLflow 模型(XGBoost 模型),带有模型签名;特征重要性;输入示例

如果启用了提前停止,最佳迭代时的指标将作为额外的步骤/迭代进行记录。