MLFlow 简介

管理深度学习实验

可以参考 这个问题 下的回答. 主要需要保存每次实验的

  • 代码 (Git 提交记录)
  • 数据 (路径), 模型
  • 超参数, 指标
  • 日志

保证实验结果好找, 便于复现实验. 这类工具 (满足上述部分功能) 有很多, 比如 TensorBoard; 自己写也可以.

介绍 MLflow

MLflow 是开源的机器学习工作流 (workflow) 管理平台, 提供了 Python, R, Java, REST API 等多种接口. 它是 Spark 团队 (他们还创建了 Databricks 公司) 2018 年的新作, 现在已经到 2.1 版本了.

With origins in academia and the open source community, Databricks was founded in 2013 by the original creators of Apache Spark™, Delta Lake and MLflow.

如果只是管理实验, 那每个人在本地上自己操作就行. MLflow 提供了中心化的管理, 有助于多人协作, 管理模型生命周期. 包含如下四个组件 (最主要的是 tracking 和 model registry).

Tracking

两个主要概念, 一个 experiment (比如某数据集上的分类任务) 包含多个 runs, 一个 run 包含 管理深度学习实验 那节提到的各种数据的记录. MLflow 也提供了 UI 查看记录, 可视化操作. 可以作为个人实验用的本地管理工具.

import mlflow

with mlflow.start_run():
    mlflow.log_param(key, value)
    for epoch in range(...):
        ...
        mlflow.log_metric(key, value, step=...)
    mlflow.log_metric(key, value)

MLflow 提供了四种记录: log_param 记录参数, log_metric 记录指标, log_artifact 记录文件, log_model 记录模型.

Projects

打包代码以便其他人 (或自动化工具) 复用. 每个 project 是一个包含代码的文件夹. 包含

  • 项目名
  • 入口 (运行项目的命令行, 参数的解释)
  • 环境 (依赖库)

Models

打包模型. 除了模型文件和环境信息文件, 还包含

  • run_id, 和特定的 run 关联
  • 模型签名 (输入输出格式)
  • 输入例

Model Registry

中心化的模型仓库, 协作管理模型生命周期.

模型用 mlflow.<model_flavor>.log_model() (model_flavor 例如 onnx, pytorch 等) 记录为 MLflow Model 后, 就可以用 Model Registry 注册. 每个注册过的模型包含来源 (哪个 experiment 的哪个 run), 版本, 状态 (staging/production/archived) 等信息. 在同个模型名注册新模型会自动增加版本.

The Staging stage is meant for model testing and validating, while the Production stage is for model versions that have completed the testing or review processes and have been deployed to applications for live scoring. An Archived model version is assumed to be inactive, at which point you can consider deleting it.

Manage model lifecycle | Databricks on AWS

然后可以参考官方文档 UI Workflow.

MLflow 也提供一些 内置的 serve 方法.

MLflow 不提供原生的多用户管理, 现在好像也没有可用的多用户管理插件, 可参考 #724, #3553, #5115.

与其他库的集成

模型 log 的 flavor 支持主流机器学习库.

原生的 mlflow.pytorch.autolog 作用在 pytorch-lightning 上.

其他比如 Huggingface Transformers Trainer 默认开启 MLflowCallback.

其他关于 MLOps 的介绍可以参考 这里 列出的材料.