如何通过 DVC 与 MLFlow 协同构建端到端的模型实验管理工作流

在 AI 基础设施建设中，模型开发的一大核心痛点是“实验一致性”。虽然 MLFlow 能够完美地记录超参数和评估指标，但它在处理 TB 级别的原始数据版本时显得力不从心；而 DVC (Data Version Control) 虽然是数据版本管理的利器，却缺乏直观的实验对比 UI。

本文将展示如何将 DVC 与 MLFlow 深度整合，通过将 DVC 的数据哈希值注入 MLFlow 实验元数据，实现“代码-数据-模型”三位一体的追溯体系。

1. 核心架构设计

• DVC: 负责追踪大型数据集、模型权重文件（存放在 S3/OSS 中），并在本地生成 .dvc 轻量化指纹文件。
• MLFlow: 负责实验追踪、参数记录、以及存储指向 DVC 指纹文件的 metadata。

2. 实战步骤

第一步：初始化 DVC 数据版本

首先，我们需要对原始数据进行版本化。

dvc init
# 假设数据存储在 data/dataset.csv
dvc add data/dataset.csv
git add data/dataset.csv.dvc .gitignore
git commit -m \"Initialize training data version\"

第二步：编写集成 DVC 指纹的训练脚本

在 Python 训练代码中，通过读取 .dvc 文件的内容，获取数据的 MD5 哈希值并记录到 MLFlow。

import mlflow
import yaml
import os

def get_dvc_hash(dvc_file_path):
    \"\"\"解析 dvc 文件获取数据哈希值\"\"\"
    if not os.path.exists(dvc_file_path):
        return \"none\"
    with open(dvc_file_path, \"r\") as f:
        data = yaml.safe_load(f)
        return data['outs'][0]['md5']

def train_model():
    # 获取当前实验使用的数据版本
    data_hash = get_dvc_hash(\"data/dataset.csv.dvc\")

    mlflow.set_tracking_uri(\"http://localhost:5000\")
    with mlflow.start_run(run_name=\"dvc_integrated_experiment\"):
        # 记录关键标签：数据哈希
        mlflow.set_tag(\"dvc.data.hash\", data_hash)

        # 模拟模型训练过程
        params = {\"learning_rate\": 0.001, \"batch_size\": 32}
        mlflow.log_params(params)

        # 模拟指标记录
        mlflow.log_metric(\"accuracy\", 0.925)

        # 记录模型文件信息
        # 注意：此处记录的是 .dvc 文件而非模型大文件本身
        mlflow.log_artifact(\"data/dataset.csv.dvc\")

if __name__ == \"__main__\":
    train_model()

第三步：运行与溯源

1. 运行实验：执行 python train.py。
2. UI 查看：打开 MLFlow UI，你会发现在 Tags 栏位中清晰地记录了 dvc.data.hash。
3. 精准溯源：如果需要复现某次实验，只需从 MLFlow 找到该 hash 值，并在本地执行：

# 在 Git 中切换到对应的 commit
git checkout <commit_id>
# 通过 DVC 拉取对应版本的数据
dvc checkout

3. 进阶：自动化模型文件的 DVC 管理

当训练结束后生成了 model.pth，我们同样可以使用 DVC 对其进行版本控制，并将生成的 model.pth.dvc 作为 MLFlow Artifact 上传。这样既保证了 MLFlow 仓库的轻量化，又确保了大型模型权重文件的安全性。

总结

通过将 DVC 产生的数据摘要注入 MLFlow 的实验流，我们成功构建了一个闭环：MLFlow 记录了“实验发生了什么”，而 DVC 确保了“实验是在什么数据上发生的”。这种模式是当前企业级 AI 基础设施实现模型可复现性的标准方案。