MLflow for Deep Learning

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

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

为什么深度学习需要 MLflow

现代深度学习的挑战

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

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

深度学习的关键特性

📊 全面的实验跟踪

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

• PyTorch、TensorFlow 和 Keras 的单行自动日志记录
• 分步指标可捕获跨 epoch 和 batch 的训练动态
• 超参数跟踪，用于模型架构选择和训练配置
• 资源监控可跟踪 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 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)

迁移学习工作流

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 指南：探索 Keras 3.0 的多后端功能与 MLflow 的集成
• 模型注册表：管理模型版本和开发阶段的转换
• MLflow 部署：将深度学习模型部署到生产环境