使用 MLflow 跟踪图像数据集

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

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

为什么追踪图像数据集？

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

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

理解图像数据集格式

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

使用原生文件格式的优点

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

缺点

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

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

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

缺点

• 由于二进制文件以文本形式写入表格条目，文件大小巨大
• 需要自定义数据加载器

COCO: 上下文中的常见对象

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

• 目标检测、关键点检测、物体分割、图像标注等。
• 基于 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

注意：验证集用于节省空间，但如果您想在完整数据集上训练/微调模型，也可以使用“train”集（需要 +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，后者也提供了用于图像数据集管理的丰富用户界面功能。
• 可视化：MLflow 的用户界面和 Databricks 未针对图像数据集标注的可视化进行优化，需要 CVAT 等工具或自定义脚本。
• 中心数据集管理：CVAT 还可以帮助管理和版本化数据集，以便在 MLflow 实验运行中重复使用。

其他资源

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

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

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