数据加载器

数据加载器是为模型提供数据的组件。数据加载器通常(但并非必须)从 数据集 中获取原始信息,并将其处理成模型所需的格式。

现有数据加载器的工作原理

Detectron2 包含一个内置的数据加载管道。了解其工作原理很重要,以备您需要编写自定义管道时使用。

Detectron2 提供了两个函数 build_detection_{train,test}_loader,这些函数可以从给定的配置创建默认数据加载器。以下是 build_detection_{train,test}_loader 的工作方式

  1. 它获取已注册数据集的名称(例如,“coco_2017_train”)并加载一个 list[dict],该列表以轻量级格式表示数据集项。这些数据集项尚未准备好供模型使用(例如,图像尚未加载到内存中,随机增强尚未应用,等等)。有关数据集格式和数据集注册的详细信息,请参阅 数据集

  2. 此列表中的每个字典都由一个函数(“映射器”)映射

    • 用户可以通过在 build_detection_{train,test}_loader 中指定 “mapper” 参数来自定义此映射函数。默认映射器是 DatasetMapper

    • 映射器的输出格式可以是任意的,只要它被此数据加载器的使用者(通常是模型)接受即可。默认映射器的输出(在批处理后)遵循在 使用模型 中记录的默认模型输入格式。

    • 映射器的作用是将数据集项的轻量级表示转换为模型可以使用的格式(包括例如读取图像、执行随机数据增强和转换为 torch 张量)。如果您想对数据执行自定义转换,则通常需要自定义映射器。

  3. 映射器的输出被批处理(简单地放入一个列表中)。

  4. 此批处理数据是数据加载器的输出。通常,它也是 model.forward() 的输入。

编写自定义数据加载器

build_detection_{train,test}_loader(mapper=) 使用不同的 “mapper” 可以满足大多数自定义数据加载用例。例如,如果您想将所有图像调整为固定大小以进行训练,请使用

import detectron2.data.transforms as T
from detectron2.data import DatasetMapper   # the default mapper
dataloader = build_detection_train_loader(cfg,
   mapper=DatasetMapper(cfg, is_train=True, augmentations=[
      T.Resize((800, 800))
   ]))
# use this dataloader instead of the default

如果默认 DatasetMapper 的参数无法满足您的需求,您可以编写自定义映射器函数并使用它,例如

from detectron2.data import detection_utils as utils
 # Show how to implement a minimal mapper, similar to the default DatasetMapper
def mapper(dataset_dict):
    dataset_dict = copy.deepcopy(dataset_dict)  # it will be modified by code below
    # can use other ways to read image
    image = utils.read_image(dataset_dict["file_name"], format="BGR")
    # See "Data Augmentation" tutorial for details usage
    auginput = T.AugInput(image)
    transform = T.Resize((800, 800))(auginput)
    image = torch.from_numpy(auginput.image.transpose(2, 0, 1))
    annos = [
        utils.transform_instance_annotations(annotation, [transform], image.shape[1:])
        for annotation in dataset_dict.pop("annotations")
    ]
    return {
       # create the format that the model expects
       "image": image,
       "instances": utils.annotations_to_instances(annos, image.shape[1:])
    }
dataloader = build_detection_train_loader(cfg, mapper=mapper)

如果您不仅想更改映射器(例如,为了实现不同的采样或批处理逻辑),build_detection_train_loader 无法满足您的需求,您将需要编写不同的数据加载器。数据加载器只是一个简单的 Python 迭代器,它会生成 模型接受的格式。您可以使用任何您喜欢的工具来实现它。

无论实现什么,都建议查看 detectron2.data 的 API 文档 以了解有关这些函数的 API 的更多信息。

使用自定义数据加载器

如果您使用 DefaultTrainer,您可以覆盖其 build_{train,test}_loader 方法以使用您自己的数据加载器。请查看 deeplab 数据加载器 获取示例。

如果您编写自己的训练循环,您可以轻松地插入您的数据加载器。