轻松掌握 MMDetection 中 Head 流程

文@000007

0 摘要

轻松把握 MMDetection 练习测试流程(二)对全部方针检测框架的练习以及测试整体流程停止的具体分析，旨在帮助大师对 MMDetection 框架停止较为周全的了解。在此根本上，本文首要针对 Model 模块中最为复杂的 Head 部分停止详解，首要内容包括：

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

项目地址：https://github.com/open-mmlab/mmdetection
接待 star

1 Head 模块整体概述

今朝 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 而言，具体以下所示：

Head 模块焦点是挪用 self.rpn_head.forward_train 和 self.roi_head.forward_train 函数，输出 losses 和其他相关数据。

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

这个比 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 中，以下所示：

每个算法的 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 函数停止分析

forward 函数比力简单，就是对多标准特征图中每个特征图别离计较分类和回归输出即可，首要复杂度在 loss 函数中，其运转流程图以下所示：

• 在 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 方式。

(2-2) AnchorFreeHead

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

(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, 其焦点代码以下所示：

从上述逻辑可以看出，StandardRoIHead 中 forward_train 函数仅仅是对内部的 bbox_head 相关函数停止挪用，例如 get_targets 和 loss，自己 StandardRoIHead 类不做具体算法逻辑计较。

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

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 函数中。

现实上仍然是挪用了 Head 模块内部的 get_bboxes 函数，处置逻辑和 dense_head 差不多( 解码+复原标准+ nms)。

3 总结

本文是 MMDetection 流程解读三部曲的终极章，重点是对最复杂的 Head 模块停止深入具体解读。经过三篇解读，读者可以领会到：

• MMDetection 框架的整体设想思惟和算法模块分别原则
• MMDetection 框架的整体练习和测试流程
• MMDetection 框架每个组件的具体代码实现进程
• 针对任何一个新复现代码，可以很快了解 MMDetection 是若何经过模块组合实现的

快速指引：

快速指引：