使用 MLflow 跟踪图像数据集

数据集追踪是构建健壮且可复现的机器学习模型的基础。在所有数据类型中，图像由于其高维度、变异性和存储需求，带来了独特的追踪挑战。在这篇文章中，我们将演示如何有效地使用 MLflow 的实验追踪功能和用户界面追踪图像数据集，为您提供实用的技术，以增强您的计算机视觉工作流的数据和模型追踪。

注意：本指南假定您熟悉 MLflow 及其追踪功能以及 PyTorch。对于初学者，请参阅 MLflow 入门教程 和本 PyTorch 视觉教程。

为什么要追踪图像数据集？

追踪图像数据集对于结构化机器学习项目至关重要。它确保了

• 高质量训练数据：与结构化数据集不同，图像数据集难以理解和阅读，但它们对模型性能的重要性相同，并且仍然需要标注、整理和特征转换。
• 可复现性：可以使用相同的数据集版本和预处理步骤复制实验。
• 数据和模型血缘：追踪维护数据使用记录，这对于遵守数据治理以及追踪训练、验证和测试中使用的模型和数据的血缘至关重要。
• 调试：追踪有助于识别可能影响模型性能的数据质量或预处理相关问题。

了解图像数据集格式

在全球机器学习社区中存在许多数据集格式。在这篇博客文章中，我们将使用计算机视觉模型中广泛使用的 COCO 格式，包括传统文件格式和同一数据集的 Hugging Face 版本。这两种格式各有优缺点，并且对于 MLflow 追踪有不同的可能性。

使用原生文件格式的优点

• dataloaders 可以在 pycocotools / torchvision 中直接使用
• 加载速度快
• 文件大小更小
• 目录结构简单
• 可以追踪和显示工件，例如作为图像

缺点

• 不可追踪
• 非结构化数据，因此搜索和探索可能会很混乱
• 不可查询

使用 Hugging Face 数据集/表格数据集的优点

• 非常结构化
• 作为训练数据集使用 mlflow.data 完全可追踪（详见下文）
• 可以在 MLflow 中为数据添加元数据

缺点

• 文件大小巨大，因为二进制数据以文本形式写入表格条目中
• 需要自定义 dataloader

COCO：上下文中的常见对象

COCO 是计算机视觉中广泛使用的数据集格式，以其丰富的注释而闻名。它支持

• 对象检测、关键点检测、Stuff 分割、图像字幕等等。
• 基于 JSON 的注释，用于存储元数据。

我们将在本博客文章中一直使用此数据集。

为了进行全面的质量检查，可视化图像数据集中的注释非常重要。您可以在 COCO 数据集官方网站 上探索数据集，以了解更多关于其中包含的数据的性质。

图像数据集具有注释，可以是图片中对象的片段或边界框。这意味着对于每张图像，都会有一个类别和一组用于每个识别对象的坐标。请参阅以下来自 COCO 数据集的示例

Hugging Face 图像数据集

Hugging Face 提供了一种简单的 Image Folder 类型，用于从本地文件创建数据集。它支持

• 文本标题和对象检测的元数据集成。
• 使用目录文件路径快速创建数据集。

这可以与多种格式（例如 COCO）一起使用，并受 MLflow 支持。

使用 MLflow 追踪数据集

了解数据集的属性对于有效的 ML 模型训练、测试和评估至关重要。这包括分析类别平衡、标注质量和其他关键特征。通过彻底审查源数据的这些方面，您可以确保数据集与机器学习任务的要求保持一致，并识别可能影响模型性能的潜在偏差或空白。因此，在项目实验和完善过程中，数据集需要与模型一起追踪。

MLflow 提供强大的工具来确保可复现性和模型血缘

• 记录数据集元数据，例如格式（例如，COCO、Hugging Face 图像文件夹）。
• 记录特征转换/数据增强步骤中使用的参数。
• 追踪数据集版本以实现可复现性。
• 存储和检索数据集相关工件，例如数据集描述。
• 将数据集与特定的模型训练运行关联起来。

使我们能够追踪源数据的关键 API 之一是使用 mlflow.log_artifacts 方法，mlflow.log_input 方法，我们将看到如何结合使用 mlflow.data 模块 可以在与 Hugging Face 一起使用时为数据集追踪添加更多结构。我们将使用 mlflow.pytorch 模块文档 来记录模型以及我们的数据集追踪。

使用计算机视觉模型和图像数据集追踪的示例

有两种方法可以记录图像数据集

• 使用 mlflow.artifacts
• 使用 mlflow.data（数据集 API）。

您还可以记录一个 评估数据集，我不会在这篇文章中涉及。

为什么是两种方法？

将图像数据集从基于文件的格式（如 COCO）转换为表格格式具有挑战性，因为大多数数据加载器都期望基于文件的 COCO 格式。记录工件提供了一种快速而直接的解决方案，无需重新格式化文件。但是，如果您不小心组织目录结构中的文件，这也会变得有点混乱。请务必为您的工件创建有意义的路径。

COCO 数据集的关键工件是 instances.json 文件，它描述了您的图像数据集的元数据和注释。例如，通过分析 category 字段，此文件可用于检查数据集中的类别平衡。

如果你不太担心这个问题，Hugging Face 可以帮助以 MLflow 的方式记录数据集。一些 Hugging Face 数据集包含丰富的元数据，可以传输到 MLflow 的追踪功能中。这就是 mlflow.data 发挥作用的地方。与记录工件相比，这为你的数据集添加了更丰富的元数据和结构，使其在给定实验运行中更容易管理和查看。如果你的数据集可以适配 Hugging Face 类型数据集并使其在你的数据加载器或训练脚本中正常工作，那么这是推荐的方法。

在这篇文章中，我将通过代码介绍这两种方法。

安装 MLflow 和其他依赖项

首先在你的 python >= 3.10 环境中安装两个代码示例的依赖项。如果只使用第一个示例，可以省略 opencv；如果只使用第二个示例，可以省略 pycocotools。

pip install mlflow datasets torch torchvision pycocotools opencv-python-headless psutil

如果您想追踪 GPU 指标，请同时安装 pynvml。

对于 Hugging Face 数据集下载，请确保也已登录。

huggingface-cli login

其中一个示例需要计算；因此，请务必开启 MLflow 系统指标 以追踪训练期间在计算上发生的情况。

export MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING=true

注意：验证集用于节省空间，但您也可以使用“训练”集，如果您想在完整数据集上训练/微调模型（需要 +25 GB 存储空间）。训练期间还使用了 Epoch 数量和数据集的子集。

将数据集作为工件与模型一起记录

由于 COCO 数据集是基于文件的，因此需要先下载文件。我们使用官方作者网站上最新版本数据集的最小版本。

# download the COCO val 2017 dataset
wget -P datasets http://images.cocodataset.org/zips/val2017.zip
unzip -q datasets/val2017.zip -d datasets
wget -P datasets http://images.cocodataset.org/annotations/annotations_trainval2017.zip
unzip -q datasets/annotations_trainval2017.zip -d datasets
rm datasets/val2017.zip & rm datasets/annotations_trainval2017.zip

我们现在可以训练模型并追踪训练数据集的工件以及同一运行中的输入。

import json

from torchvision.datasets import CocoDetection
from torchvision import models
import mlflow

# Load a COCO Dataset (val used to limit size)
img_folder = "datasets/val2017"
coco_annotation_file = "datasets/annotations/instances_val2017.json"

# Download dataset
dataset = CocoDetection(img_folder, coco_annotation_file)

# Load a pre-trained model from COCO
model = models.detection.fasterrcnn_resnet50_fpn(weights='COCO_V1')

# Set experiment name
mlflow.set_experiment("coco_experiment")

# Save dataset artifacts and model
with mlflow.start_run():

  # log dataset
  with open(coco_annotation_file, 'r') as f:
      dataset_metadata = json.load(f)
  mlflow.log_dict(dataset_metadata, "coco_annotation_file")

  # log images
  mlflow.log_artifact(img_folder, artifact_path="images")

  # log model
  mlflow.pytorch.log_model(model, "model")

  # register model with a meaningful name
  mlflow.register_model(
      "runs:/{}/model".format(mlflow.active_run().info.run_id),
      "fasterrcnn_resnet50_fpn_coco_2017_model"
  )

我们可以进入 MLflow 用户界面，查看已在模型实验运行下注册的数据集。

支持图像和文本文件的可视化。

数据集与模型一起记录

我们可以使用 Hugging Face 数据集以更结构化的方式进行此操作，并利用一种方便的方式来读取数据。这样，我们将 MLflow 追踪的数据集、训练指标和模型全部放在同一个实验运行中！

import numpy as np
import cv2
import io
import mlflow
from torchvision import models
from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights
import os

os.environ["MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING"] = "true"

# Load the COCO dataset from Hugging Face
dataset = load_dataset("detection-datasets/coco", split="val")

# Transform to MLFlow Dataset
mlflow_dataset = mlflow.data.huggingface_dataset.from_huggingface(dataset)

# For this example we create a subset of the dataset with the first 100 rows
subset_dataset = dataset.select(range(100))

# Load a pre-trained object detection / segmentation model
model = models.detection.fasterrcnn_resnet50_fpn(weights=FasterRCNN_ResNet50_FPN_Weights.DEFAULT)
# Let’s fine-tune it, log dataset, metrics, and model in an MLflow Experiment run

mlflow.set_experiment("hg_image_experiment")

with mlflow.start_run():

  # log training dataset in model training run
  mlflow.log_input(mlflow_dataset, context="training")
  optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=0.0005)

  for epoch in range(1): # We train for 1 epoch in this example

    print(f"Training object detection model, epoch {epoch+1}...")

    for row in subset_dataset: # We run a subset of the dataset to save time

      # In this example we are not using a dataloader but just converting image bytes to ndarray
      image_bytes = io.BytesIO()
      row["image"].save(image_bytes, format="JPEG")
      image_bytes = image_bytes.getvalue()
      if isinstance(image_bytes, bytes):
          image_array = np.frombuffer(image_bytes, np.uint8)
          image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
          image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
      else:
          raise TypeError("Expected bytes object for 'image', got {}".format(type(image_bytes)))
      image = np.array(image)

      # Prepare annotations as target
      annotations = row["objects"]
      target = []
      for i in range(len(annotations['category'])):
          d = {}
          d['boxes'] = torch.tensor(annotations['bbox'][i], dtype=torch.float32).reshape(-1, 4)  # Ensure shape [N, 4]
          d['labels'] = torch.tensor([annotations['category'][i]], dtype=torch.int64)  # Wrap in a list for correct shape
          target.append(d)

      # Convert the image to a PyTorch tensor and normalize it
      image_tensor = torch.tensor(image, dtype=torch.float32).permute(2, 0, 1) / 255.0

      # Perform forward pass in batches of one
      input_batch = [image_tensor]
      output = model(input_batch, target)

      # Compute loss
      loss_dict = output[0] if isinstance(output, list) else output
      loss = sum(loss for loss in loss_dict.values())

      # Backpropagation and optimization step
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

      # Pretty print the loss
      print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

      mlflow.log_metrics({"epoch": epoch+1})
      mlflow.log_metrics({"loss": loss.item()})

    # finally log model
    mlflow.pytorch.log_model(
      model,
      "model",
      input_example=input_batch
    )

我们已经展示了如何以表格格式处理图像，以简化 Hugging Face 数据集在训练运行中的使用。

在第二个实验下，您现在将拥有一个已记录的数据集。

局限性

虽然 MLflow 本身功能强大，但它需要支持。请考虑以下限制

• 存储开销：记录大型数据集可能需要大量存储空间。
• 注释复杂性：管理复杂的注释可能需要自定义脚本，例如 pycocotools 或开源工具，例如 CVAT，它还提供了广泛的 UI 功能，用于图像数据集管理。
• 可视化：MLflow 的 UI 和 Databricks 尚未针对图像数据集注释的可视化进行优化，需要 CVAT 或自定义脚本等工具。
• 中心数据集管理：CVAT 还可以帮助管理和版本化数据集，以便在 MLflow 实验运行中使用。

其他资源

• MLflow 数据集模块使用教程.
• COCO 数据集格式
• Hugging Face 图像数据集
• MLflow 模型数据集评估教程
• 本例中计算机视觉模型的文档

我们希望本指南能帮助您简化 MLflow 的图像数据集追踪，并为您提供一些关于图像数据集的新思路。祝您 ML 模型训练愉快！

永远不要让您的 GPU/CPU 冷却下来。在 MLflow UI 中检查模型训练期间的系统指标。