如何集成MLflow实验跟踪

集成 MLflow 实验跟踪（Experiment Tracking） 通常用于记录模型训练过程中的参数、指标、模型、日志和制品。下面按 从零到可落地 的方式说明，适用于 本地 / 单机 / 团队 / 云端 多种场景。

一、MLflow 实验跟踪的核心概念

二、安装 MLflow

pip install mlflow

如果你使用 sklearn / PyTorch / TensorFlow，可一并安装：

pip install mlflow scikit-learn

三、最简单示例（本地实验跟踪）

1️⃣ 启动 MLflow UI

mlflow ui

访问：

http://127.0.0.1:5000

2️⃣ 在训练代码中记录实验

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
# 设置实验
mlflow.set_experiment("iris_rf_experiment")
with mlflow.start_run():
# 参数
params = {
"n_estimators": 100,
"max_depth": 5
}
mlflow.log_params(params)
# 数据
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
# 模型
model = RandomForestClassifier(**params)
model.fit(X_train, y_train)
# 指标
acc = accuracy_score(y_test, model.predict(X_test))
mlflow.log_metric("accuracy", acc)
# 记录模型
mlflow.sklearn.log_model(model, "model")
print(f"Accuracy: {acc}")

四、设置 Tracking URI（关键）

✅ 本地（默认）

mlflow.set_tracking_uri("file:///path/to/mlruns")

✅ 远程服务器（推荐团队使用）

mlflow.set_tracking_uri("http://mlflow-server:5000")

✅ 环境变量方式

export MLFLOW_TRACKING_URI=http://mlflow-server:5000

五、组织实验（最佳实践）

1️⃣ 一个项目一个 Experiment

mlflow.set_experiment("recommendation_system")

2️⃣ 给 Run 起名字

with mlflow.start_run(run_name="rf_100_trees"):
...

3️⃣ 记录更多 Artifact

mlflow.log_artifact("confusion_matrix.png")
mlflow.log_dict({"features": ["a", "b"]}, "features.json")

六、与常见框架集成

✅ PyTorch

import mlflow.pytorch
mlflow.pytorch.log_model(model, "model")

✅ TensorFlow / Keras

import mlflow.tensorflow
mlflow.tensorflow.log_model(model, "model")

✅ AutoML / Hyperparameter Tuning

for lr in [0.01, 0.001]:
with mlflow.start_run():
mlflow.log_param("lr", lr)
...

七、在团队/生产环境中部署 MLflow

推荐架构

训练脚本
↓
MLflow Tracking Server
↓
后端存储（MySQL / PostgreSQL）
↓
制品存储（S3 / MinIO）

启动 Tracking Server

mlflow server \
--backend-store-uri mysql+pymysql://user:pass@host/mlflow \
--default-artifact-root s3://mlflow-bucket/ \
--host 0.0.0.0 \
--port 5000

八、常见问题

❓ Run 没显示？

• 是否 mlflow.start_run()
• Tracking URI 是否一致

❓ 模型太大？

• 使用 S3 / MinIO 存 artifact
• 不要 log 整个数据集

❓ 多人协作？

• 统一 Tracking Server
• 统一 experiment 命名规范

九、总结一句话

MLflow 实验跟踪 = 用 start_run() 包裹训练代码 + 记录参数 / 指标 / 模型 + 统一 Tracking URI
如果你愿意，我可以：

• ✅ 给你一个 完整项目模板
• ✅ 结合 PyTorch / TensorFlow 写示例
• ✅ 帮你设计 团队级 MLflow 架构
• ✅ 对接 Airflow / Kubeflow

你现在的场景是 个人实验、团队使用，还是生产部署？