使用 MLflow Tracking 进行自动日志记录
自动日志记录(Auto logging)是一项强大的功能,无需显式编写日志语句即可记录指标、参数和模型。您只需在训练代码前调用 mlflow.autolog() 即可。
import mlflow
mlflow.autolog()
with mlflow.start_run():
# your training code goes here
...
这将使 MLflow 能够自动记录有关运行的各种信息,包括:
- 指标 (Metrics) - MLflow 会根据您使用的模型和库,预先选择一组指标进行记录。
- 参数 (Parameters) - 训练指定的超参数,以及如果未明确设置时由库提供的默认值。
- 模型签名 (Model Signature) - 记录 模型签名 实例,该实例描述了模型的输入和输出模式。
- 工件 (Artifacts) - 例如模型检查点。
- 数据集 (Dataset) - 用于训练的数据集对象(如果适用),例如 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 server --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 风格(flavor)仅为 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() 来显式启用(或禁用)特定库的自动日志记录。
以下列表涵盖了 MLflow 中支持自动日志记录的最流行库:
还有许多其他集成支持自动日志记录,并且支持的库列表在不断增加。请参阅专门的页面以获取关于特定库是否支持自动日志记录的进一步指南。
对于自动将模型保存为工件的库,将记录用于依赖项管理的附加文件。
Keras/TensorFlow
在训练代码前调用通用自动日志记录函数或 mlflow.tensorflow.autolog(),以启用指标和参数的自动日志记录。例如,尝试运行 Keras/Tensorflow 示例。
请注意,仅支持 tensorflow>=2.3 版本。与 tf.estimator 和 EarlyStopping 相关的相应指标会被自动记录。
自动日志记录捕获以下信息:
| 框架 | 指标 | 参数 | 标签 | 工件 |
|---|---|---|---|---|
tf.keras | 训练损失;验证损失;用户指定的指标 | fit() 参数;优化器名称;学习率;epsilon | -- | 训练开始时的模型摘要;MLflow 模型(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 会自动创建一个运行来记录信息。此外,一旦通过调用 tf.keras.fit() 完成训练,MLflow 就会自动结束该运行。
如果 autolog() 捕获数据时已存在运行,MLflow 将记录到该运行中,但不会在训练后自动结束该运行。如果您希望为新运行启动一个新的运行上下文,则必须手动停止该运行。
LightGBM
在训练代码前调用通用自动日志记录函数 mlflow.lightgbm.autolog(),以启用指标和参数的自动日志记录。
自动日志记录捕获以下信息:
| 框架 | 指标 | 参数 | 标签 | 工件 |
|---|---|---|---|---|
| LightGBM | 用户指定的指标 | lightgbm.train 参数 | -- | 训练结束时的 MLflow 模型(LightGBM 模型)及模型签名;特征重要性;输入示例; |
如果启用了提前停止(early stopping),最佳迭代时的指标将记录为额外的步骤/迭代。
Paddle
在训练代码前调用通用自动日志记录函数 mlflow.paddle.autolog(),以启用指标和参数的自动日志记录。
自动日志记录捕获以下信息:
| 框架 | 指标 | 参数 | 标签 | 工件 |
|---|---|---|---|---|
| Paddle | 用户指定的指标 | paddle.Model.fit 参数 | -- | 训练结束时的 MLflow 模型(Paddle 模型)及模型签名 |
PySpark
在训练代码前调用 mlflow.pyspark.ml.autolog() 以启用指标、参数和模型的自动日志记录。请参阅 PySpark 使用示例。
pyspark ml 估计器的自动日志记录捕获以下信息:
| 指标 | 参数 | 标签 | 工件 |
|---|---|---|---|
通过 Evaluator.evaluate 获得的训练后指标 | 通过 Estimator.fit 获得的参数 |
|
|
PyTorch
在 PyTorch Lightning 训练代码前调用通用自动日志记录函数 mlflow.pytorch.autolog(),以启用指标、参数和模型的自动日志记录。示例用法请见此处。请注意,目前 PyTorch 自动日志记录仅支持使用 PyTorch Lightning 训练的模型。
自动日志记录会在调用 pytorch_lightning.trainer.Trainer.fit 时触发,并捕获以下信息:
| 框架/模块 | 指标 | 参数 | 标签 | 工件 |
|---|---|---|---|---|
pytorch_lightning.trainer.Trainer | 训练损失;验证损失;average_test_accuracy;用户定义的指标 | fit() 参数;优化器名称;学习率;epsilon。 | -- | 训练开始时的模型摘要;训练结束时的 MLflow 模型(PyTorch 模型); |
pytorch_lightning.callbacks.earlystopping | 训练损失;验证损失;average_test_accuracy;用户定义的指标。来自 EarlyStopping 回调的指标。例如,stopped_epoch、restored_epoch、restore_best_weight 等 | fit() 参数;优化器名称;学习率;epsilon。来自 EarlyStopping 回调的参数。例如,min_delta、patience、baseline、restore_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
- ...
包含以下数据:
| 运行类型 | 指标 | 参数 | 标签 | 工件 |
|---|---|---|---|---|
| 父运行 (Parent) | 训练得分 |
|
|
|
| 子运行 (Child) | 每个参数组合的 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 模型)及模型签名;特征重要性;输入示例 |
如果启用了提前停止(early stopping),最佳迭代时的指标将记录为额外的步骤/迭代。