0 摘要

轻松掌握 MMDetection 训练测试流程(二)对整个目标检测框架的训练以及测试整体流程进行的详细分析，旨在帮助大家对 MMDetection 框架进行较为全面的理解。在此基础上，本文主要针对 Model 模块中最为复杂的 Head 部分进行详解，主要内容包括：

  • Head 模块整体概述
  • Head 模块构建流程
  • Head 模块源码分析

    1 Head 模块整体概述

    v2-ba9edb24f8cbf10ee77bacb7f10befa7_720w.jpg
    目前 MMDetection 中 Head 模块主要是按照 stage 来划分，主要包括两个 package: dense_heads 和 roi_heads , 分别对应 two-stage 算法中的第一和第二个 stage 模块，如果是 one-stage 算法则仅仅有 dense_heads 而已。

(1) dense_heads dense_heads 部分主要是按照 anchor-based 和 anchor-free 来划分，对应的类是 AnchorHead 和 AnchorFreeHead, 这两个类主要区别是 AnchorHead 会额外需要 anchor_generator 配置，用于生成默认 anchor。

同时可以看到有些类并没有直接继承这两个基类，例如 YOLOV3Head。原因是在该类中大部分函数处理逻辑都需要复写，为了简单就直接继承了 BaseDenseHead，而对于 SABLRetinaHead 而言，由于 SABL 是类似 anchor-based 和 anchor-free 混合的算法，故直接继承 BaseDenseHead 是最合适的做法。用户如果要进行扩展开发，可以依据开发便捷度自由选择最合适的基类进行继承。

(2) roi_heads roi_heads 部分主要是按照第二阶段内部的 stage 个数来划分，经典的 Faster R-CNN 采用的是 StandardRoIHead，表示进行一次回归即可，而对于 Cascade R-CNN，其第二阶段内部也包括多个 stage 回归阶段，实现了 CascadeRoIHead，即可以构建任意次数的分类回归结果。

总结来说，每个 Head 内部都可能包括:

  1. RoI 特征提取器 roi_extractor
  2. 共享模块 shared_heads
  3. bbox 分类回归模块 bbox_heads
  4. mask 预测模块 mask_heads

其中1、3是必备模块。

2 Head 模块构建流程

为了方便理解，首先需要回顾下 MMDetection 训练和测试流程，然后再对每个 Head 模块进行深入分析。

2.1 MMDetection 训练和测试流程

(1) 训练流程

对应 two-stage 而言，具体如下所示：

  1. #============= mmdet/models/detectors/two_stage.py/TwoStageDetector ============
  2. def forward_train(...):
  3.     # 先进行 backbone+neck 的特征提取
  4.     x = self.extract_feat(img)
  5.     losses = dict()
  6.     # RPN forward and loss
  7.     if self.with_rpn:
  8.         # 训练 RPN
  9.         proposal_cfg = self.train_cfg.get('rpn_proposal',
  10.                                         self.test_cfg.rpn)
  11.         # 主要是调用 rpn_head 内部的 forward_train 方法
  12.         rpn_losses, proposal_list = self.rpn_head.forward_train(x,...)
  13.         losses.update(rpn_losses)
  14.     else:
  15.         proposal_list = proposals
  16.     # 第二阶段，主要是调用 roi_head 内部的 forward_train 方法
  17.     roi_losses = self.roi_head.forward_train(x, ...)
  18.     losses.update(roi_losses)
  19.     return losses

Head 模块核心是调用 self.rpn_head.forward_train 和 self.roi_head.forward_train 函数，输出 losses 和其他相关数据。

对于 one-stage 而言，具体如下所示：

  1. #============= mmdet/models/detectors/single_stage.py/SingleStageDetector ============
  2. def forward_train(...):
  3.     super(SingleStageDetector, self).forward_train(img, img_metas)
  4.     # 先进行 backbone+neck 的特征提取
  5.     x = self.extract_feat(img)
  6.     # 主要是调用 bbox_head 内部的 forward_train 方法
  7.     losses = self.bbox_head.forward_train(x, ...)
  8.     return losses

这个比 two-stage head 模块简单，因为其只有第一个 stage，对应的函数是 self.bbox_head.forward_train。


(2) 测试流程

  1. 调用 MMDataParallel 或 MMDistributedDataParallel 中的 forward 方法
  2. 调用 base.py 中的 forward 方法
  3. 调用 base.py 中的 self.forward_test 方法
  4. 如果是单尺度测试，则会调用 TwoStageDetector 或 SingleStageDetector 中的 simple_test 方法，如果是多尺度测试，则调用 aug_test 方法
  5. 最终调用的是每个具体 Head 模块的 simple_test 或者 aug_test 方法(one-stage 和 two-stage 的 head 调用逻辑有些区别)


可以看出在测试阶段，主要是调用了 Head 模块自身的 simple_test 或 aug_test 方法。

2.2 dense_heads 模块训练和测试流程

(1) 训练流程


dense_heads 训练流程最外层函数是 forward_train, 其实现是在mmdet/models/dense_heads/base_dense_head.py/BaseDenseHead 中，如下所示：

  1. def forward_train(self,
  2.                   x,
  3.                   img_metas,
  4.                   gt_bboxes,
  5.                   gt_labels=None,
  6.                   gt_bboxes_ignore=None,
  7.                   proposal_cfg=None,
  8.                   **kwargs):
  9.     # 调用各个子类实现的 forward 方法
  10.     outs = self(x)
  11.     if gt_labels is None:
  12.         loss_inputs = outs + (gt_bboxes, img_metas)
  13.     else:
  14.         loss_inputs = outs + (gt_bboxes, gt_labels, img_metas)
  15.     # 调用各个子类实现的 loss 计算方法
  16.     losses = self.loss(*loss_inputs, gt_bboxes_ignore=gt_bboxes_ignore)
  17.     if proposal_cfg is None:
  18.         return losses
  19.     else:
  20.         # two-stage 算法还需要返回 proposal
  21.         proposal_list = self.get_bboxes(*outs, img_metas, cfg=proposal_cfg)
  22.         return losses, proposal_list


每个算法的 Head 子类一般不会重写上述方法，但是每个 Head 子类都会重写 forward 和 loss 方法，其中 forward 方法用于运行 head 网络部分输出分类回归分支的特征图，而 loss 方法接收 forward 输出，并且结合 label 计算 loss。


(1-1) BaseDenseHead

BaseDenseHead 基类过于简单，对于 anchor-based 和 anchor-free 算法又进一步进行了继承，得到 AnchorHead 或者 AnchorFreeHead 类。在目前的各类算法实现中，绝大部分子类都是继承自 AnchorHead 或者 AnchorFreeHead，其提供了一些相关的默认操作，如果直接继承 BaseDenseHead 则子类需要重写大部分算法逻辑。

(1-2) AnchorHead

首先分析 AnchorHead，其主要是封装了 anchor 生成过程。下面对 forward 和 loss 函数进行分析

  1. # BBoxTestMixin 是多尺度测试时候调用
  2. class AnchorHead(BaseDenseHead, BBoxTestMixin):
  3.     # feats 是 backbone+neck 输出的多个尺度图
  4.     def forward(self, feats):
  5.         # 对每张特征图单独计算预测输出
  6.         return multi_apply(self.forward_single, feats)
  8.     # head 模块分类回归分支输出
  9.     def forward_single(self, x):
  10.         cls_score = self.conv_cls(x)
  11.         bbox_pred = self.conv_reg(x)
  12.         return cls_score, bbox_pred


forward 函数比较简单，就是对多尺度特征图中每个特征图分别计算分类和回归输出即可，主要复杂度在 loss 函数中，其运行流程图如下所示：

v2-6eb148b2d43b9e30471caf886701afda_720w.jpg

  • 在 loss 函数中首先会调用 get_anchors 函数得到默认 anchor 列表。而 get_anchors 函数内部会先计算多尺度特征图上每个特征点位置的 anchor，然后再计算有效 anchor 标志(因为在组织 batch 时候有些图片会进行左上角 padding，这部分像素人为加的，不需要考虑 anchor)
  • 然后基于 anchor、gt bbox 以及其他必备信息调用 get_targets 函数计算每个预测分支对应的 target。get_targets 函数内部会调用 multi_apply(_get_targets_single) 函数对每张图片单独计算 target，而 _get_targets_single 函数实现的功能比较多，包括：bbox assigner、bbox sampler 和 bbox encoder 三个关键环节
  • 在得到 targets 后，调用 loss_single 函数计算每个输出尺度的 loss 值，最终返回各个分支的 loss


    (1-3) AnchorFreeHead

AnchorFreeHead 逻辑比 AnchorHead 简单很多，主要是因为 anchor-free 类算法比 anchor-based 算法更加灵活多变，而且少了复杂的 anchor 生成过程，其 forward 方法实现和 AnchorHead 完全相同，而 loss 方法没有实现，其子类必须实现。


(2) 测试流程


前面说过在测试流程中，最终会调用 Head 模块的 simple_test 或 aug_test 方法分别进行单尺度和多尺度测试，涉及到具体代码层面，one-stage 和 two-stage 调用函数有区别，但是最终调用的依然是 Head 模块的 get_bboxes 方法。


(2-1) AnchorHead

在单尺度测试模式下，对于 one-stage 而言，是直接调用 self.bbox_head.get_bboxes 方法，如果是 AnchorHead，其流程是：

  1. 遍历每个特征尺度输出分支，利用 nms_pre 配置参数对该层预测结果按照 scores 值进行从大到小进行 topk 截取，保留 scores 最高的前 nms_pre 的预测结果
  2. 对保留的预测结果进行 bbox 解码还原操作
  3. 还原到最原始图片尺度
  4. 如果需要进行 nms，则对所有分支预测保留结果进行统一 nms 即可，否则直接属于多尺度预测结果


对于 two-stage 而言，其第一阶段 Head 推理是直接调用了 simple_test_rpn 方法，该方法内部最终也是调用了 AnchorHead 中的get_bboxes 方法。

  1. # mmdet/models/dense_heads/rpn_test_mixin.py/RPNTestMixin
  2. def simple_test_rpn(self, x, img_metas):
  3.     rpn_outs = self(x)
  4.     proposal_list = self.get_bboxes(*rpn_outs, img_metas)
  5.     return proposal_list


(2-2) AnchorFreeHead

AnchorFreeHead 比较灵活， get_bboxes 都是由具体算法子类实现。

  1. def get_bboxes(self,
  2.                cls_scores,
  3.                bbox_preds,
  4.                img_metas,
  5.                cfg=None,
  6.                rescale=None):
  7.     raise NotImplementedError

(2-3) 多尺度测试


除了 RPN 算法的多尺度测试是在mmdet/models/dense_heads/rpn_test_mixin.py，其余 Head 多尺度测试都是在 mmdet/models/dense_heads/dense_test_mixins.py/BBoxTestMixin 中实现，其思路是对多尺度图片中每张图片单独运行 get_bboxes，然后还原到原图尺度，最后把多尺度图片预测结果合并进行统一 nms。

2.3 roi_heads 模块训练和测试流程

以最常用的 StandardRoIHead 为例进行分析。


(1) 训练流程

训练流程最外层依然是调用 forward_train, 其核心代码如下所示：

  1. def forward_train(self,
  2.                   x,
  3.                   img_metas,
  4.                   proposal_list,
  5.                   gt_bboxes,
  6.                   gt_labels,
  7.                   ...):
  8.     if self.with_bbox or self.with_mask:
  9.         num_imgs = len(img_metas)
  10.         sampling_results = []
  11.         for i in range(num_imgs):
  12.             # 对每张图片进行 bbox 正负样本属性分配
  13.             assign_result = self.bbox_assigner.assign(
  14.                 proposal_list[i], ...)
  15.             # 然后进行正负样本采样
  16.             sampling_result = self.bbox_sampler.sample(
  17.                 assign_result,
  18.                 proposal_list[i],
  19.                 ...)
  20.             sampling_results.append(sampling_result)
  21.     losses = dict()
  23.     if self.with_bbox:
  24.         # bbox 分支 forward，返回 loss
  25.         bbox_results = self._bbox_forward_train(...)
  26.         losses.update(bbox_results['loss_bbox'])
  28.     if self.with_mask:
  29.         # mask 分支 forward,返回 loss
  30.     return losses
  33. def _bbox_forward_train(self, x, sampling_results, gt_bboxes, gt_labels,
  34.                         img_metas):
  35.     rois = bbox2roi([res.bboxes for res in sampling_results])
  36.     # forward
  37.     bbox_results = self._bbox_forward(x, rois)
  38.     # 计算 target
  39.     bbox_targets = self.bbox_head.get_targets(...)  
  40.     # 计算 loss                                          
  41.     loss_bbox = self.bbox_head.loss(...)
  42.     return ...    
  44. def _bbox_forward(self, x, rois):
  45.     # roi 提取
  46.     bbox_feats = self.bbox_roi_extractor(
  47.         x[:self.bbox_roi_extractor.num_inputs], rois)
  48.     # bbox head 网络前向
  49.     cls_score, bbox_pred = self.bbox_head(bbox_feats)
  50.     return ...

从上述逻辑可以看出，StandardRoIHead 中 forward_train 函数仅仅是对内部的 bbox_head 相关函数进行调用，例如 get_targets 和 loss，本身 StandardRoIHead 类不做具体算法逻辑计算。

可以参考 Faster R-CNN 配置文件理解 StandardRoIHead 和 bbox_head 的关系：

  1. roi_head=dict(
  2.     type='StandardRoIHead',
  3.     bbox_roi_extractor=dict(
  4.         type='SingleRoIExtractor',
  5.         roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
  6.         out_channels=256,
  7.         featmap_strides=[4, 8, 16, 32]),
  8.     bbox_head=dict(
  9.         type='Shared2FCBBoxHead',
  10.         in_channels=256,
  11.         fc_out_channels=1024,
  12.         roi_feat_size=7,
  13.         num_classes=80,
  14.         bbox_coder=dict(
  15.             type='DeltaXYWHBBoxCoder',
  16.             target_means=[0., 0., 0., 0.],
  17.             target_stds=[0.1, 0.1, 0.2, 0.2]),
  18.         reg_class_agnostic=False,
  19.         loss_cls=dict(
  20.             type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
  21.         loss_bbox=dict(type='L1Loss', loss_weight=1.0))))


StandardRoIHead 类包装了 bbox_roi_extractor 和 bbox_head 的实例，前者用于 RoI 特征提取，后者才是真正计算分类和回归的逻辑。在 bbox_head 中除了网络模型有些变换外，loss计算过程是非常类似的，其 get_targets 和 loss 计算过程都是封装在基类 mmdet/models/roi_heads/bbox_heads/bbox_head.py 中。


(2) 测试流程

测试流程是调用 Head 模块的 simple_test 和 aug_test 函数，单尺度测试 bbox 相关实现代码在 mmdet/models/roi_heads/test_mixins.py/BBoxTestMixin 的 simple_test_bboxes 函数中。

  1. def simple_test_bboxes(self,
  2.                        x,
  3.                        ...):
  4.     rois = bbox2roi(proposals)
  5.     # roi 提取+ forward，输出预测结果
  6.     bbox_results = self._bbox_forward(x, rois)
  7.     cls_score = bbox_results['cls_score']
  8.     bbox_pred = bbox_results['bbox_pred']
  9.     det_bboxes = []
  10.     det_labels = []
  11.     for i in range(len(proposals)):
  12.         # 对预测结果进行解码输出 bbox 和对应 label
  13.         det_bbox, det_label = self.bbox_head.get_bboxes(...)
  14.         det_bboxes.append(det_bbox)
  15.         det_labels.append(det_label)
  16.     return det_bboxes, det_labels

实际上依然是调用了 Head 模块内部的 get_bboxes 函数，处理逻辑和 dense_head 差不多( 解码+还原尺度+ nms)。

3 总结


本文是 MMDetection 流程解读三部曲的最终章，重点是对最复杂的 Head 模块进行深入详细解读。通过三篇解读，读者可以了解到：

  • MMDetection 框架的整体设计思想和算法模块划分原则
  • MMDetection 框架的整体训练和测试流程
  • MMDetection 框架每个组件的详细代码实现过程
  • 针对任何一个新复现代码，能够很快理解 MMDetection 是如何通过模块组合实现的

    参考

https://zhuanlan.zhihu.com/p/343433169