[杂记]mmdetection3.x中的数据流与基本流程详解(数据集读取, 数据增强, 训练)

---

之前跑了一下mmdetection 3.x自带的一些算法, 但是具体的代码细节总是看了就忘, 所以想做一些笔记, 方便初学者参考. 其实比较不能忍的是, 官网的文档还是空的…

这次想写其中的数据流是如何运作的, 包括从读取数据集的样本与真值, 到数据增强, 再到模型的forward当中.

---

0. MMDetection整体组成部分

让我们首先回顾一下C++的标准模板库(STL)是怎样设计的. STL的三个核心组件是容器, 算法与迭代器. 容器, 例如vector, queue等等, 他们是负责存储数据的, 算法是负责进行一些操作, 例如排序, 查找等等. 而迭代器是容器与算法之间的桥梁, 也就是算法可以通过迭代器去访问容器, 使得算法可以独立于容器的类型进行操作. 三个部分相辅相成, 就达到了泛型编程的理念.

再让我们回顾一下一套深度学习的代码包含什么部分. 从大的方面来说, 需要有数据的读取与增强(DataLoader), 模型的定义, 损失函数的计算, 负责梯度传播的优化器, 在验证(测试)集上的评估等. 同理, MMDetection也是按照这种方式来的, 并且每个部分接口相通, 就可以实现更广义的模型定义和训练方式.

在mmengine/registry/__init__.py中, 我们可以看到, MMEngine(或者说MMDetection)总体有这些类型的模块:

from .root import (DATA_SAMPLERS, DATASETS, EVALUATOR, FUNCTIONS, HOOKS,
                   INFERENCERS, LOG_PROCESSORS, LOOPS, METRICS, MODEL_WRAPPERS,
                   MODELS, OPTIM_WRAPPER_CONSTRUCTORS, OPTIM_WRAPPERS,
                   OPTIMIZERS, PARAM_SCHEDULERS, RUNNER_CONSTRUCTORS, RUNNERS,
                   TASK_UTILS, TRANSFORMS, VISBACKENDS, VISUALIZERS,
                   WEIGHT_INITIALIZERS)

那么以上这么多模块可以分成几类, 分别负责什么呢? 按照我个人的理解, MMDetection的整体组成部分可以表示为下图:

为了节省空间, 优化器相关并未画出

1. 认识config文件

mmdetection设计的核心思想是通过字典来配置整个的训练过程和模型定义, 这些字典放在一个.py的config文件中. 一般来说，config文件最重要的就是数据加载（train_dataloader, val_dataloader和test_dataloader）, 模型定义(model)和训练与测试过程(train_pipeline, test_pipeline). 除此之外, 还有一些训练, 测试配置(train_cfg, test_cfg)等等. 具体config的例子可以参照官网Learn about configs.

需要注意的是, mmdetection中字典定义class的方式, 往往是键type表示类的名字, 之后的其他键都是类初始化需要的参数. 例如, 如果我想自定义一个模型, 叫做MyModel, 定义在当前目录下的./models/my_model.py中, 定义方式如下:

from mmdet.registry import MODELS  # 自定义模型, 需要在模型库中"注册", 初始化时才能找到定义
from mmdet.models.mot.base import BaseMOTModel  # 一个模型基类

@MODELS.register_module()  # 装饰器 在模型库中"注册"
class MyModel(BaseMOTModel):
    def __init__(self, 
                 arg1=..., 
                 arg2=..., 
                 arg3=...):
        ...
	def loss(self, inputs, data_samples):  # 前向传播, inputs是输入tensor, data_samples是包含标签的列表
		...

如果按上述方式定义了模型, 那么在我们的配置文件中, 就是这个样子:

# 必须将自定义类的py文件导入 这样可以自动register自定义模型 否则模型初始化时找不到

custom_imports = dict(
    imports=['models.my_model'],
    allow_failed_imports=False)

# 现在就可以愉快的传参了
models=dict(
    type='MyModel', 
    arg1=1, 
    arg2=[16, 128], 
    arg3=dict(channel=256), 
    ...
)

同样, 我们可以自定义DataLoader, Loss, 等等.

此外, dict是可以嵌套的, 例如mmdetection将检测模型分成了backbone, neck和head三部分, 那么如果我们又自定义了一个Head, 叫MyHead:

from mmdet.registry import MODELS  # 自定义模型, 需要在模型库中"注册", 初始化时才能找到定义
from mmengine.model import BaseModule  # 一个模型基类

@MODELS.register_module()  # 装饰器 在模型库中"注册"
class MyHead(BaseModule):
    def __init__(self, 
                 arg4=...):
        ...

这样, 如果MyModel的前向传播过程中需要一个head, 则代码大致是这个样子:

from mmdet.registry import MODELS  # 自定义模型, 需要在模型库中"注册", 初始化时才能找到定义
from mmdet.models.mot.base import BaseMOTModel  # 一个模型基类

@MODELS.register_module()  # 装饰器 在模型库中"注册"
class MyModel(BaseMOTModel):
    def __init__(self, 
                 arg1=..., 
                 arg2=..., 
                 arg3=...,
                 head=...):
        self.head = MODELS.build(head)  # 建立Head的模型, 类型是nn.Module
        ...
	def loss(self, inputs, data_samples):  # 前向传播, inputs是输入tensor, data_samples是包含标签的列表
		...  # 一些其他过程
		ret = self.head(inputs)  # forward
		...  # 后处理

配置文件中对应更改为:

如果按上述方式定义了模型, 那么在我们的配置文件中, 就是这个样子:

custom_imports = dict(
    imports=['models.my_model', '自定义HEAD所在的py文件'],
    allow_failed_imports=False)

models=dict(
    type='MyModel', 
    arg1=1, 
    arg2=[16, 128], 
    arg3=dict(channel=256), 
    head=dict(  # 定义head
        type='MyHead',
        arg4=256,
        ...
	)
    ...
)

篇幅所限, 自定义损失函数, 数据增强之类的就不一一列举了.

2. 数据流

我们接下来以检测与跟踪任务为例, 看看数据到底是如何被读入的. 我们以训练过程说明.

在训练过程中, 我们会初始化一个RUNNER类, 其读入我们的config文件并依次完成各种(模型, 数据加载, 优化器, 钩子等等)的初始化. 我们以官方提供的train.py为例:

runner = Runner.from_cfg(cfg)

from_cfg()是一个类方法(classmethod), 在其中我们实例化了Runner类.

随后, 我们调用Runner的train()方法进行训练. 首先, 我们实例化训练循环:

        self._train_loop = self.build_train_loop(
            self._train_loop)  # type: ignore

训练循环就属于LOOP类型.

在这里, 我们以最常用的EpochBasedTrainLoop为例. 在EpochBasedTrainLoop的初始化函数中, 根据config文件中的train_dataloader字典实例化出torch的DataLoader类():

        data_loader = DataLoader(
            dataset=dataset,
            sampler=sampler if batch_sampler is None else None,
            batch_sampler=batch_sampler,
            collate_fn=collate_fn,
            worker_init_fn=init_fn,
            **dataloader_cfg)
        return data_loader

当然, 我们知道torch的DataLoader类在调用的时候, 会调用到dataset(类别是torch.utils.data.Dataset)的__getitem__方法. 因此, 我们从__getitem__入手来探索数据流.

在MMDetection的设计中, 数据集的类都是继承于MMengine中的BaseDataset, 其中的__getitem__是这样写的:

    def __getitem__(self, idx: int) -> dict:
        if not self._fully_initialized:
            print_log(
                'Please call `full_init()` method manually to accelerate '
                'the speed.',
                logger='current',
                level=logging.WARNING)
            self.full_init()

        if self.test_mode:
            data = self.prepare_data(idx)
            if data is None:
                raise Exception('Test time pipline should not get `None` '
                                'data_sample')
            return data

        for _ in range(self.max_refetch + 1):
            data = self.prepare_data(idx)
            # Broken images or random augmentations may cause the returned data
            # to be None
            if data is None:
                idx = self._rand_another()
                continue
            return data

        raise Exception(f'Cannot find valid image after {self.max_refetch}! '
                        'Please check your image path and pipeline')

我们可以看到, 在__getitem__中最核心的是self.prepare_data(idx). 按照这种思路一级一级向上查找, 我们就可以总结出如下图的数据读取流程:

其中, 数据增强pipeline是一系列类型为TRANSFORMS类的列表, 再每经过一次数据增强时, 字典都会被更新.

我们以较为常用的随机便宜(RandomShift)来说, 其是这样定义的:

@TRANSFORMS.register_module()
class RandomShift(BaseTransform):

    def __init__(self,
        ...

    @autocast_box_type()
    def transform(self, results: dict) -> dict:  # transform方法, 更新字典, 图像与对应的边界框等都需要被更新
        """Transform function to random shift images, bounding boxes.

        Args:
            results (dict): Result dict from loading pipeline.

        Returns:
            dict: Shift results.
        """
        if self._random_prob() < self.prob:
            img_shape = results['img'].shape[:2]

            random_shift_x = random.randint(-self.max_shift_px,
                                            self.max_shift_px)
            random_shift_y = random.randint(-self.max_shift_px,
                                            self.max_shift_px)
            new_x = max(0, random_shift_x)
            ori_x = max(0, -random_shift_x)
            new_y = max(0, random_shift_y)
            ori_y = max(0, -random_shift_y)

            # TODO: support mask and semantic segmentation maps.
            bboxes = results['gt_bboxes'].clone()
            bboxes.translate_([random_shift_x, random_shift_y])

            # clip border
            bboxes.clip_(img_shape)

            # remove invalid bboxes
            valid_inds = (bboxes.widths > self.filter_thr_px).numpy() & (
                bboxes.heights > self.filter_thr_px).numpy()
            # If the shift does not contain any gt-bbox area, skip this
            # image.
            if not valid_inds.any():
                return results
            bboxes = bboxes[valid_inds]
            results['gt_bboxes'] = bboxes
            results['gt_bboxes_labels'] = results['gt_bboxes_labels'][
                valid_inds]

            if results.get('gt_ignore_flags', None) is not None:
                results['gt_ignore_flags'] = \
                    results['gt_ignore_flags'][valid_inds]

            # shift img
            img = results['img']
            new_img = np.zeros_like(img)
            img_h, img_w = img.shape[:2]
            new_h = img_h - np.abs(random_shift_y)
            new_w = img_w - np.abs(random_shift_x)
            new_img[new_y:new_y + new_h, new_x:new_x + new_w] \
                = img[ori_y:ori_y + new_h, ori_x:ori_x + new_w]
            results['img'] = new_img

        return results

需要注意的是, 经过pipeline后, 字典最终会被更新成如下形式:

dict = {'inputs': torch.Tensor, 
        'data_samples': DetDataSample或TrackDataSample等}

其中'inputs'键对应的值就是转换为tensor的图片, 而'data_samples'键对应的值是表示样本的类, 在检测任务中, 是DetDataSample, 跟踪任务中, 是TrackDataSample. DetDataSample类有许多成员, 包括该样本(图片)的目标的边界框真值, 分割真值等:

class DetDataSample(BaseDataElement):
    """A data structure interface of MMDetection. They are used as interfaces
    between different components.

    The attributes in ``DetDataSample`` are divided into several parts:

        - ``proposals``(InstanceData): Region proposals used in two-stage
            detectors.
        - ``gt_instances``(InstanceData): Ground truth of instance annotations.
        - ``pred_instances``(InstanceData): Instances of detection predictions.
        - ``pred_track_instances``(InstanceData): Instances of tracking
            predictions.
        - ``ignored_instances``(InstanceData): Instances to be ignored during
            training/testing.
        - ``gt_panoptic_seg``(PixelData): Ground truth of panoptic
            segmentation.
        - ``pred_panoptic_seg``(PixelData): Prediction of panoptic
           segmentation.
        - ``gt_sem_seg``(PixelData): Ground truth of semantic segmentation.
        - ``pred_sem_seg``(PixelData): Prediction of semantic segmentation.

以上过程可以借用MMEngine文档里的一个图说明:

最终, 模型的forward, loss, predict等方法都是接收inputs: torch.Tensor与data_samples作为输入, 例如:

def loss(self, inputs: Tensor, data_samples: TrackSampleList,
             **kwargs) -> Union[dict, tuple]: