跳至主内容

MLflow 用于深度学习

深度学习 已经彻底改变了人工智能,使计算机视觉、自然语言处理、生成式 AI 以及无数其他领域实现了突破性能力。随着模型变得越来越复杂,管理深度学习实验的复杂性也变得愈发具有挑战性。

MLflow 提供了一个全面的解决方案,用于在所有主流框架中跟踪、管理和部署深度学习模型。无论你是在微调 Transformer 模型、训练计算机视觉模型,还是开发自定义神经网络,MLflow 强大的工具包都能简化你从实验到生产的工作流程。

Why Deep Learning Needs MLflow

现代深度学习的挑战

  • 🔄 迭代开发: 深度学习需要对模型架构、超参数和训练方案进行大量实验
  • 📊 复杂指标: 模型在训练步骤中生成大量指标,需要被跟踪和比较
  • 💾 大型工件: 模型、检查点和可视化需要系统化的存储和版本控制
  • 🧩 框架多样性: 团队通常使用 PyTorch、TensorFlow、Keras 以及其他专门的库
  • 🔬 可重复性危机: 如果没有适当的跟踪,复现结果几乎不可能
  • 👥 团队协作:多位研究人员需要能够查看实验并在彼此的工作基础上继续开展
  • 🚀 部署复杂性:将成功的实验推向生产会带来新的挑战

MLflow通过一个与框架无关的平台,为整个深度学习生命周期带来结构性和清晰性。

深度学习的关键特性

📊 全面的实验跟踪

MLflow 的跟踪功能专为深度学习的迭代特性量身定制:

  • 一行自动日志记录 用于 PyTorch、TensorFlow 和 Keras
  • 基于步骤的指标 捕捉跨训练周期和批次的训练动态
  • 超参数跟踪 用于架构选择和训练配置
  • 资源监控 跟踪 GPU 利用率、内存消耗和训练时间
Advanced Tracking Capabilities

超越基本指标

MLflow 的跟踪系统支持深度学习工作流的特定需求:

  • 模型架构记录: 自动捕获神经网络结构和参数数量
  • 数据集跟踪: 记录数据集版本、预处理步骤和增强参数
  • 可视化调试: 存储样本预测、注意力图和其他可视化产物
  • 分布式训练:监控分布式训练设置中跨多个节点的指标
  • 自定义工件: 记录混淆矩阵、嵌入投影和其他专用的可视化
  • 硬件性能分析: 跟踪 GPU/TPU 利用率、内存消耗和吞吐量指标
  • 早停点: 记录早停发生的时间并保存最佳模型状态

一目了然地比较训练收敛情况

可视化多个深度学习运行,以快速识别哪些配置在训练迭代中实现了更优的性能。

Training convergence comparison

🏆 简化的模型管理

深度学习模型是需要谨慎管理的宝贵资产:

  • 版本化模型注册表 为您所有的模型提供一个中央存储库
  • 模型溯源 跟踪从数据到部署的完整历史
  • 元数据注释 存储架构细节、训练数据集和性能指标
  • 阶段转换 管理模型贯穿开发、暂存和生产阶段
  • 团队权限 控制谁可以查看、修改和部署模型
  • 依赖管理 确保所有必需的软件包都与模型一起被跟踪
Model Registry for Teams

协作式模型开发

MLflow Model Registry 通过以下方式提高团队的生产力:

  • 阶段迁移请求: 团队成员可以请求模型晋升并附上记录的理由
  • 审批工作流: 在生产部署中实施需要审批的治理(仅限托管 MLflow)
  • 性能基线:在模型能够进入生产之前设置阈值要求
  • 回滚能力:在出现问题时快速还原到以前的版本
  • 活动提要: 跟踪谁在何时对模型进行了更改(仅限托管 MLflow)
  • Webhook 集成: 根据注册表事件触发 CI/CD 流水线和通知(仅限托管 MLflow)
  • 模型文档: 将全面的文档与模型工件一起存储

🚀 简化部署

轻松将成功的实验投入生产:

  • 一致的推理 APIs 在所有深度学习框架中
  • 支持 GPU 的部署,用于计算密集型模型
  • 批处理和实时服务选项,满足不同应用需求
  • Docker 容器化,用于便携、隔离的环境
  • 无服务器部署 用于在您的云服务提供商基础设施内提供可扩展且具有成本效益的服务
  • 边缘部署 支持移动和物联网应用
Advanced Deployment Options

超越基础服务

MLflow 支持复杂的深度学习部署场景:

  • 模型集成: 部署多个模型,使用投票或平均机制
  • 自定义预处理/后处理: 将数据转换管道附加到您的模型
  • Optimized Inference: Support for quantization, pruning, and other optimization techniques
  • 监控集成: 连接到可观测性平台,用于生产环境跟踪
  • 硬件加速: 利用 GPU/TPU 资源在云服务提供商的基础设施中进行高吞吐量推理
  • 可扩展架构: 应对可变负载,具备自动伸缩能力(仅限托管 MLflow)
  • 多框架部署: 在同一服务环境中混合来自不同框架的模型

框架集成

MLflow 为所有主流深度学习框架提供原生支持,使您可以使用首选工具,同时获得统一的实验跟踪和模型管理的好处。

开始使用

Quick Setup Guide

1. 安装 MLflow

pip install mlflow

确保已安装相应的深度学习(DL)集成包。例如,对于支持图像模型的 PyTorch:

pip install torch torchvision

2. 启动跟踪服务器(可选)

# Start a local tracking server
mlflow server --host 0.0.0.0 --port 5000

3. 启用 Autologging

import mlflow

# For TensorFlow/Keras
mlflow.tensorflow.autolog()

# For PyTorch Lightning
mlflow.pytorch.autolog()

# For all supported frameworks
mlflow.autolog()

4. 像平常一样训练你的模型

# Your existing training code works unchanged!
model.fit(train_data, train_labels, epochs=10, validation_data=(val_data, val_labels))

5. 查看结果

打开 MLflow UI 以查看你跟踪的实验:

mlflow ui

或者如果使用跟踪服务器:

http://localhost:5000

现实世界中的应用

使用 MLflow 的深度学习推动了跨行业的广泛应用:

  • 🖼️ 计算机视觉:跟踪目标检测、图像分割和分类模型的性能
  • 🔊 语音识别: 监控声学模型训练并比较不同架构的词错误率
  • 📝 自然语言处理: 管理大型语言模型的微调并评估在下游任务上的性能
  • 🎮 强化学习: 跟踪智能体在各次训练运行中的性能、奖励和与环境的交互
  • 🧬 基因组学: 组织用于分析基因序列和蛋白质结构的深度学习模型
  • 📊 金融预测:比较用于时间序列分析和风险评估的预测模型
  • 🏭 制造业: 部署计算机视觉模型用于质量控制和预测性维护
  • 🏥 医疗保健: 通过严格的版本控制和审批流程管理医学影像模型

高级主题

Distributed Training Integration

MLflow integrates seamlessly with distributed training frameworks:

  • Horovod: 在分布式 TensorFlow 和 PyTorch 训练中跟踪指标
  • PyTorch DDP: 监控分布式数据并行训练
  • TensorFlow 分布式策略: 从多 GPU 和多节点训练中记录指标
  • Ray: 与 Ray 的分布式计算生态系统集成

使用 PyTorch DDP 的示例:

import mlflow
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel

mlflow.pytorch.autolog()

# Initialize process group
dist.init_process_group(backend="nccl")

# Create model and move to GPU with DDP wrapper
model = DistributedDataParallel(model.to(rank))

# MLflow tracking works normally with DDP
with mlflow.start_run():
trainer.fit(model)
Hyperparameter Optimization

MLflow 与流行的超参数优化框架集成:

  • Optuna: 跟踪试验并可视化优化结果
  • Ray Tune: 监控分布式超参数搜索
  • Weights & Biases Sweeps: 将 W&B sweeps 与 MLflow 跟踪同步
  • HyperOpt: 组织并比较超参数搜索结果

使用 Optuna 的示例:

import mlflow
import optuna


def objective(trial):
with mlflow.start_run(nested=True):
# Suggest hyperparameters
lr = trial.suggest_float("lr", 1e-5, 1e-1, log=True)
batch_size = trial.suggest_categorical("batch_size", [16, 32, 64, 128])

# Log parameters to MLflow
mlflow.log_params({"lr": lr, "batch_size": batch_size})

# Train model
model = create_model(lr)
result = train_model(model, batch_size)

# Log results
mlflow.log_metrics({"accuracy": result["accuracy"]})

return result["accuracy"]


# Create study
with mlflow.start_run():
study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=100)

# Log best parameters
mlflow.log_params({f"best_{k}": v for k, v in study.best_params.items()})
mlflow.log_metric("best_accuracy", study.best_value)
Transfer Learning Workflows

MLflow 帮助组织迁移学习和微调工作流:

  • 基础模型注册表: 维护预训练模型的目录
  • 微调跟踪: 在将模型适配到新任务时监控性能
  • 冻结层分析: 比较不同的层冻结策略
  • 学习率调度: 跟踪不同学习率策略对微调的影响

示例:跟踪一次微调运行:

import mlflow
import torch
from transformers import AutoModelForSequenceClassification

with mlflow.start_run():
# Log base model information
base_model_name = "bert-base-uncased"
mlflow.log_param("base_model", base_model_name)

# Create and customize model for fine-tuning
model = AutoModelForSequenceClassification.from_pretrained(base_model_name)

# Log which layers are frozen
frozen_layers = ["embeddings", "encoder.layer.0", "encoder.layer.1"]
mlflow.log_param("frozen_layers", frozen_layers)

# Freeze specified layers
for name, param in model.named_parameters():
if any(layer in name for layer in frozen_layers):
param.requires_grad = False

# Log trainable parameter count
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
total_params = sum(p.numel() for p in model.parameters())
mlflow.log_params(
{
"trainable_params": trainable_params,
"total_params": total_params,
"trainable_percentage": trainable_params / total_params,
}
)

# Fine-tune and track results...

了解更多

在我们针对各框架的指南中深入了解 MLflow 的深度学习功能: