Detectron源码解析-Fast R-CNN Heads

源码文件

./detectron/modeling/fast_rcnn_heads.py

Fast R-CNN Heads

本文件定义了 Fast R-CNN Heads 的结构, 即 backbone 网络之上的两个并列的预测层: 分类层和回归层. Fast R-CNN 的设计如下:

1
2
3
4
... -> RoI ----\                               /-> box cls output -> cls loss
-> RoIFeatureXform -> box head
... -> Feature / \-> box reg output -> reg loss
Map

首先, Fast R-CNN 会利用 RoI pooling 生成固定长度的 RoI 的特征表示, 然后利用此特征进行物体的类别预测和边框回归预测.

fast_rcnn_heads.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
# ./detectron/modeling/fast_rcnn_head.py

# ...
from detectron.core.config import cfg
# ...


# ----------------------------------- #
# Fast R-CNN outputs and losses
# ----------------------------------- #

def add_fast_rcnn_outputs(model, blob_in, dim):
# 添加 RoI 分类器和 bounding box 回归器的输出
# ...

def add_fast_rcnn_losses(model):
# 添加损失函数
# ...


# ----------------------------------- #
# Box heads
# ----------------------------------- #
def add_roi_2mlp_head(model, blob_in, dim_in, spatial_scale):
# 添加具有两层隐藏层的 ReLU MLP
# ...

def add_roi_Xconv1fc_head(model, blob_in, dim_in, spatial_scale):
# 添加一层 X conv + 一层 fc 作为 head
# ...

def add_roi_Xconv1fc_gn_head(model, blob_in, dim_in, spatial_scale):
# 添加一层 X conv + 一层 fc 作为 head, 同时使用 GN
# ...

导入的包及函数

老样子, 我们先看一下这个文件使用了哪些包和函数

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

# 常规包
from __future__ import absolute_import
from __future__ import division
from __future__ import print_functioin
from __future__ import unicode_literals

from detectron.core.config import cfg
from detectron.utils.c2 import const_fill
from detectron.utils.c2 import gauss_fill
from detectron.utils.net import get_group_on
import detectron.utils.blob as blob_utils

Fast R-CNN outputs and losses

add_fast_rcnn_outputs()会在./detectron/modeling/model_builder.py中的_add_fast_rcnn_head()函数中调用, 如果此时模型处于训练状态, 还会同时调用add_fast_rcnn_losses()函数. 下面我们就先来看一下这两个函数的内部实现. 首先是add_fast_rcnn_outputs()函数.

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
def add_fast_rcnn_outputs(model, blob_in, dim):
# 依据 Fast R-CNN 的设计, 创建 box 分类层
model.FC(
blob_in,
'cls_score', # blob_out, 以名称存在于 workspace 当中
dim, # dim_in
model.num_classes, # dim_out
weight_init=gauss_fill(0.01), # 该函数来自于 detectron.utils.c2 文件
bias_init=const_fill(0.0)
)
if not model.train: # == test
# 在推演的时候, 仅仅添加softmax
# 在训练的时候, 需要将softmax和交叉熵结合
model.Softmax('cls_score', 'cls_prob', engine='CUDNN')
# Box 回归层
num_bbox_reg_classes=( # 回归层的类别数, 当只要求分成前景和后景时, 为2, 当要求对分成对应类别时, 为 num_classes
2 if cfg.MODEL.CLS_AGNOSTIC_BBOX_REG else model.num_classes
)
model.FC(
blob_in,
'bbox_pred',
dim, # dim_in
num_bbox_reg_classes*4, # dim_out
weight_init=gauss_fill(0.001),
bias_init=const_fill(0.0)
)

可以看到, 上面的函数是根据 Fast R-CNN 中两个并列的预测层进行定义的, 这两个预测层分别为物体类别置信度的分类预测层, 以及物体边框坐标点的回归预测层, 其中回归预测层在 Fast R-CNN 中的类别默认值为2, 即类别不可知(class-agnostic), 只分前景和背景, 并且预测出前景框和背景框的坐标.

下面是 Fast R-CNN 的损失定义

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
def add_fast_rcnn_losses(model):
# 添加 RoI 分类损失
cls_prob, loss_cls = model.net.SoftmaxWithLoss(
['cls_score', 'labels_int32'], ['cls_prob', 'loss_cls'],
scale=model.GetLossScale()
)
# 添加 bounding box 回归损失
loss_bbox = model.net.SmoothL1Loss(
[
'bbox_pred', 'bbox_targets', 'bbox_inside_weights',
'bbox_outside_weights'
],
'loss_bbox',
scale=model.GetLossScale()
)

# 求梯度
loss_gradients = blob_utils.get_loss_gradients(model, [loss_cls, loss_bbox])