用于深度学习的 MLflow

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

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

为什么深度学习需要 MLflow

现代深度学习的挑战

• 🔄 迭代开发：深度学习需要对架构、超参数和训练方案进行广泛的实验
• 📊 复杂指标：模型在训练步骤中生成大量指标，必须对其进行跟踪和比较
• 💾 大型工件：模型、检查点和可视化需要系统存储和版本控制
• 🧩 框架多样性：团队经常使用 PyTorch、TensorFlow、Keras 和其他专用库
• 🔬 可重现性危机：如果没有适当的跟踪，重现结果几乎不可能
• 👥 团队协作：多位研究人员需要了解实验情况，并能够在彼此的工作基础上进行构建
• 🚀 部署复杂性：从成功的实验到生产引入了新的挑战

MLflow 通过一个与框架无关的平台解决了这些挑战，该平台为整个深度学习生命周期带来了结构和清晰性。

深度学习的关键特性

📊 全面的实验跟踪

MLflow 的跟踪功能专为深度学习的迭代性质而设计

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

高级跟踪功能

超越基本指标

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

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

图表比较

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

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

图表自定义

自定义可视化以获得更深入的见解

定制图表以专注于关键指标和训练阶段，帮助您准确地找出某些模型何时以及为什么优于其他模型。

运行比较

分析参数关系

通过 MLflow 的综合比较视图，探索参数交互及其对模型性能的影响。

统计评估

超参数的统计见解

使用箱线图可视化快速确定哪些超参数值始终能带来更好的性能。

实时跟踪

实时监控训练

通过实时更新的指标观看您的深度学习模型训练，无需手动检查进度。

模型比较

模型比较

使用 MLflow UI 跟踪您的所有跨 epoch 的 DL 检查点。比较性能并根据任何指标快速找到最佳检查点。

🏆 简化的模型管理

深度学习模型是宝贵的资产，需要仔细管理

• 版本控制模型注册表为您的所有模型提供了一个中心存储库
• 模型沿袭跟踪从数据到部署的完整历史记录
• 元数据注释存储架构详细信息、训练数据集和性能指标
• 阶段转换管理模型通过开发、暂存和生产阶段
• 团队权限控制谁可以查看、修改和部署模型
• 依赖项管理确保所有必需的包都与模型一起跟踪

团队模型注册表

协作模型开发

MLflow 模型注册表通过以下方式提高团队生产力

• 转换请求：团队成员可以请求模型升级，并提供书面理由
• 审批工作流程：实施治理，需要对生产部署进行审批（仅限托管 MLflow）
• 性能基准：设置模型可以提升到生产环境之前的阈值要求
• 回滚功能：如果出现问题，可以快速恢复到以前的版本
• 活动源：跟踪谁在何时更改了模型（仅限托管 MLflow）
• Webhook 集成：根据注册表事件触发 CI/CD 管道和通知（仅限托管 MLflow）
• 模型文档：将综合文档与模型工件一起存储

🚀 简化的部署

轻松地从成功的实验转移到生产

• 跨所有深度学习框架的一致推理 API
• 适用于计算密集型模型的 GPU 就绪部署
• 适用于不同应用程序需求的批量和实时服务选项
• 适用于可移植、隔离环境的 Docker 容器化
• 适用于云提供商基础设施中可扩展、经济高效的服务的无服务器部署
• 适用于移动和 IoT 应用程序的边缘部署支持

高级部署选项

超越基本服务

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

• 模型集成：部署具有投票或平均机制的多个模型
• 自定义预处理/后处理：将数据转换管道附加到您的模型
• 优化推理：支持量化、剪枝和其他优化技术
• 监控集成：连接到可观察性平台以进行生产跟踪
• 硬件加速：利用 GPU/TPU 资源在云提供商基础设施中实现高吞吐量推理
• 可扩展架构：通过自动缩放功能处理可变负载（仅限托管 MLflow）
• 多框架部署：在同一服务环境中混合来自不同框架的模型

框架集成

MLflow 为所有主要深度学习框架提供本机支持，允许您使用您喜欢的工具，同时获得统一实验跟踪和模型管理的好处。

使用一键自动日志记录无缝跟踪 TensorFlow 实验。在集中式存储库中捕获训练指标、模型架构和 TensorBoard 可视化。

将 MLflow 与 PyTorch 灵活的深度学习生态系统集成。从自定义训练循环记录指标，保存模型检查点，并简化生产部署。

利用 Keras 3.0 的多后端功能和全面的 MLflow 跟踪。通过一致的实验管理跨 TensorFlow、PyTorch 和 JAX 后端监控训练。

在 spaCy NLP 模型的整个生命周期中对其进行跟踪和管理。记录训练指标，比较模型版本，并将语言处理管道部署到生产环境。

开始使用

快速设置指南

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. 启用自动日志记录

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

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

https://:5000

实际应用

与 MLflow 的深度学习为各行业的广泛应用提供支持

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

高级主题

分布式训练集成

MLflow 与分布式训练框架无缝集成

• 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)

超参数优化

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

• Optuna：跟踪试验并可视化优化结果
• Ray Tune：监控分布式超参数扫描
• Weights & Biases Sweeps：将 W&B 扫描与 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)

迁移学习工作流程

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 的深度学习功能

• TensorFlow 指南：掌握 MLflow 与 TensorFlow 和 Keras 的集成
• PyTorch 指南：了解如何跟踪自定义 PyTorch 训练循环
• Keras 指南：通过 MLflow 探索 Keras 3.0 的多后端功能
• 模型注册表：管理模型版本和通过开发阶段的转换
• MLflow 部署：将深度学习模型部署到生产环境