Maskrcnn-Benchmark骨干网络分析

不论是在训练脚本文件 train_net.py 还是在测试脚本文件 test_net.py 中, 都调用了 build_detection_model(cfg) 函数来创建模型, 该函数封装了模型定义的内部细节, 使得我们可以通过cfg配置文件轻松的组合出不同类型的模型, 为了能够更好的了解模型的内部细节, 我们有必要知道这些模型是如何被定义, 又是如何组合到一起的, 为此我们需要对 MaskrcnnBenchmark 的 modeling 文件夹进行解析, 该文件夹的结构位于 ./maskrcnn_benchmark/modeling/ 文件夹下。

首先detector模型定义入口

第一部分是detector文件夹，该文件夹中的两个文件（detectors.py，generalized_rcnn.py）定义了是整个modeling模块的入口，解析如下：

detector.py

主要功能就是根据给定的配置信息实例化一个 class GeneralizedRCNN 的对象, 代码如下所示:

# ./maskrcnn_benchmark/modeling/detector/detectors.py

from .generalized_rcnn import GeneralizedRCNN
#字典类型
_DETECTION_META_ARCHITECTURES = {"GeneralizedRCNN": GeneralizedRCNN}

# 该函数是创建模型的入口函数, 也是唯一的模型创建函数
def build_detection_model(cfg):
    meta_arch = _DETECTION_META_ARCHITECTURES[cfg.MODEL.META_ARCHITECTURE]
    # 下面的语句等价于
    # return GeneralizedRCNN(cfg)
    return meta_arch(cfg)

上面的代码利用配置信息cfg实例化了一个 GeneralizedRCNN 类。

generalized_rcnn.py
该文件定义了 Maskrcnn-Benchmark 的 GeneralizedRCNN 类, 用于表示各种组合后的目标检测模型, 代码解析如下:

import torch
from torch import nn

# 该函数定义于 ./maskrcnn_benchmark/structures/image_list.py 文件中
from maskrcnn_benchmark.structures.image_list import to_image_list

from ..backbone import build_backbone
from ..rpn.rpn import build_rpn
from ..roi_heads.roi_heads import build_roi_heads

# 定义类的具体实现
class GeneralizedRCNN(nn.Module):
    # 该类是 MaskrcnnBenchmark 中所有模型的共同抽象, 目前支持 boxes 和 masks 两种形式的标签
    # 该类主要包含以下三个部分:
    # - backbone
    # - rpn(option)
    # - heads: 利用前面网络输出的 features 和 proposals 来计算 detections / masks.

    def __init__(self, cfg): # 根据配置信息初始化模型
        super(GeneralizedRCNN, self).__init__()

        # 根据配置信息创建 backbone 网络
        self.backbone = build_backbone(cfg)

        # 根据配置信息创建 rpn 网络
        self.rpn = build_rpn(cfg)

        # 根据配置信息创建 roi_heads
        self.roi_heads = build_roi_heads(cfg)

    def forward(self, images, targets=None): # 定义模型的前向传播过程
        # images (list[Tensor] or ImageList)
        # targets (list[BoxList])
        # 返回值: result (list[BoxList] or dict[Tensor])
        # 在训练阶段, 返回字典类型的模型损失, 在测试阶段, 返回模型的预测结果.

        # 当 training 设置为 True 时, 必须提供 targets.
        if self.training and targets is None:
            raise ValueError("In training mode, targets should be passed")

        images = to_image_list(images) # 将图片的数据类型转换成 ImageList

        # 利用 backbone 网络获取图片的 features
        features =  self.backbone(images.tensors)

        # 利用 rpn 网络获取 proposals 和相应的 loss
        proposals, proposal_losses = self.rpn(images, features, targets)
        if self.roi_heads: # 如何 roi_heads 不为 None 的话, 就计算其输出的结果
            x, result, detector_losses = self.roi_heads(features, proposals, targets)
        else:
            # RPN-only models don't have roi_heads
            x = features
            result = proposals
            detector_losses = {}
        if self.training: # 在训练模式下, 输出损失值
            losses = {}
            losses.update(detector_losses)
            losses.update(proposal_losses)

        return result # 如果不在训练模式下, 则输出模型的预测结果.

Maskrcnn-Benchmark 模型的创建主要依赖于三个函数, 即 build_backbone(cfg), build_rpn(cfg), build_roi_heads(cfg).