实验追踪与模型版本管理MLflow/DVC/WB实战引言当你训练了100个模型如何快速找到那个在测试集上准确率最高、且使用了learning_rate0.05的版本实验追踪Experiment Tracking就是解决这个问题的核心实践。本文将对比三大主流工具——MLflow、DVC和Weights Biases并通过实战代码展示如何管理模型的全生命周期。核心概念三个层次的管理层次工具解决的问题实验追踪MLflow/WB记录参数、指标、模型文件数据版本DVC跟踪数据集变更像Git管理代码一样管理数据模型注册MLflow Model Registry模型阶段管理Staging→Production→Archived工具对比维度MLflowDVCWB开源✅✅❌免费版有限自托管✅✅远程存储❌数据版本❌✅❌可视化中等弱极强协作中等Git原生强学习曲线低中等低实战代码一、MLflow实验追踪1. 安装与配置pipinstallmlflow2.8.0# 启动MLflow服务器mlflow server--host0.0.0.0--port5000\--backend-store-uri sqlite:///mlflow.db\--default-artifact-root ./mlruns2. 完整的实验追踪示例importmlflowimportmlflow.sklearnfrommlflow.modelsimportinfer_signatureimportpandasaspdimportnumpyasnpfromsklearn.ensembleimportRandomForestClassifier,GradientBoostingClassifierfromsklearn.model_selectionimporttrain_test_split,cross_val_scorefromsklearn.metricsimport(accuracy_score,precision_score,recall_score,f1_score,roc_auc_score,confusion_matrix,classification_report)fromsklearn.datasetsimportmake_classificationimportmatplotlib.pyplotaspltimportseabornassnsimportjson# 设置追踪服务器mlflow.set_tracking_uri(http://localhost:5000)defcreate_experiment_with_tags(name:str,tags:dict)-str:创建带标签的实验experiment_idmlflow.create_experiment(name,tagstags)returnexperiment_iddeflog_confusion_matrix(y_true,y_pred,artifact_path:str):记录混淆矩阵为图片cmconfusion_matrix(y_true,y_pred)plt.figure(figsize(8,6))sns.heatmap(cm,annotTrue,fmtd,cmapBlues)plt.xlabel(Predicted)plt.ylabel(Actual)plt.title(Confusion Matrix)plt.tight_layout()plt.savefig(f{artifact_path}/confusion_matrix.png)plt.close()deftrain_and_track(model_class,model_params:dict,X_train,X_test,y_train,y_test,experiment_name:str,run_name:str):训练模型并记录所有实验信息mlflow.set_experiment(experiment_name)withmlflow.start_run(run_namerun_name)asrun:# 记录参数mlflow.log_params(model_params)# 记录数据集信息mlflow.log_param(train_samples,len(X_train))mlflow.log_param(test_samples,len(X_test))mlflow.log_param(n_features,X_train.shape[1])# 训练模型modelmodel_class(**model_params)model.fit(X_train,y_train)# 预测y_predmodel.predict(X_test)y_probamodel.predict_proba(X_test)[:,1]ifhasattr(model,predict_proba)elseNone# 计算指标metrics{accuracy:accuracy_score(y_test,y_pred),precision:precision_score(y_test,y_pred,averageweighted),recall:recall_score(y_test,y_pred,averageweighted),f1:f1_score(y_test,y_pred,averageweighted),}ify_probaisnotNone:metrics[roc_auc]roc_auc_score(y_test,y_proba)# 交叉验证cv_scorescross_val_score(model,X_train,y_train,cv5,scoringaccuracy)metrics[cv_mean]cv_scores.mean()metrics[cv_std]cv_scores.std()# 记录指标mlflow.log_metrics(metrics)# 记录混淆矩阵log_confusion_matrix(y_test,y_pred,plots)mlflow.log_artifacts(plots)# 记录特征重要性ifhasattr(model,feature_importances_):importancedict(zip([ffeature_{i}foriinrange(X_train.shape[1])],model.feature_importances_.tolist()))mlflow.log_dict(importance,feature_importance.json)# 记录模型带签名signatureinfer_signature(X_train,y_pred)mlflow.sklearn.log_model(model,model,signaturesignature,input_exampleX_train[:5])# 记录分类报告reportclassification_report(y_test,y_pred,output_dictTrue)mlflow.log_dict(report,classification_report.json)print(fRun ID:{run.info.run_id})print(fMetrics:{json.dumps(metrics,indent2)})returnrun.info.run_id,metrics# 主流程if__name____main__:# 生成数据X,ymake_classification(n_samples10000,n_features20,n_informative15,n_redundant3,random_state42)X_train,X_test,y_train,y_testtrain_test_split(X,y,test_size0.2,random_state42)# 实验1随机森林run_id_1,metrics_1train_and_track(RandomForestClassifier,{n_estimators:200,max_depth:10,random_state:42},X_train,X_test,y_train,y_test,model_comparison,random_forest_v1)# 实验2GBDTrun_id_2,metrics_2train_and_track(GradientBoostingClassifier,{n_estimators:200,max_depth:5,learning_rate:0.1},X_train,X_test,y_train,y_test,model_comparison,gbdt_v1)# 对比结果print(\n 模型对比 )print(f随机森林 F1:{metrics_1[f1]:.4f})print(fGBDT F1:{metrics_2[f1]:.4f})3. 模型注册与阶段管理importmlflowfrommlflow.trackingimportMlflowClient clientMlflowClient(http://localhost:5000)# 注册模型model_urifruns:/{run_id_2}/modelmodel_versionmlflow.register_model(model_uri,churn_prediction_model)# 添加描述client.update_model_version(namechurn_prediction_model,versionmodel_version.version,descriptionGBDT模型F10.92用于客户流失预测)# 设置别名MLflow 2.x推荐方式client.set_registered_model_alias(namechurn_prediction_model,aliaschampion,versionmodel_version.version)# 模型阶段转换# 旧版API仍然可用client.transition_model_version_stage(namechurn_prediction_model,versionmodel_version.version,stageProduction)# 加载生产模型production_modelmlflow.pyfunc.load_model(models:/churn_prediction_modelchampion)二、DVC数据版本管理1. 安装与初始化pipinstalldvc3.25.0# 初始化DVCcdyour_projectgitinit dvc initgitadd.dvc .dvcignoregitcommit-mInitialize DVC2. 数据版本控制# 配置远程存储S3示例dvc remoteadd-dstorage s3://my-ml-bucket/dvc-store dvc remote modify storage endpointurl http://minio:9000# 添加数据文件到DVC追踪dvcadddata/training_data.csvgitadddata/training_data.csv.dvc .gitignoregitcommit-mAdd training data v1# 修改数据后重新版本化dvcadddata/training_data.csvgitadddata/training_data.csv.dvcgitcommit-mUpdate training data v2 - add new samples# 推送数据到远程dvc push3. DVC Pipeline定义# dvc.yamlstages:prepare:cmd:python scripts/prepare_data.pydeps:-scripts/prepare_data.py-data/raw/training_data.csvouts:-data/processed/train.csv-data/processed/test.csvparams:-prepare.test_sizetrain:cmd:python scripts/train.pydeps:-scripts/train.py-data/processed/train.csvouts:-models/model.pklparams:-train.n_estimators-train.max_depthmetrics:-metrics/train_metrics.json:cache:falseplots:-plots/confusion_matrix.pngevaluate:cmd:python scripts/evaluate.pydeps:-scripts/evaluate.py-models/model.pkl-data/processed/test.csvmetrics:-metrics/test_metrics.json:cache:false# 运行Pipelinedvc repro# 查看指标对比dvc metrics show dvc metricsdiff# 查看参数和指标的变化dvc plots show dvc plotsdiff4. 数据版本回溯# 查看数据历史gitlog--onelinedata/training_data.csv.dvc# 回溯到旧版本gitcheckout HEAD~1 -- data/training_data.csv.dvc dvc checkout三、WB实验追踪1. 基础使用importwandbimporttorchimporttorch.nnasnnfromtorch.utils.dataimportDataLoader,TensorDataset# 初始化项目wandb.init(projectml-experiment-tracking,namepytorch_mlp_v1,config{learning_rate:0.001,batch_size:64,epochs:50,architecture:MLP,hidden_layers:[256,128,64],optimizer:Adam,dropout:0.3})configwandb.config# 模型定义classMLP(nn.Module):def__init__(self,input_dim,hidden_dims,dropout):super().__init__()layers[]prev_diminput_dimfordiminhidden_dims:layers.extend([nn.Linear(prev_dim,dim),nn.BatchNorm1d(dim),nn.ReLU(),nn.Dropout(dropout)])prev_dimdim layers.append(nn.Linear(prev_dim,1))self.netnn.Sequential(*layers)defforward(self,x):returnself.net(x)modelMLP(20,config.hidden_layers,config.dropout)optimizertorch.optim.Adam(model.parameters(),lrconfig.learning_rate)criterionnn.BCEWithLogitsLoss()# 训练循环带WB记录forepochinrange(config.epochs):model.train()total_loss0forbatch_X,batch_yintrain_loader:optimizer.zero_grad()outputmodel(batch_X).squeeze()losscriterion(output,batch_y)loss.backward()optimizer.step()total_lossloss.item()# 验证model.eval()withtorch.no_grad():val_outputmodel(X_val_tensor).squeeze()val_losscriterion(val_output,y_val_tensor)val_pred(torch.sigmoid(val_output)0.5).float()val_acc(val_predy_val_tensor).float().mean()# 记录到WBwandb.log({train_loss:total_loss/len(train_loader),val_loss:val_loss.item(),val_accuracy:val_acc.item(),epoch:epoch,learning_rate:config.learning_rate})# 记录最终模型wandb.log_model(models/best_model.pth,pytorch_mlp)# 记录混淆矩阵wandb.log({confusion_matrix:wandb.plot.confusion_matrix(y_truey_test.tolist(),predsval_pred.tolist(),class_names[No Churn,Churn])})wandb.finish()2. 超参数搜索Sweeps# sweep_config.yamlprogram:train.pymethod:bayesmetric:name:val_accuracygoal:maximizeparameters:learning_rate:min:0.0001max:0.01batch_size:values:[32,64,128]hidden_layers:values:[[256,128],[512,256,128],[1024,512,256]]dropout:min:0.1max:0.5optimizer:values:[adam,sgd,adamw]importwandb sweep_idwandb.sweep(sweep_config,projecthyperparameter-search)deftrain_sweep():runwandb.init()configwandb.config# 训练代码...wandb.finish()wandb.agent(sweep_id,functiontrain_sweep,count50)进阶技巧1. 版本化Pipeline# 创建实验分支gitcheckout-bexperiment/lr-0.001 dvc exp run --set-paramtrain.learning_rate0.001# 对比实验dvc exp show dvc exp compare2. 模型血缘追踪# 记录数据版本与模型的关联mlflow.set_tag(dvc_data_version,abc123)mlflow.set_tag(git_commit,def456)mlflow.set_tag(training_data_hash,sha256:...)常见问题与避坑问题原因解决方案MLflow UI加载慢Artifact过大压缩模型文件使用模型签名DVC push失败远程存储配置错误dvc remote modify检查凭证WB上传超时网络问题设置wandb.init(settingswandb.Settings(init_timeout300))实验太多难管理无命名规范使用{model}_{version}_{date}命名模型文件冲突多人同时训练使用DVC分支实验命名区分总结实验追踪不是可选项而是ML项目的基础建设。选择适合团队的工具组合小团队快速起步MLflow本地服务器数据密集型项目MLflow DVC需要强大可视化WB企业级全功能MLflow DVC 自建平台