maskrcnn-benchmark-master（七）：build_roi_heads()函数

前言 

截至上一篇博客，我们已经介绍完RPN的整个流程，RPN的作用是提取Proposals传入ROI_heads部分，然而ROI有哪些heads呢？根据不同的任务，代码分为了box_head，keypoint_head，mask_head。

box_head：进行bounding box的回归以及类别的分类任务。

keypoint_head：进行关键点的检测。（这部分代码我已经从项目中去除掉了）

mask_head：进行mask的分割任务。

 一、CombinedROIHeads类

接下来我们将开启ROI_heads的介绍篇章，build_roi_heads()函数在your_project/maskrcnn_benchmark/modeling/roi_heads/roi_heads.py文件中，我们打开这个文件,首先引入眼帘的是CombinedROIHeads类，通过这个类的名字，我们易可知这个类的作用是将box_head，keypoint_head，mask_head这几个模块都整合在一起，我们来看一下相关代码（关于keypoint_head的代码我已经删掉了~）：

class CombinedROIHeads(torch.nn.ModuleDict):
    """
    Combines a set of individual heads (for box prediction or masks) into a single
    head.

    合并许多个单独的head为一个统一的heads
    """

    def __init__(self, cfg, heads):
        super(CombinedROIHeads, self).__init__(heads)
        self.cfg = cfg.clone()
        # 如果box和mask的head的特征共享，则将box head的features 赋值给mask head
        if cfg.MODEL.MASK_ON and cfg.MODEL.ROI_MASK_HEAD.SHARE_BOX_FEATURE_EXTRACTOR:
            self.mask.feature_extractor = self.box.feature_extractor

    def forward(self, features, proposals, targets=None):
        losses = {}
        # TODO rename x to roi_box_features, 
        # if it doesn't increase memory consumption
        # box head的loss
        # self.box就是一个box_head的对象（后续的博客再介绍） 
        # 返回结果是box_head部分提取的特征，detections是检测的结果，loss_box是损失函数
        x, detections, loss_box = self.box(features, proposals, targets)
        losses.update(loss_box)

        # 如果存在mask 分支
        if self.cfg.MODEL.MASK_ON:
            mask_features = features
            # optimization: during training, if we share the feature extractor between
            # the box and the mask heads, then we can reuse the features already computed
            # 优化：在训练阶段，如果我们共享了box head 和 mask head的特征提取器， 
            # 我们可以重复使用box head所计算的feature用于mask head
            if (
                self.training
                and self.cfg.MODEL.ROI_MASK_HEAD.SHARE_BOX_FEATURE_EXTRACTOR
            ):
                # 如果mask_feature共享box的特征 
                # 就将box_head部分提取的特征赋予mask_features
                mask_features = x
            # During training, self.box() will return the unaltered proposals as "detections"
            # this makes the API consistent during training and testing
            # 训练阶段， self.box() 会返回未经变换的proposals作为检测结果
            # 将mask的的检测结果加入detections中，并计算mask的loss返回。
            x, detections, loss_mask = self.mask(mask_features, detections, targets)
            losses.update(loss_mask)

        return x, detections, losses


# 从上述代码可以看出box_head和mask_head和之前介绍的rpn_heads很相像，
# 返回的结果都包含有检测的结果和loss
# rpn_head返回:Proposals(相当于RPN检测的bounding box 和类别结果), rpn_loss
# box_head返回:提取的特征x, 检测的bounding box和类别分类结果detections, box_loss
# mask_head返回:提取的特征x, 检测的mask结果并加上之前的box_head的检测结果, mask_loss
# 因此推断box_head对象和mask_head对象中应该也是包含有相应的loss计算文件和inference文件的

二、build_roi_heads()函数

接着我们看build_roi_heads()函数的相关代码:

# 创建roi heads
def build_roi_heads(cfg, in_channels):
    # individually create the heads, that will be combined together
    # afterwards
    roi_heads = []
    if cfg.MODEL.RETINANET_ON:
        return []

    # 根据配置文件依次添加各个head
    if not cfg.MODEL.RPN_ONLY:
        # 添加boxes head
        # 通过build_roi_box_head()创建roi_box_head分支
        roi_heads.append(("box", build_roi_box_head(cfg, in_channels)))
    if cfg.MODEL.MASK_ON:
        # 添加mask head
        # 通过build_roi_mask_head()创建roi_mask_head分支
        roi_heads.append(("mask", build_roi_mask_head(cfg, in_channels)))


    # combine individual heads in a single module
    # 将独立的分支进行合并
    if roi_heads:
        roi_heads = CombinedROIHeads(cfg, roi_heads)

    return roi_heads

从build_roi_heads代码可以看出，它主要是通过build_roi_box_head()函数和build_roi_mask_head()函数来生成box_head和mask_head分支，最后通过CombinedROIHeads类将这两个分支进行整合，整体的结构简图如下所示：

build_roi_box_head（）函数将放在下一篇博客介绍:

maskrcnn-benchmark-master（八）：build_roi_box_head()函数

build_roi_mask_head（）函数待续~