Detectron 源码解析-模型创建

发表于   |   更新于 | 分类于 | 评论数： | 阅读次数：
本文字数： 22k | 阅读时长 ≈ 20 分钟

源码文件

  • detectron/modeling/model_builder.py

model_builder.py 文件总览

该文件的路径为 ./detectron/modeling/model_builder.py, 顾名思义, 该文件主要用于模型的创建, Detectron 支持的模型种类有很多, 因此该文件的代码量较大, 将近700行. 通常情况下, 一个模型是下列元素的一个笛卡尔积(cartesian product):

  • backbone: 如VGG16, ResNet, ResNeXt
  • FPN: on or off(使用或者不使用)
  • RPN only: 提供候选区域框
  • Fixed proposals for Fast R-CNN, RFCN, Mask R-CNN(可以带有keypoints)
  • 端到端模型: RPN + Fast R-CNN, Mask R-CNN, …
  • 不同的“head”
  • 大量配置选项(configuration options)

最终的模型就是将上面这些“基本元素”组合起来构成的完整的模型. 可以看出, 这种构建模型的方式非常灵活, 但是缺点是刚开始的时候比较难理解. 不过放心, 我们一起来慢慢学习 :)

我们可以通过搭配不同的基本模型来组成一个新的模型, 例如, 如果我们想要搭建一个用 ResNet-50-C4 网络作为 backbone 的 Fast R-CNN 模型的话, 那么我们就可以在配置信息中进行定制:

1
2
3
4
5
# .yaml
MODEL:
    TYPE: generalized_rcnn
    CONV_BODY: ResNet.add_ResNet50_conv4_body
    ROI_HEAD: ResNet.add_ResNet_roi_conv5_head

下面我们对model_builder.py文件先进行一个概览, 了解一下该文件都有哪些函数, 以及实现了哪些功能.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
# ./detectron/modeling/model_builder.py

from __future__ import absolute_import
# ... 导入各种包及函数


def generalized_rcnn(model):
    #... 返回一个通用模型, 可以处理 Fast, Faster, FPN, MASK 等模型

def rfcn(model):
    # 调入 build_generic_detection_model 函数

def retinanet(model):
    # 调入 build_generic_detection_model 函数

#-------分割线------#
#Helper functions for building various re-usable network bits
#-------分割线------#

def create(model_type_func, train=False, gpu_id=0):
    #... 创建通用模型, 并将其分派给某个特定的模型搭建函数(model building functions)

def get_func(func_name):
    #... 通过名字返回一个函数对象

def build_generic_detection_model(...):
    def _single_gpu_build_func(model):
        #... 在单一的GPU上搭建模型

def _narrow_to_fpn_roi_levels(blobs, spatial_scales):
    #... 返回用于 RoI 头部的 blobs 和 spatial scales

def _add_fast_rcnn_head(...):
    #... 将 Fast R-CNN 的头部添加到模型中

def _add_roi_mask_head(...):
    #... 将 mask prediction head 添加到模型中

def _add_roi_keypoint_head(...):
    #... 将 keypoint prediction head 添加到模型中

def build_generic_rfcn_model(...):
    def _single_gpu_build_func(model):
        #... 该函数已调入到 build_generic_detection_model 函数中

def build_generic_retinanet_model(...):
    def _single_gpu_build_func(model):
        # #... 该函数已调入到 build_generic_detection_model 函数中

#----------分割线----------#
# Network inputs
#----------分割线----------#

def add_training_inputs(model, roidb=None):
    #... 创建用于训练网络的 input ops 和 blobs. 必须在 model_builder.create() 之后调用

def add_inference_inputs(model):
    #... 创建用于测试网络(inference)的 inputs blobs
    def create_input_blobs_for_net(net_def):
        #...

# ----------分割线---------- #
# DEPRECATED FUNCTIONALITY BELOW
# ----------分割线---------- #

# 余下的代码已经被弃用(deprecated), 故不做介绍
# ...
# 下面的模型创建代码没有通过 generic composable 来创建模型
# 而是通过分开定义(硬编码, hardcoded)的方式独立创建, 模型之间比较独立.

根据以上函数功能及相互之间的关键, 我们将从以下几个方面对模型创建的源码进行解析:

  • 包及函数的导入: 简单介绍导入的相关包和函数, 它们的功能, 以及在文件中的使用情况, 同时会给出相关的详细解析的链接
  • 模型创建核心函数: 对 create(), get_func(), build_generic_detection_model() 等关键函数进行解析
  • 模型头部解析: 对于不同的任务, 通常具有不同的网络模型 heads, 因此, 本部分主要介绍这些 heads 的创建.
  • 网络模型输入: 除了模型之外, 还要定义相应的输入接口以便模型训练和模型推演.

包及函数的导入

下面我们对文件导入的包和相关的函数做一个简单的介绍, 并给出其相应的详细解析的链接, 有兴趣的朋友可以点击链接直接查看相关的信息.

首先是一些常用包, 这部分简单看一下就可以, 了解一下使用了哪些常用包

1
2
3
4
5
6
7
8
9
10
11
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import copy
import importlib
import logging

from caffe2.python import core
from caffe2.python import workspace

接下来是 Detectron 内部的包及函数, 我们简要介绍一下这些包和函数的相关作用, 以及本文件调用了包的哪些属性和函数.

1
2
3
4
5
6
7
8
9
10
11
12
13
from detectron.core.config import cfg
from detectron.modeling.detector import DetectionModelHelper
from detectron.roi_data.loader import RoIDataLoader
import detectron.modeling.fast_rcnn_heads as fast_rcnn_heads
import detectron.modeling.keypoint_rcnn_heads as keypoint_rcnn_heads
import detectron.modeling.mask_rcnn_heads as mask_rcnn_heads
import detectron.modeling.name_compat as name_compat
import detectron.modeling.optimizer as optim
import detectron.modeling.retinanet_heads as retinanet_heads
import detectron.modeling.rfcn_heads as rfcn_heads
import detectron.modeling.rpn_heads as rpn_heads
import detectron.roi_data.minibatch as roi_data_minibatch
import detectron.utils.c2 as c2_utils

上面的包及函数的主要功能如下所示:

  • from detectron.core.config import cfg: 导入全局配置信息, 详细解析:全局配置选项及相关函数
  • from detectron.modeling.detector import DetectionModelHelper: DetectionModelHelper 类, 用于表示 Detectron 模型, 详细解析:DetectionModelHelper
  • from detectron.roi_data.loader import RoIDataLoader: RoIDataLoader 类, 完成了 Detectron 加载数据的功能, 详细解析:数据载入源码解析
  • import detectron.modeling.fast_rcnn_heads as fast_rcnn_head: 定义了 Fast R-CNN 的头部, 详细解析: Fast R-CNN Heads
  • import detectron.modeling.keypoint_rcnn_heads as keypoint_rcnn_heads: 定义了 Keypoint R-CNN 的头部, 详细解析: Keypoint R-CNN Heads
  • import detectron.modeling.mask_rcnn_heads as mask_rcnn_heads: 定义了 Mask R-CNN 的头部, 详细解析: Mask R-CNN Heads
  • import detectron.modeling.name_compat as name_compat: 对旧版的名称做了一些修改, 本部分代码用于提供与旧版名称的兼容性.
  • import detectron.modeling.optmizer as optim: 构建网络优化器
  • import detectron.modeling.retinanet_heads as retinaneet_heads: 定义了 RetinaNet 的头部, 详细解析: RetinaNet Heads
  • import detectron.modeling.rfcn_heads as rfcn_heads: 定义了 RFCN 的头部. 详细解析: 其他网络的Heads定义
  • import detectron.modeling.rpn_heads as rpn_heads: 定义了 RPN 的头部. 详细解析: 其他网络的Heads定义
  • import detectron.roi_data.minibatch as roi_data_minibatch: 构建数据集的 mini batch. 详细解析: 数据加载解析
  • import detectron.utils.c2 as c2_utils: 详细解析: 工具函数解读

模型创建核心函数

model_builder.create() 方法

在调用detectron/utils/train.py文件中的train_model()方法时,在第一行代码中调用了同属该文件的create_model()方法, 而在该方法中, 核心的创建语句为model = model_builder.create(cfg.MODEL.TYPE, train=True), 其中, model_builder.create(...) 是位于文件 detectron/modeling/model_builder.py中create(...)方法, 该方法的详细实现如下所示:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
# detectron/modeling/model_builder.py
def generalized_rcnn(model):
    #... 通用模型
def rfcn(model):
    #...
def retinanet(model):
    #...
def create(model_type_func, train=False, gpu_id = 0):
    # 通用的模型创建函数, 该函数可以继续分派到特定的模型创建函数中
    # 默认情况下, 该函数将以并行模式(并行数量取决于cfg.NUM_GPUS)生成数据
    # 但是, 你可以将其限制在特定的GPU上进行(通过gpu_id), 在测试阶段使用optimizer.build_data_parallel_model()

    # model_type_func 通常为 cfg.MODEL.TYPE, 即: generalized_rcnn.
    # 在 test_engin.py 中的调用形式:
    # model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)

    # from detectron.modeling.detector import DetectionModelHelper
    model = DetectionModelHelper(
        name=model_type_func, # 对于示例: model_type_func=generalized_rcnn
        train=train,
        num_classes=cfg.MODEL.NUM_CLASSES,
        init_params=train
    )
    model.only_build_forward_pass = Fasle
    model.target_gpu_id = gpu_id
    return get_func(model_type_func)(model) # get_func(...)函数解析就在下面

#...

上面代码中 model 对象是 DetectionModelHelper 类的一个实例, 关于此类的详细解析可以看DetectionModelHelper 解析.
接着, 设置了该对象的两个属性, 分别为 model.only_build_forward_pass 以及 model.target_gpu_id. 接着调用 get_func 获取到相应的函数, 并将 model 对象作为参数传递给了该函数, 然后将函数的返回值作为 create() 函数的返回值返回, 返回的是一个模型对象 model(常用名).
在获得模型对象 model 后, 将其返回. 通常情况下, 会对返回的模型对象进行权重初始化, 如下所示:

1
2
3
4
5
6
7
from detectron.modeling import model_builder
import detectron.utils.net as net_utils

model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id) # gpu_id = 0
net_utils.initialize_gpu_from_weights_file(
    model, weights_file, gpu_id=gpu_id
) # gpu_id = 0

get_func() 函数

由上面的 create() 函数可知, 最终的返回语句在返回直接会利用 get_func() 函数获取名字对于的函数变量, get_func() 函数的解析如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
# ./detectron/modeling/model_builder.py


def get_func(func_name):
    # 该函数会通过name返回一个函数对象,
    # func_name必须等于该module中的函数, 或者是与base 'modeling'相关的函数的路径

    # 通常: func_name = cfg.MODEL.TYPE = generalized_rcnn
    if func_name == '':
        return None
    # import detectron.modeling.name_compat as name_compat
    # 因为名字做过改动, 这句话是为了兼容性而存在的
    new_func_name = name_compat.get_new_name(func_name)
    if new_func_name != func_name:
    # 对于本例: new_func_name = func_name, 名字不变, 依然为: cfg.MODEL.TYPE = generalized_rcnn
    # 因此不会出现下面的警告
        logger.warn(
            'Remapping old function name:{} -> {}'.
            format(func_name, new_func_name)
        )
        func_name = new_func_name
    try:
        parts = func_name.split('.')
        if len(parts) == 1:
            return globals()[parts[0]]
        module_name = 'detectron.modeling.' + '.'.join(parts[:-1])
        # 会根据module导入位于'detectrong/modeling'中的module
        # fast_rcnn_heads, FPN, mask_rcnn_heads 等等
        module = importlib.import_module(module_name)
        return getattr(module, parts[-1])
    except Exception:
        logger.error('Failed to find function: {}'.format(func_name))
        raise

如果单纯看上面的函数, 大部分人都不太能理解网络到底是通过什么函数创建的, 因为这里我们好像没有看到任何有关 fast rcnn 或者 mask rcnn 的函数调用信息. 实际上, 这也是 Caffe2 的意思变成风格, 即通过配置文件信息来搭建网络, 首先, 我们要知道, 在当前文件中, 存在这大量的全局变量, 这些变量都得是数据变量, 有的是函数名变量. 因此, 本文件中的 get_func() 函数成了配置文件信息与模型函数之间的一道桥梁, 该函数可以根据传入的不同的func_name变量来获得相应的全局函数名, 然后, 将这些函数返回, 最终完成模型搭建. 用一个例子作说明.
下面的语句是 Detectron 官网的标准示例之一.

1
2
3
4
5
6
python tools/infer_simple.py \
    --cfg configs/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml \
    --output-dir /tmp/detectron-visualizations \
    --image-ext jpg \
    --wts https://s3-us-west-2.amazonaws.com/detectron/35861858/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml.02_32_51.SgT4y1cO/output/train/coco_2014_train:coco_2014_valminusminival/generalized_rcnn/model_final.pkl \
    demo

对于上面的指令, 通过--cfg指定了配置文件的位置, 看一下看文件的简单构成:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
MODEL:
  TYPE: generalized_rcnn
  CONV_BODY: FPN.add_fpn_ResNet101_conv5_body
  NUM_CLASSES: 81
  FASTER_RCNN: True
  MASK_ON: True
NUM_GPUS: 8
SOLVER:
  WEIGHT_DECAY: 0.0001
  LR_POLICY: steps_with_decay
  BASE_LR: 0.02
  GAMMA: 0.1
  MAX_ITER: 180000
  STEPS: [0, 120000, 160000]
FPN:
  FPN_ON: True
  MULTILEVEL_ROIS: True
  MULTILEVEL_RPN: True
FAST_RCNN:
  ROI_BOX_HEAD: fast_rcnn_heads.add_roi_2mlp_head
  ROI_XFORM_METHOD: RoIAlign
  ROI_XFORM_RESOLUTION: 7
  ROI_XFORM_SAMPLING_RATIO: 2
MRCNN:
  ROI_MASK_HEAD: mask_rcnn_heads.mask_rcnn_fcn_head_v1up4convs
  RESOLUTION: 28  # (output mask resolution) default 14
  ROI_XFORM_METHOD: RoIAlign
  ROI_XFORM_RESOLUTION: 14  # default 7
  ROI_XFORM_SAMPLING_RATIO: 2  # default 0
  DILATION: 1  # default 2
  CONV_INIT: MSRAFill  # default GaussianFill
TRAIN:
  WEIGHTS: https://s3-us-west-2.amazonaws.com/detectron/ImageNetPretrained/MSRA/R-101.pkl
  DATASETS: ('coco_2014_train', 'coco_2014_valminusminival')
  SCALES: (800,)
  MAX_SIZE: 1333
  BATCH_SIZE_PER_IM: 512
  RPN_PRE_NMS_TOP_N: 2000  # Per FPN level
TEST:
  DATASETS: ('coco_2014_minival',)
  SCALE: 800
  MAX_SIZE: 1333
  NMS: 0.5
  RPN_PRE_NMS_TOP_N: 1000  # Per FPN level
  RPN_POST_NMS_TOP_N: 1000
OUTPUT_DIR: .

可以看到, 在该配置文件中, 网络模型的 CONV_BODY 为: FPN.add_fpn_ResNet101_conv5_body, Fast_RCNN 的头部(用于检测框)为: fast_rcnn_heads.add_roi_2mlp_head, Mask RCNN 的头部为(用于实例分割)为: mask_rcnn_heads.mask_rcnn_fcn_head_v1up4convs. 我们可以在get_func()函数中添加一条输出语句, 来看看创建的网络是否和我们预期的一直, 添加代码如下:

1
2
3
4
5
6
7
8
9
# ./detectron/modeling/model_builder.py

def get_func(func_name):
    if func_name == '':
        return None
    new_func_name = name_compat.get_new_name(func_name)
    print("查看func_name: ", func_name, ", 查看new_func_name: ", new_func_name)

    #...

然后运行上面指令(最好模型创建后的位置加上断点, 因为我们这里只想看模型的创建信息), 输出结果如下:

1
2
3
4
查看func_name: generalized_rcnn, 查看new_func_name: generalized_rcnn
查看func_name: FPN.add_fpn_ResNet101_conv5_body,  查看new_func_name: FPN.add_fpn_ResNet101_conv5_body
查看func_name: fast_rcnn_heads.add_roi_2mlp_head, 查看new_func_name: fast_rcnn_heads.add_roi_2mlp_head
查看func_name: mask_rcnn_heads.mask_rcnn_fcn_head_v1up4convs, 查看new_func_name: mask_rcnn_heads.mask_rcnn_fcn_head_v1up4convs

可以看到, 除了最开始的generalized_rcnn外, 模型还根据配置文件创建了其余的部分. 但是我们并没有显式的看到在哪里使用了这些配置信息, 那么, 到底是在哪里使用了这些配置信息呢? 实际上, 正是get_func()返回的函数generalized_rcnn()起到了关键的作用, 我们知道, get_func()返回的是一个函数地址, 因此, create()函数最后的语句return get_func(model_type_func)(model), 其实也可以看做是return generalized_rcnn(model).

generalized_rcnn() 函数

该函数正是本文件中模型创建部分的第一个函数, 其代码解析如下:

1
2
3
4
5
6
7
8
9
10
11
# ./detectron/modeling/model_builder.py

def generalized_rcnn(model):
    return build_generic_detection_model(
        model,
        get_func(cfg.MODEL.CONV_BODY),
        add_roi_box_head_func=get_func(cfg.FAST_RCNN.ROI_BOX_HEAD),
        add_roi_mask_head_func=get_func(cfg.MRCNN.ROI_MASK_HEAD),
        add_roi_keypoint_head_func=get_func(cfg.KRCNN.ROI_KEYPOINTS_HEAD),
        freeze_conv_body=cfg.TRAIN.FREEZE_CONV_BODY
    )

我们根据示例的指令来对上面的参数简单说明一下

  • model: 由 DetectionModelHelper 创建的一个实例对象.
  • get_func(cfg.MODEL.CONV_BODY): 参数值为FPN.add_fpn_ResNet101_conv5_body, 因此会导入detectron.modeling.FPN, 同时返回FPN.add_fpn_ResNet101_conv5_body()函数.
  • get_func(cfg.FAST_RCNN.ROI_BOX_HEAD): 参数值为fast_rcnn_heads.add_roi_2mlp_head, 因此会导入detectron.modeling.fast_rcnn_heads.add_roi_2mlp_head, 同时返回fast_rcnn_head.add_roi_2mlp_head()函数.
  • get_func(cfg.MRCNN.ROI_MASK_HEAD): 参数值为mask_rcnn_heads.mask_rcnn_fcn_head_v1up4convs, 因此会导入detectron.modeling.mask_rcnn_heads, 同时返回mask_rcnn_heads.mask_rcnn_fcn_head_v1up4convs()函数.
  • get_func(cfg.KRCNN.ROI_KEYPOINTS_HEAD): 参数值为空, 因此不会导入任何文件, 直接返回None.
  • cfg.TRAIN.FREEZE_CONV_BODY: 参数值为默认值False.

build_generic_detection_model() 函数

可以看出, 上面的generalized_rcnn函数是直接调用了build_generic_detection_model() 函数, 该函数内部又新定义了一个函数, 概览如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# ./detectron/modeling/model_builder.py

def build_generic_detection_model(
    model,
    add_conv_body_func,
    add_roi_box_head_hunc=None,
    add_roi_mask_head_func=None,
    add_roi_keypoint_head_func=None,
    freeze_conv_body=False
):
    def _single_gpu_build_func(model):
        #...
    optim.build_data_parallel_model(model, _single_gpu_build_func)
    return model

可以看到, build_generic_detection_model()函数主要是利用函数内部的_single_gpu_build_func()完成模型的搭建, 顾名思义, 该函数主要用于在单个 GPU 上搭建模型, 但是也可以通过在多个不同 GPU 上循环调用来建立多个模型. 在函数的最后, 利用optim的build_data_parallel_model()函数. 该函数位于./detectron/modeling/optimizer.py文件中, 主要提供数据的并行载入功能.

下面我们结合函数中各个参数的含义, 给出build_generic_detection_model()函数代码的详细解析:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
# ./detectron/modeling/model_builder.py

def build_generic_detection_model(
    model,
    add_conv_body_func,
    add_roi_box_head_hunc=None,
    add_roi_mask_head_func=None,
    add_roi_keypoint_head_func=None,
    freeze_conv_body=False
):
    def _single_gpu_build_func(model):
        # 首先添加 conv body(也就是backbone), 如 ResNet-50, VGG等等
        # 根据参数知, 本例中的add_conv_body_func, 实际上就是FPN.add_fpn_ResNet101_conv5_body()
        blob_conv, dim_conv, spatial_scale_conv = add_conv_body_func(model)

        # 固定backbone, 默认为不固定
        if freeze_conv_body:
            # BlobReferenceList() 函数确保参数会以blob引用列表的形式返回.
            for b in c2_utils.BlobReferenceList(blob_conv):
                model.StopGradient(b, b)

        if not model.train: # == inference
            # 创建一个可以在一张图片上执行的 conv body
            # (不带有 RPN 或其他网络头)
            model.conv_body_net = model.net.Clone('conv_body_net')

        # 定义网络头损失函数的梯度字典, 用于存储相应的梯度
        head_loss_gradients={
            'rpn': None,
            'box': None,
            'mask': None,
            'keypoints': None,
        }

# 以下带有_ON的配置大部分默认为False, 所以只有当用户自定义的配置文件显示打开时, 才回执行相关代码
        if cfg.RPN.RPN_ON:
            # 添加 RPN head
            # rpn_heads为.py文件
            # model 参数为函数参数
            # 其余三个参数均来自于 add_conv_body_func(model)
            head_loss_gradients['rpn'] = rpn_heads.add_generic_rpn_outputs(
                model, blob_conv, dim_conv, spatial_scale_conv
            )

        if cfg.FPN_ON:
            # 将添加了RPN head以后,
            # 将FPN的blobs和scales限制到RoI heads使用的大小
            # _narrow_to_fpn_roi_levels函数位于本文件当中, 用于控制大小
            # 参数同样来自于backbone网络
            blob_conv, spatial_scale_conv = _narrow_to_fpn_roi_levels(
                blob_conv, spatial_scale_conv
            )

        if not cfg.MODEL.RPN_ONLY: # 默认为 Fasle
            # 添加 Fast R-CNN head
            # _add_fast_rcnn_head为本文件的函数
            # add_roi_box_head_func为函数参数
            # 其余的3个参数为backbone的返回值
            head_loss_gradients['box'] = _add_fast_rcnn_head(
                model, add_roi_box_head_func, blob_conv, dim_conv, spatial_scale_conv
            )

        if cfg.MODEL.MASK_ON:
            # 添加 Mask 头部
            # _add_roi_mask_head为本文件的函数
            # add_roi_mask_head_func为函数参数
            head_loss_gradients['mask'] = _add_roi_mask_head(
                model, add_roi_mask_head_func, blob_conv, dim_conv, spatial_scale_conv
            )

        if cfg.MODEL.KEYPOINTS_ON:
            head_loss_gradients['keypoints'] = _add_roi_keypoint_head(
                model, add_roi_keypoint_head_func, blob_conv, dim_conv, spatial_scale_conv
            )

        if model.train:
            loss_gradients = {}
            for lg in head_loss_gradients.values():
                if lg is not None:
                    loss_gradients.update(lg)
            return loss_gradients
        else:
            return None

    optim.build_data_parallel_model(model, _single_gpu_build_func)
    return model

模型头部解析

在模型创建函数build_generic_detection_model()中, 将许多子功能单独作为一个函数实现, 例如限制RPN网络处理的blobs尺寸, 为模型添加各种头部等, 下面我们就对这部分代码做简要分析. 更详细的分析需要根据不同的论文及模型展开讨论, 这些讨论我会在讲解各个模型的时候进行, 因此, 这里我们就只做简单了解, 脑子里有个概念就可以.

先来看看用于控制 FPN 和 RoI 之间的图谱大小的代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def _narrow_to_fpn_roi_levels(blobs, spatial_scales):
    # 只返回用于RoI heads的 blobs 和 spatial scales
    # 输入的`blbs`和`spatial scales`也许会包含额外的用于RPN的部分,
    # 但是这部分并不用于RoI heads, 因此需要去掉

    # 此处代码仅仅支持RPN和RoI的min levels相同的情况
    assert cfg.FPN.RPN_MIN_LEVEL == cfg.FPN.ROI_MIN_LEVEL
    # RPN max level可以 >= RoI max level
    assert cfg.FPN.RPN_MAX_LEVEL >= cfg.FPN.ROI_MAX_LEVEL
    # 因为FPN RPN的max level可能大于FPN RoI max level, 因此, 我们需要这些大于的情况去除
    # blobs 是按照 max/最粗糙 到 min/最精细 的级别排列的
    # cfg.FPN.ROI_MAX_LEVEL 默认为5
    # cfg.FPN.ROI_MIN_LEVEL 默认为2
    num_roi_levels = cfg.FPN.ROI_MAX_LEVEL - cfg.FPN.ROI_MIN_LEVEL + 1
    # 截断, 将前面的较粗糙的blobs和spatial_scales舍去
    return blobs[-num_roi_levels:], spatial_scale_conv[-num_roi_levels:]

添加 Fast R-CNN 的网络头部

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def _add_fast_rcnn_head(
    model, add_roi_box_head_func, blob_in, dim_in, spatial_scale_in
):
    # 调用add_roi_box_head_func函数
    # 该函数位于fast_rcnn_head.py中, 详细的讲解可以查看关于Fast R-CNN的代码解析.
    blob_frcn, dim_frcn = add_roi_box_head_func(
        model, blob_in, dim_in, spatial_scale_in
    )

    # fast_rcnn_heads 为 fast_rcnn_heads文件
    fast_rcnn_heads.add_fast_rcnn_outputs(model, blob_frcn, dim_frcn)

    if model.train:
        loss_gradients = fast_rcnn_heads.add_fast_rcnn_losses(model)
    else:
        loss_gradients = None
    return loss_gradients

添加 Mask R-CNN 头部, 代码逻辑和上面的 Fast R-CNN 差不多.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def _add_roi_mask_head(
    model, add_roi_mask_head_func, blob_in, dim_in, spatial_scale_in
):
    # 在添加 mask head 之前先捕获 model graph
    bbox_net = copy.deepcopy(model.net.Proto())
    # 添加 mask head
    # 调用函数add_roi_mask_head_func
    blob_mask_head, dim_mask_head = add_roi_mask_head_func(
        model, blob_in, dim_in, spatial_scale_in
    )
    # 添加 mask output
    # mask_rcnn_heads为.py文件
    blob_mask = mask_rcnn_heads.add_mask_rcnn_outputs(
        model, blob_mask_head, dim_mask_head
    )

    if not model.train: # == inference
        # 模型推演的时候会用到一系列 box predictions, 然后会预测 mask
        # 这需要将网络的box和mask预测分离开
        # 因此我们将mask预测网络提取出来, 将它作为一个独立的网络存储起来
        # 然后将整个网络重新存储成 bbox-only 网络
        model.mask_net, blob_mask = c2_utils.SuffixNet(
            'mask_net', model.net, len(bbox_net.op), blob_mask
        )
        model.net._net = bbox_net
        loss_gradients = None
    else:
        loss_gradients = mask_rcnn_heads.add_mask_rcnn_losses(model, blob_mask)
    return loss_gradients

添加 keypoint 头部

1
2
3
4
def _add_roi_keypoint_head(
    model, add_roi_keypoint_head_func, blob_in, dim_in, spatial_scale_in
):
    #... 由于我不太关键关键点检测相关任务, 所以这部分的代码不讨论了, 不过整体逻辑和上面是一样的

这部分还有两个函数, 但是根据注释来看, 它们的功能都被build_generic_detection_model()函数代替了, 如下所示:

1
2
3
4
5
6
7
def build_generic_rfcn_model(model, add_conv_body_func, dim_reduce=False):
    # fold this function into build_generic_detection_model
    # ...

def build_generic_retinanet_model(model, add_conv_body_func, freeze_conv_body=False):
    # fold this function into build_generic_detection_model
    # ...

网络模型输入

首先是add_training_inputs函数, 该函数会创建由于网络训练的 input ops 和 blobs.
通常情况下, 我们会同时创建 input ops 和剩下的网络, 但是, 创建 input ops 依赖于加载数据集, 在面对COCO数据集时, 这可能需要好几分钟. 然而, Detectron 的实现希望尽量避免等待, 这样一来在 debug 的时候, 可以更快的进行调试(so debugging can fail fast). 因此, 这里先创建了一个 不带有 input ops 的网络来加载数据集, 然后在加载了数据集以后, 才将 input ops 添加到网络中. 由于 Detectron 推迟了 input ops 的创建, 因此就额外进行一些操作来将 input ops 放置在网络 op list 的最开始.

注意, 该函数必在model_builder.create()被调用之后才能使用.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def add_training_inputs(model, roidb=None):
    assert model.train, 'Training inputs can only be added to a trainable model'

    if roidb is not None:
        # 如果你希望更容易的调试, 则可以将 cfg.DATA_LOADER.NUM_THREADS的值设为1
        model.roi_data_loader = RoIDataLoader(
            roidb,
            num_loaders=cfg.DATA_LOADER.NUM_THREADS,
            minibatch_queue_size=cfg.DATA_LOADER.MINIBATCH_QUEUE_SIZE,
            blobs_queue_capacity=cfg.DATA_LOADER.BLOBS_QUEUE_CAPACITY
        )
    orig_num_op = len(model.net._net.op)
    blob_names = roi_data_minibatch.get_minibatch_blob_names(is_training=True)
    for gpu_id in range(cfg.NUM_GPUS):
        with c2_utils.NamedCudaScope(gpu_id):
            for blob_name in blob_names:
                workspace.CreateBlob(core.ScopedNamed(blob_name))
            model.net.DequeueBlobs(
                model.roi_data_loader._blobs_queue_name, blob_names
            )
    # A little op surgery to move input ops to the start of the net
    diff = len(model.net._net.op) - orig_num_op
    new_op = model.net._net.op[-diff:] + model.net._net.op[:-diff]
    del model.net._net.op[:]
    model.net._net.op.extend(new_op)

下面是创建用于模型推演的网络的 input blobs

1
2
3
4
5
6
7
8
9
10
11
12
def add_inference_inputs(model):

    def create_input_blobs_for_net(net_def):
        for op in net_def.op:
            for blob_in in op.input:
                if not workspace.HasBlob(blob_in):
                    workspace.CreateBlob(blob_in)
    create_input_blobs_for_net(model.net.Proto())
    if cfg.MODEL.MASK_ON:
        create_input_blobs_for_net(model.mask_net.Proto())
    if cfg.MODEL.KEYPOINTS_ON:
        create_input_blobs_for_net(model.keypoint_net.Proto())

关于数据载入部分代码大多都是调用./detectron/roi_data/loader.py和./detectron/roi_data/minibatch.py文件中的类和函数实现的, 这里我们仅仅当做接口来使用, 如果你想要了解数据加载的具体实现方法, 可以查看这篇解析