怎样集成MLflow实验跟踪

下面按“从零到可用”的思路，系统讲一下如何集成 MLflow 实验跟踪（Experiment Tracking），涵盖本地、代码、以及和常见框架/平台的集成方式。

一、MLflow 实验跟踪能做什么？

MLflow Tracking 主要记录：

• 参数（params）
• 指标（metrics）
• 模型（artifacts）
• 运行环境（tags、source）

核心概念：

• Experiment：一组实验
• Run：一次训练过程
• Tracking Server：存储实验的地方（本地 / 远程）

二、安装 MLflow

pip install mlflow

三、最简单：本地实验跟踪

1️⃣ 写训练代码

import mlflow
mlflow.set_experiment("my-first-exp")
with mlflow.start_run():
mlflow.log_param("lr", 0.01)
mlflow.log_metric("accuracy", 0.95)
mlflow.log_artifact("model.pkl")

2️⃣ 启动 UI 查看

mlflow ui

浏览器打开：

http://localhost:5000

✅ 适合：个人调试、论文实验

四、远程 Tracking Server（推荐）

1️⃣ 启动 Tracking Server

mlflow server \
--host 0.0.0.0 \
--port 5000 \
--backend-store-uri sqlite:///mlflow.db \
--default-artifact-root s3://my-bucket/mlflow

2️⃣ 客户端连接

import mlflow
mlflow.set_tracking_uri("http://:5000")
mlflow.set_experiment("remote-exp")
with mlflow.start_run():
mlflow.log_param("epochs", 10)
mlflow.log_metric("loss", 0.2)

✅ 适合：团队、多人协作

五、与训练框架集成（自动记录）

✅ Scikit-learn

import mlflow.sklearn
mlflow.sklearn.autolog()
model.fit(X_train, y_train)

✅ PyTorch / Lightning

import mlflow.pytorch
mlflow.pytorch.autolog()

✅ TensorFlow / Keras

import mlflow.tensorflow
mlflow.tensorflow.autolog()

✅ 自动记录：

• 超参
• loss / accuracy
• 模型文件

六、与 ML 平台集成

1️⃣ Databricks（企业级）

mlflow.set_tracking_uri("databricks")

2️⃣ Kubernetes

• Tracking Server 部署为 Service
• Artifact 用 S3 / MinIO

七、常见最佳实践

✅ 每个实验一个 experiment
✅ 用 tags 区分环境（dev / prod）
✅ 模型统一存 artifact

✅ CI 中自动记录实验

八、典型架构示意

训练代码
↓
MLflow Client
↓
Tracking Server
├─ Metadata (SQLite / Postgres)
└─ Artifacts (S3 / HDFS)

如果你愿意，我可以：

• 给你一个 完整 PyTorch + MLflow 示例
• 帮你 设计团队级 MLflow 架构
• 讲 MLflow + Model Registry 模型管理

你现在是个人用还是团队/生产环境？