使用 MLflow 跟踪图像数据集

数据集跟踪对于构建健壮且可重现的机器学习模型至关重要。在所有数据类型中，由于图像的高维度、可变性和存储要求，它们带来了独特的跟踪挑战。在这篇博文中，我们将演示如何使用 MLflow 的实验跟踪功能和 UI 有效地跟踪图像数据集，为您提供实用的技术，以增强计算机视觉工作流程中对数据和模型的跟踪。

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

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

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

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

了解图像数据集格式

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

使用原生文件格式的好处

• pycocotools / torchvision 中可用的 dataloaders
• 加载速度快
• 文件大小较小
• 简单的目录结构
• 可跟踪和展示工件，例如作为图像

缺点

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

使用 Hugging Face 数据集 / 表格数据集的好处

• 高度结构化
• 使用 mlflow.data 作为训练数据集可完全跟踪（下文有更多信息）
• 可以为 MLflow 中的数据添加元数据

缺点

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

COCO：背景中的常见物体

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

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

我们将在整篇博文中使用这个数据集。

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

图像数据集具有注释，这些注释可以是图像中物体的分割或边界框。这意味着对于每张图像，都会有一个类别和每个已识别对象的一组坐标。请看 COCO 数据集的以下示例

Hugging Face 图像数据集

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

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

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

使用 MLflow 跟踪数据集

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

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

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

启用我们跟踪源数据的关键 API 之一是通过使用 mlflow.log_artifacts 方法、mlflow.log_input 方法，我们将看到与 Hugging Face 一起使用时，如何结合使用 mlflow.data 模块可以为数据集跟踪增加更多结构。我们将使用 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 UI，在模型的实验运行下看到注册的数据集。

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

将数据集与模型一起记录

我们可以使用 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 中查看模型训练期间的系统指标。