MLflow 实验追踪与模型管理全流程
Lisa Tan | 2026-09-01T08:57:31 | Python, AI
介绍 MLflow 的 Tracking、Projects、Models、Registry 四大组件,演示从实验记录到模型上线的完整工作流。
# MLflow 实验追踪与模型管理全流程 ## MLflow 四大组件 - **Tracking**: 记录参数、指标、产物 - **Projects**: 可复现的运行环境 - **Models**: 统一模型打包格式 - **Registry**: 模型版本管理与生命周期 ## 实验追踪 ```python import mlflow import mlflow.sklearn from sklearn.ensemble import RandomForestClassifier from sklearn.model_selection import cross_val_score mlflow.set_tracking_uri("http://localhost:5000") mlflow.set_experiment("customer-churn-prediction") params_grid = [ {"n_estimators": 100, "max_depth": 10, "min_samples_split": 5}, {"n_estimators": 200, "max_depth": 15, "min_samples_split": 3}, {"n_estimators": 300, "max_depth": 20, "min_samples_split": 2}, ] for params in params_grid: with mlflow.start_run(run_name="rf-" + str(params["n_estimators"])): mlflow.log_params(params) model = RandomForestClassifier(**params, random_state=42, n_jobs=-1) scores = cross_val_score(model, X_train, y_train, cv=5, scoring="f1") mlflow.log_metric("f1_mean", scores.mean()) mlflow.log_metric("f1_std", scores.std()) model.fit(X_train, y_train) mlflow.sklearn.log_model(model, "model") # 记录特征重要性图 fig = plot_feature_importance(model, feature_names) mlflow.log_figure(fig, "feature_importance.png") ``` ## 模型注册与阶段管理 ```python from mlflow import MlflowClient client = MlflowClient() # 注册最佳模型 best_run = mlflow.search_runs( order_by=["metrics.f1_mean DESC"], max_results=1 ).iloc[0] model_uri = "runs:/{}/model".format(best_run.run_id) result = mlflow.register_model(model_uri, "ChurnPredictor") # 推进到 Staging client.transition_model_version_stage( name="ChurnPredictor", version=result.version, stage="Staging" ) # 验证通过后推进到 Production client.transition_model_version_stage( name="ChurnPredictor", version=result.version, stage="Production" ) ``` ## 模型加载与推理 ```python # 从 Registry 加载生产版本 model = mlflow.pyfunc.load_model("models:/ChurnPredictor/Production") # 批量预测 predictions = model.predict(X_test) ``` MLflow 支持与 Kubernetes、Seldon Core 等部署平台集成,实现从实验到生产的无缝衔接。推荐在团队中搭建统一的 MLflow Server 来管理所有 ML 实验。