AI模型版本管理实践笔记
AI模型版本管理我没按教科书顺序做。
先解决眼前的阻塞,再回头补原理。
一开始的问题:模型文件到底放哪儿?
刚开始做机器学习项目时,我和很多新手一样,把模型文件直接塞进 Git 仓库:
project/
├── models/
│ ├── model_v1.pth
│ ├── model_v2.pth
│ └── model_final.pth
├── src/
└── .git/
看起来挺正常的,对吧?但很快问题就来了:
- 文件太大:一个 BERT-base 模型就 400MB,加上训练过程中的 checkpoint,仓库很快就破了 1GB
- push 吃力:每次上传模型都要等半天,网络稍微抖一下就得重来
- 协作困难:同事 clone 项目得等半小时,项目换电脑简直是灾难
- 版本混乱:
model_final.pth是哪个版本的代码训练出来的?model_v2.pth对应什么超参数?
最惨的一次是项目 repo 爆到 8GB,整个团队 clone 下来都要 20 分钟。有个新同事搞了半天才发现本地磁盘不够,直接放弃了入职第一个任务。
尝试方案一:模型文件不进 Git,手动管理
意识到问题后,我把模型文件从 Git 里删了,改用外部存储:
project/
├── src/
└── .gitignore
└── models/
然后在共享服务器上建了个目录,大家约定:
/shared/models/
├── bert-base-v1-20250101.pth
├── bert-base-v2-20250105.pth
└── bert-base-final.pth
结果呢?新的问题又来了:
- 没有元数据:
bert-base-v1-20250101.pth是什么版本训练出来的?超参数是什么? - 缺少追溯:线上出问题时,找不到对应的代码版本
- 依赖人工:每次训练完要手动上传,文件名全靠约定,很容易搞错
- 没有历史:覆盖了就是覆盖了,想回到两周前的模型发现已经没了
有个线上服务出了问题,想回退到上一版模型。结果发现"上一版"到底是谁上传的、用啥代码训练的、怎么加载的,全靠回忆。团队几个人对着文件名争论了半小时,最后只能重新训练一个模型顶上。
尝试方案二:用 DVC 管理大文件
这时候了解到了 DVC(Data Version Control),它是专门为机器学习项目设计的版本控制工具,能和 Git 配合管理大文件。
DVC 基本配置
先安装 DVC:
pip install dvc
初始化项目:
cd /path/to/project
dvc init
git add .dvc .dvcignore
git commit -m "Initialize DVC"
配置远程存储(我用的是 S3,也可以用本地目录或其他存储):
dvc remote add -d myremote s3://my-bucket/ml-models
dvc remote modify myremote access_key_id YOUR_ACCESS_KEY
dvc remote modify myremote secret_access_key YOUR_SECRET_KEY
添加模型文件到 DVC
把模型文件加入 DVC 管理:
dvc add models/bert-base-v1.pth
git add models/bert-base-v1.pth.dvc .gitignore
git commit -m "Add BERT model v1"
DVC 会生成一个 .dvc 文件:
outs:
- md5: a3f5e8b2c9d1f4a7e8b3c6d5f2a1e8b4
path: models/bert-base-v1.pth
size: 419430400
这个 .dvc 文件很小,可以放进 Git,而实际的大文件会存到远程存储。
追踪模型训练流程
DVC 的真正威力在于能追踪整个训练流程:
# dvc.yaml
stages:
train:
cmd: python src/train.py --config config/bert_v1.yaml
deps:
- src/train.py
- data/train.csv
- config/bert_v1.yaml
params:
- train.epochs
- train.learning_rate
- train.batch_size
outs:
- models/bert-base-v1.pth:
cache: true
metrics:
- metrics/train_metrics.json:
cache: false
运行训练:
dvc repro
DVC 会自动:
- 检查依赖是否有变化
- 只有变化时才重新训练
- 把输出模型文件注册到缓存
- 记录版本信息和依赖关系
踩过的坑
用 DVC 也不是一帆风顺,踩过的坑包括:
坑一:缓存目录爆炸
一开始不知道 .dvc/cache 会越积越大,磁盘直接爆满:
# 清理缓存
dvc cache dir
cd .dvc/cache
dvc gc
坑二:多机协作时的缓存同步
团队多人协作时,各自的 .dvc/cache 不同步,导致同一个文件要反复下载:
# 配置共享缓存
dvc cache dir /shared/dvc-cache
chmod -R 777 /shared/dvc-cache
坑三:远程存储权限问题
S3 权限配置没做好,导致某些用户无法 push/pull:
# 检查远程配置
dvc remote list
dvc remote status myremote
# 测试连接
dvc remote verify myremote
尝试方案三:用 MLflow 管理模型元数据和版本
DVC 解决了大文件存储和依赖追踪,但模型版本管理的另一块是元数据管理——模型对应的超参数、性能指标、训练配置等。这时候 MLflow 就派上用场了。
MLflow 基本配置
安装 MLflow:
pip install mlflow
启动 MLflow Tracking Server:
mlflow server --backend-store-uri sqlite:///mlflow.db \
--default-artifact-root ./mlruns \
--host 0.0.0.0 \
--port 5000
在训练脚本中集成 MLflow
import mlflow
import mlflow.pytorch
from torch import nn, optim
# 设置实验
mlflow.set_experiment("bert-text-classification")
with mlflow.start_run():
# 记录参数
mlflow.log_params({
"epochs": 10,
"learning_rate": 2e-5,
"batch_size": 32,
"model_name": "bert-base-uncased"
})
# 训练代码...
model = train_model()
# 记录指标
mlflow.log_metrics({
"train_loss": 0.234,
"train_accuracy": 0.923,
"val_loss": 0.289,
"val_accuracy": 0.901
})
# 记录模型
mlflow.pytorch.log_model(model, "model")
# 记录其他文件(配置、图表等)
mlflow.log_artifact("config/bert_v1.yaml")
mlflow.log_artifact("metrics/training_history.png")
# 打印 run ID
print(f"Run ID: {mlflow.active_run().info.run_id}")
模型版本管理
MLflow 的 Models Registry 提供了模型版本管理功能:
from mlflow.tracking import MlflowClient
client = MlflowClient()
# 注册模型
model_uri = f"runs:/{run_id}/model"
model_version = mlflow.register_model(model_uri, "bert-classifier")
# 标记为生产版本
client.transition_model_version_stage(
name="bert-classifier",
version=model_version.version,
stage="Production"
)
# 获取生产版本
production_model = mlflow.pyfunc.load_model(
model_uri="models:/bert-classifier/Production"
)
踩过的坑
坑一:artifact 存储路径配置错误
一开始 artifact 存在本地,换了机器就找不到了:
# 正确配置 artifact URI
mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_artifact_location("s3://my-bucket/mlflow-artifacts")
坑二:并发写入冲突
多个训练任务同时运行时,MLflow 有时会报并发写入冲突:
# 加上重试机制
from tenacity import retry, stop_after_attempt, wait_exponential
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10))
def log_metrics_safely(metrics):
mlflow.log_metrics(metrics)
坑三:模型加载版本不匹配
训练代码和推理代码的库版本不一致,导致加载失败:
# 记录环境信息
import json
import pip
environment = {
"python": platform.python_version(),
"packages": [str(pkg) for pkg in pip.get_installed_distributions()]
}
mlflow.log_dict(environment, "environment.json")
完整方案:DVC + MLflow + Git 的组合
经过多次折腾,最后形成了这套组合拳:
架构总览
目录结构
project/
├── .dvc/
│ └── config # DVC 配置
├── .git/
├── data/
│ └── train.csv # 数据集(DVC 管理)
├── src/
│ ├── train.py # 训练脚本
│ └── inference.py # 推理脚本
├── models/
│ └── .gitignore # 忽略模型文件
├── config/
│ └── bert_v1.yaml # 训练配置
├── dvc.yaml # DVC 流程定义
├── params.yaml # 超参数定义
└── requirements.txt
实际工作流程
训练新模型:
# 1. 更新参数
vim params.yaml
# 2. 运行训练(DVC + MLflow)
dvc repro
# 3. 检查结果
mlflow ui # 打开 http://localhost:5000 查看
# 4. 如果满意,注册模型
python scripts/register_model.py --run-id <run_id>
部署模型:
# 1. 拉取模型文件
dvc pull models/bert-base-v1.pth
# 2. 从 MLflow 加载元数据
python scripts/get_model_metadata.py --model-version 1
# 3. 部署服务
docker-compose up
版本回退:
# 1. 找到目标版本
git log --oneline | grep "bert v1"
# 2. 检出对应代码
git checkout <commit_hash>
# 3. 拉取对应模型
dvc pull
# 4. 从 MLflow 获取对应 run 的元数据
mlflow runs get <run_id>
一些思考和建议
哪些东西应该版本管理?
必须管:
- 模型权重文件(
.pth、.h5、.pkl) - 训练配置和超参数
- 数据集版本(至少是数据集的版本标识)
- 依赖环境(
requirements.txt、environment.yaml)
可以管:
- 训练日志和指标
- 验证集和测试集的划分信息
- 模型评估报告
没必要管:
- 训练过程中的临时 checkpoint(除非很贵)
- 原始数据(除非数据很小且经常变)
- 重复的实验结果(保留最好的几个即可)
什么时候该用 DVC,什么时候用 MLflow?
用 DVC:
- 文件很大(>100MB)
- 需要追踪依赖关系和复现训练流程
- 团队协作,需要共享大文件
用 MLflow:
- 需要记录和比较实验指标
- 需要模型注册和版本管理
- 需要可视化实验结果
两者都用:
- 大型生产项目,既需要管理大文件,又需要完整的实验追踪
一些血泪教训
- 不要等到出问题了才建版本管理:项目刚开始就搭建好,成本低很多
- 命名要规范:
model_final.pth是最糟糕的名字,model_bert_v3_20250115_acc_0.923.pth好很多 - 定期清理:保留最新的 3-5 个版本即可,其他的可以归档或删除
- 文档要跟上:模型文件旁边一定要有对应的
README.md或元数据文件 - 自动化一切:手动记录总会出错,尽可能用脚本自动记录
结语
模型版本管理这事,没有银弹。小项目可能简单的外部存储就够了,大项目可能需要完整的 MLOps 平台。关键是要意识到这个问题的重要性,并且根据项目规模选择合适的方案。
我现在回想起那个把模型文件误删的晚上,虽然很惨,但那次教训确实让我在后续项目中少走了很多弯路。技术选择没有绝对的对错,但"有意识的选择"和"无意识的积累"之间,往往就是这些看似不起眼的坑决定的。
希望这些踩过的坑能帮你在模型版本管理这条路上少掉进几个坑里。当然,新的坑肯定还在前面等着——毕竟这就是技术人员的日常。
版权声明: 本文首发于 指尖魔法屋-AI模型版本管理实践笔记(https://blog.thinkmoon.cn/post/190-ml-model-versioning-practice-dvc-mlflow-git/) 转载或引用必须申明原指尖魔法屋来源及源地址!
评论
使用 GitHub 账号登录后即可留言,支持 Markdown。