使用 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 flavor 为 PyTorch 启用自动日志记录，或 (2) 使用 scikit-learn flavor 并设置 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()）来显式启用（或禁用）特定库的自动日志记录。
以下库支持自动日志记录
对于自动将模型保存为 artifact 的 flavors，将记录用于依赖项管理的附加文件。
Fastai
在您的训练代码之前调用通用的自动日志记录函数 mlflow.fastai.autolog()，以启用指标和参数的自动日志记录。请参阅 Fastai 的示例用法。
自动日志记录会捕获以下信息
| 框架 | 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|---|
| fastai | 用户指定的指标 | 将优化器数据记录为参数。例如，epochs, lr, opt_func 等；记录 EarlyStoppingCallback 和 OneCycleScheduler 回调的参数 | -- | 模型检查点被记录到“models”目录；训练结束时记录 MLflow 模型 (fastai Learner 模型)；模型摘要文本被记录。 |
Keras/TensorFlow
在您的训练代码之前调用通用自动日志记录函数或 mlflow.tensorflow.autolog()，以启用指标和参数的自动日志记录。例如，尝试运行 Keras/TensorFlow 示例。
请注意，仅支持 tensorflow>=2.3 版本。与 tf.estimator 和 EarlyStopping 关联的相应指标会自动记录。例如，尝试运行 Keras/TensorFlow 示例。
自动日志记录会捕获以下信息
| 框架 | 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|---|
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 将日志记录到该运行，但不会在训练后自动结束该运行。如果您希望启动新的运行上下文以日志记录到新的运行，则必须手动停止该运行。
LangChain
在您的训练代码之前调用通用的自动日志记录函数 mlflow.langchain.autolog()，以启用跟踪的自动日志记录。请参阅 LangChain 自动日志记录 以了解更多详细信息。
自动日志记录会捕获以下信息
| 框架 | 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|---|
| LangChain | -- | -- | -- |
|
LlamaIndex
在您的训练代码之前调用通用的自动日志记录函数 mlflow.llama_index.autolog()，以启用跟踪的自动日志记录。
自动日志记录会捕获以下信息
| 框架 | 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|---|
| LlamaIndex | -- | -- | -- |
|
LightGBM
在您的训练代码之前调用通用的自动日志记录函数 mlflow.lightgbm.autolog()，以启用指标和参数的自动日志记录。
自动日志记录会捕获以下信息
| 框架 | 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|---|
| LightGBM | 用户指定的指标 | lightgbm.train 参数 | -- | 训练结束时记录带有模型签名的 MLflow 模型 (LightGBM 模型)；特征重要性；输入示例； |
如果启用了 early stopping，则最佳迭代的指标将作为额外的步骤/迭代记录。
OpenAI
在您的训练代码之前调用通用的自动日志记录函数 mlflow.openai.autolog()，以启用 artifact 的自动日志记录。请参阅 OpenAI 的示例用法。
自动日志记录会捕获以下信息
| 框架 | 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|---|
| OpenAI | -- | -- | -- |
|
Paddle
在您的训练代码之前调用通用的自动日志记录函数 mlflow.paddle.autolog()，以启用指标和参数的自动日志记录。
自动日志记录会捕获以下信息
| 框架 | 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|---|
| Paddle | 用户指定的指标 | paddle.Model.fit 参数 | -- | 训练结束时记录带有模型签名的 MLflow 模型 (Paddle 模型) |
PySpark
在您的训练代码之前调用 mlflow.pyspark.ml.autolog()，以启用指标、参数和模型的自动日志记录。请参阅 PySpark 的示例用法。
pyspark ml estimators 的自动日志记录会捕获以下信息
| 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|
通过 Evaluator.evaluate 获取的训练后指标 | 通过 Estimator.fit 获取的参数 |
|
|
PyTorch
在您的 PyTorch Lightning 训练代码之前调用通用自动日志记录函数 mlflow.pytorch.autolog()，以启用指标、参数和模型的自动日志记录。请参阅此处的示例用法。请注意，目前 PyTorch 自动日志记录仅支持使用 PyTorch Lightning 训练的模型。
自动日志记录会在调用 pytorch_lightning.trainer.Trainer.fit 时触发，并捕获以下信息
| 框架/模块 | 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|---|
pytorch_lightning.trainer.Trainer | 训练损失；验证损失；average_test_accuracy；用户定义的指标 | fit() 参数；优化器名称；学习率；epsilon。 | -- | 训练开始时记录模型摘要，训练结束时记录 MLflow 模型 (PyTorch 模型)； |
pytorch_lightning.callbacks.earlystopping | 训练损失；验证损失；average_test_accuracy；用户定义的指标。来自 EarlyStopping 回调的指标。例如，spotted_epoch, restored_epoch, restore_best_weight 等 | fit() 参数；优化器名称；学习率；epsilon。来自 EarlyStopping 回调的参数。例如，min_delta, patience, baseline, restore_best_weights 等 | -- | 训练开始时记录模型摘要；训练结束时记录 MLflow 模型 (PyTorch 模型)；如果由于 early stopping 回调而停止训练，则记录最佳 PyTorch 模型检查点。 |
如果在 autolog() 捕获数据时没有活动的运行存在，MLflow 将自动创建一个运行来记录信息，并在对 pytorch_lightning.trainer.Trainer.fit() 的调用完成后结束该运行。
如果在 autolog() 捕获数据时已有运行存在，MLflow 将日志记录到该运行，但不会在训练后自动结束该运行。
- 在使用
pytorch_lightning.trainer.Trainer.fit()时，用户未显式传递的参数（使用默认值的参数）目前不会自动记录 - 在多优化器场景中（例如使用自编码器），仅记录第一个优化器的参数
Scikit-learn
在您的训练代码之前调用 mlflow.sklearn.autolog()，以启用 sklearn 指标、参数和模型的自动日志记录。请参阅此处的示例用法。
estimators（例如 LinearRegression）和 meta estimators（例如 Pipeline）的自动日志记录会创建一个单独的运行并记录
| 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|
通过 estimator.score 获取的训练得分 | 通过 estimator.get_params 获取的参数 |
| 拟合的 estimator |
参数搜索 estimators（例如 GridSearchCV）的自动日志记录会创建一个单独的父运行和嵌套的子运行
- Parent run
- Child run 1
- Child run 2
- ...
包含以下数据
| 运行类型 | 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|---|
| 父级 | 训练得分 |
|
|
|
| 子级 | 每个参数组合的 CV 测试得分 | 每个参数组合 |
| -- |
Spark
使用附加了 mlflow-spark JAR（例如 SparkSession.builder.config("spark.jars.packages", "org.mlflow.mlflow-spark")）的 SparkSession 初始化，然后调用通用的自动日志记录函数 mlflow.spark.autolog()，以在读取时启用 Spark datasource 信息的自动日志记录，无需显式日志语句。请注意，目前尚不支持 Spark ML (MLlib) 模型的自动日志记录。
自动日志记录会捕获以下信息
| 框架 | 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|---|
| Spark | -- | -- | 包含源路径、版本、格式的单个标签。该标签每 datasource 包含一行 | -- |
- 此外，Spark datasource 自动日志记录是异步发生的 - 因此，在启动短时 MLflow 运行导致 datasource 信息未被记录时，可能会（尽管不太可能）出现竞态条件。
对于 Pyspark 3.2.0 或更高版本，Spark datasource 自动日志记录需要将 PYSPARK_PIN_THREAD 环境变量设置为 false。
Statsmodels
在您的训练代码之前调用通用的自动日志记录函数 mlflow.statsmodels.autolog()，以启用指标和参数的自动日志记录。
自动日志记录会捕获以下信息
| 框架 | 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|---|
| Statsmodels | 用户指定的指标 | statsmodels.base.model.Model.fit 参数 | -- | 训练结束时记录 MLflow 模型 (statsmodels.base.wrapper.ResultsWrapper) |
- 每个重写 fit 的模型子类都会预期并记录自己的参数。
XGBoost
在您的训练代码之前调用通用的自动日志记录函数 mlflow.xgboost.autolog()，以启用指标和参数的自动日志记录。
自动日志记录会捕获以下信息
| 框架 | 指标 | 参数 | 标签 | Artifacts |
|---|---|---|---|---|
| XGBoost | 用户指定的指标 | xgboost.train 参数 | -- | 训练结束时记录带有模型签名的 MLflow 模型 (XGBoost 模型)；特征重要性；输入示例 |
如果启用了 early stopping，则最佳迭代的指标将作为额外的步骤/迭代记录。