深度学习读取数据很慢_Detectron2代码学习3 -- 数据加载-CSDN博客

庞子奇：Detectron2 代码学习 1 -- 整体结构

2. Detectron2 代码学习 2 -- 检测模型实现

3. 数据加载与实现 (本篇)

4. (计划中) 模型细节与其它设计部分

3. 数据加载与数据接口

3.1 整体框架

从何开始阅读数据加载部分的代码呢？我们来到之前分析过的engine部分，在defaults.py里面的import可以发现构建数据集的部分是

 from detectron2.data import (
     MetadataCatalog,
     build_detection_test_loader,
     build_detection_train_loader,
 )

因此我们需要到data/中从这几个函数入手分析detectron2是如何处理数据加载的。

build_detection_train_loader和build_detection_test_loader很自然地都在data/build.py中。我们以build_detection_train_loader为例进行分析。

在build函数中，主要分为以下几步：

通过get_detection_dataset_dicts得到数据集的基本信息。通过get...函数的注释可以看到，它最终返回了一个列表，列表中的每一项都是标准的数据集字典格式。我们可以把它理解为一个列表，列表中的每一项都是数据集的一个元素，而每个元素都是用字典格式描述的。
通过DatasetFromList将之前得到的列表转化为pytorch格式的dataset，也就是torch.dataset。在pytorch中开发者提供了统一的dataset父类，通过继承这一统一的父类会对后续数据加载等提供更加方便良好的支持。
数据集提供的数据格式不一定能够满足detectron2的模型对于接口的要求。但是detectron2希望能够方便地支持多种任务、多个数据集上的训练、测试，所以在数据集和模型之间需要“胶水”来完成这个任务，这里的胶水就是mapper和MapDataset，它们将数据集中的数据变换成detectron2能够使用的格式。
为了构建一个dataloader，我们还需要指明从数据集中采样数据遵循的原则。例如最简单的方式就是从所有样本中平均采样。实现这一任务，即指明如何采样的，就是Sampler，我们需要在其中定义采样数据遵循的规范。
最后，dataset提供了数据的来源，sampler提供了采样数据的规则，cfg中包含了其它运行过程的定义，我们通过build_batch_data_loader就创建了一个dataloader。在访问过程中，我们只要有一个迭代器，就可以从dataloader中方便地加载数据了。

至此，我们完成了对于detectron2数据加载框架的分析。从这里可以看到几个深度学习，特别是pytorch中数据加载模块的重要层次：

Dataset：定义数据的来源，其中需要包含如何将原始数据变成模型接受的格式
Sampler：从数据集中采样数据的顺序需要遵循的规则
DataLoader：在前两者的基础上，由Pytorch封装的加载数据的接口。

3.2 `get_detection_dataset_dicts`

在此函数中，最终返回值是dataset_dicts，也就是最终描述dataset的列表。在函数中获得dataset_dicts的语句是dataset_dicts = [DatasetCatalog.get(dataset_name) for dataset_name in dataset_names]，也就是说，针对每一个dataset_name，通过DatasetCatalog获得对应的数据集的数据信息。

3.2.1 `DatasetCatalog`

DatasetCatalog的实现在data/catalog.py中，catalog的中文释义是“目录”，而这个类的作用也正是存储数据集的信息，并且指定了访问数据集的方法。在DatasetCatalog主要起到的就是一个索引的作用，在register函数中，我们将能够“返回数据集信息的函数func”和“数据集的名字name”绑定，这样在get方法中通过数据集的名字就可以调用提取数据的函数了。那么下面的问题就是：提取数据的函数从哪里来呢？在data/datasets中，我们看到了一个文件register_coco.py。其中实现了方法register_coco_instance，我们可以看到它调用了DatasetCatalog.register，将COCO数据集和load_coco_json方法绑定到了一起。因此，我们下一步以COCO数据集为例，分析load_coco_json的实现和它返回的结果。

3.2.2 加载COCO数据集

在load_coco_json中，通过coco标注文件的路径json_file、图片文件夹所在位置image_root，最终构建出了需要访问的coco数据集的信息。在这里，对于深度学习初学者需要提到很重要的一点：在构建数据集的函数中，往往不会去读取图片，因为计算机的内存不足以承载所有的图片。因此普遍的解决办法，是在构建数据集的阶段记录下图片的路径，并在训练过程中OnLine地对图片进行读取。所以，在这个函数中，我们也同样不会对图片进行读取。

在这个函数的实现中大量地使用了COCO数据集本身提供的接口，我们不会关心这些接口的实现，仅仅把它们看作是使用者在自己的数据中能够自行定义的方法函数。在整个函数的思路中主要包含两个部分：

第一个部分是从COCO数据中读取所有的信息
第二个部分是将上面读取的信息变成合乎形式的dataset_dicts

针对第一部分，我们首先在COCO(json_file)中用标注信息初始化了COCO数据集。在这里额外扩展一下，json文件用作储存数据集的标注信息是非常常用而且好用的一种手段。之后，detectron2分别通过coco_api.loadCats和coco_api.loadImgs获得了数据集中的类别(Category)和图片(Images)。更进一步，为了将图片和对应的标注对应起来，anns = [coco_api.imgToAnns[img_id] for img_id in img_ids]按照图片的ID得到了每张图片上的标注。

针对第二部分，其实就是一个非常简单的遍历算法：遍历列表中每张图片和它对应的标注anno，把它们组织成本图片的信息record。record中包含了图片的基本信息，例如图片的长、宽和ID，也包含了图片中每个物体的信息，例如它的类别与2D Bbox的标注。在最后，把所有的Record组织成一个列表就实现了对dataset_dict的构建。

3.3 `DatsetFromList`

在Pytorch中有针对数据集的统一接口，它可以方便地转化成统一的DataLoader接口，方便在训练过程中读取数据。因此，我们需要关心如何将上面拿到的dataset_dict转化成一个pytorch的统一dataset。

DatasetFromList的实现在data/common.py中，它继承了Pytorch对数据集的统一接口torch.utils.data.Dataset，将输入的dataset_dicts转化成可以Pytorch读取的数据集。它的主要组成部分有如下几个：

通过__init__读取数据
通过__getitem__定义读取数据的规范

在读取数据的部分其实非常简单，只需要把之前准备好的dataset_dicts存储到类中即可。

__getitem__是torch.utils.data.Dataset中定义的必须自己实现的函数，它的功能是根据给定的index读取数据集中的数据，因此我们可以把它看作是整个实现的核心。在self._lst[idx]中就体现了它的逻辑：把整个数据集想象成装载着数据的列表，我们只需要根据序号访问列表中的内容即可。

3.4 MapDataset

在上一部分中，我们将dataset_dicts变成了一个pytorch dataset，但是我们仍然没有解决如何读取数据，特别是图片数据的问题。在MapDataset中我们就需要处理从数据集索引到真正数据的过程。一个示例的实现可看data/common.py中的MapDataset，其初始化接受一个Pytorch Dataset，也就是在DatasetFromList中得到的结果，利用新的函数读取数据。这个函数区别于DatasetFromList的核心在于它的映射函数，一个示例实现在data/dataset_mapper.py中的DatasetMapper，它的主要目标就是实现self.map_func中映射的作用。

DatasetMapper的核心在__call__函数，它的主要功能是把类变成一个函数一样可以直接调用。在__call__中依次进行了如下步骤：

image = utils.read_image(dataset_dict["file_name"], format=self.image_format)，通过文件名称读取图片，并由utils.check_image_size进行一些检查工作
对图片进行数据增强，同时把数据增强的信息存储到transforms中
aug_input = T.StandardAugInput(image, sem_seg=sem_seg_gt)
transforms = aug_input.apply_augmentations(self.augmentations)
将图片转化成tensor，作为'image'的Key存储到图片数据中
dataset_dict["image"] = torch.as_tensor(np.ascontiguousarray(image.transpose(2, 0, 1)))
针对之前做过的图片变换transforms，我们可能需要对一些标注进行修改。例如我们如果对图片进行了裁剪(Crop)操作，可能BBox的坐标位置也会相对发生改变。这部分都是通过一些工具函数实现的，例如utils.transform_instance_annotations，不再进行详细分析
通过以上步骤，我们实现了一个读取数据的映射函数。

3.5 `Sampler`

在Pytorch数据加载的过程中，需要指定数据采样的规范。例如最简单的方式就是把这个过程看成是均匀采样，每次随机从所有实例中挑选一个；也有复杂一些的需求，例如需要对不同的类别做数量上的平衡，那么就需要对不同的实例赋予不同的权重。因此，我们需要研究一下detectron2中Sampler的写法，具体研究data/samplers/distributed_sampler.py中的RepeatFactorTrainingSampler。在这个Sampler中，它会对不同的样本赋予不一样的权重，目标就是对抗不同类别之间的不平衡现象。

pytorch的Sampler中，最主要的部分在于__iter__函数。它的作用就是返回所在训练/测试的Batch需要使用的数据的序号。在sampler的实现中，就是从所有序号的列表self._infinite_indices()种得到想要的序号。在这里，我们需要额外解释一下yield的作用。按照我的理解，yield包含对内、对外两部分功能，对外的部分yield的作用相当于return，把结果返回；对内的作用yield相当于一个生成器，它存储了这一次函数执行的上下文，在下次函数被调用的时候，就从上一次结束的位置开始。因此，yield相当于帮助我们实现了一个迭代器的功能。

在看完了__iter__之后，我们还需要进一步查看sampler是如何找到的所有序号的的indices，包括如何实现的RepeatFactor，所以我们下面会研究self._infinite_indices()的实现方式。在查看RepeatFactorTrainingSampler之前，我们先参考一下最普通的TraningSampler的实现。普通TrainingSampler的实现方式就是随机产生了一个固定self._size的列表作为indices的列表，因此实现了对于数据集中对象的随机采样。那么在RepeatFactorTrainingSampler中，它的做法就是按照一定的比例对权重更大的index进行重复，之后随机从这个新的列表中采样，由此实现了按照权重采样的经过。

综上，我们探清楚了Sampler的设计方式和写法。

3.6 小结

在定义了dataset、sampler之后，我们就已经基本完成了对数据加载过程的定义。在build_batch_dataloader中detectron2展示了如何通过调用Pytorch的统一接口实现对dataloader的创建。在pytorch中，dataloader的创建逻辑是接受传入的dataset——数据的来源、sampler——采样的方式，在需要的时候online地获得数据，而获得数据的函数就是dataset的__getitem__方法。在上面的分析中，我们就对于如下的几个层次进行分析，自底向上地分析了如下几个层次：