Skip to content

MLflow

MLflow 是一个开源的机器学习生命周期管理平台，由 Databricks 开发。它提供实验追踪、模型打包、模型注册和部署等功能，特别适合需要自托管的企业环境。

核心组件

MLflow 包含四个主要组件：

组件 功能 说明
MLflow Tracking 实验追踪 记录参数、指标、代码版本、模型文件
MLflow Projects 项目打包 用标准格式描述如何运行 ML 代码
MLflow Models 模型打包 统一的模型格式，支持多种部署方式
Model Registry 模型注册 模型版本管理、审批流程、生命周期管理

基本使用

安装

pip install mlflow

实验追踪

import mlflow

# 设置实验名
mlflow.set_experiment("image-classification")

with mlflow.start_run(run_name="resnet50-baseline"):
    # 记录超参数
    mlflow.log_param("learning_rate", 1e-3)
    mlflow.log_param("batch_size", 32)
    mlflow.log_param("optimizer", "AdamW")

    for epoch in range(num_epochs):
        train_loss = train_one_epoch(model, train_loader, optimizer)
        val_loss, val_acc = evaluate(model, val_loader)

        # 记录指标
        mlflow.log_metric("train_loss", train_loss, step=epoch)
        mlflow.log_metric("val_loss", val_loss, step=epoch)
        mlflow.log_metric("val_accuracy", val_acc, step=epoch)

    # 保存模型
    mlflow.pytorch.log_model(model, "model")

    # 记录文件
    mlflow.log_artifact("config.yaml")

启动 UI

mlflow ui --port 5000
# 访问 http://localhost:5000

模型注册与部署

Model Registry

# 注册模型
result = mlflow.register_model(
    model_uri="runs:/<run_id>/model",
    name="image-classifier"
)

# 模型阶段管理
from mlflow import MlflowClient

client = MlflowClient()
client.transition_model_version_stage(
    name="image-classifier",
    version=1,
    stage="Production"  # None, Staging, Production, Archived
)

模型部署

# 加载已注册的模型
import mlflow.pytorch

model = mlflow.pytorch.load_model("models:/image-classifier/Production")

# 通过 REST API 部署
# mlflow models serve -m "models:/image-classifier/Production" -p 5001

PyTorch 自动记录

MLflow 支持 PyTorch Lightning 的自动记录：

import mlflow

mlflow.pytorch.autolog()

# 之后的 Lightning Trainer 会自动记录所有指标和模型
trainer = pl.Trainer(max_epochs=10)
trainer.fit(model, train_loader, val_loader)

适用场景

  • 企业自托管：不希望数据上传到第三方云服务
  • ML Pipeline：与 Airflow、Kubeflow 等编排工具集成
  • 模型治理：需要模型版本管理、审批流程、阶段转换
  • 多框架支持：团队使用多种 ML 框架（PyTorch、TensorFlow、Sklearn 等）

参考

评论 #