Detectron 源码解析-检测模型辅助器(DetectionModelHelper)

DetectionModelHelper 类概览

./detectron/modeling/model_builder.py文件中的, Detectron 使用了 create() 函数来创建目标检测模型, 其中最主要的部分是使用了 detectron/modeling/detector.py文件中的 class DetectionModelHelper()类, 该类是 Detectron 所有类型模型的最基本结构, 下面, 我们就对该类进行解析. 首先, 还是老样子, 我们大致浏览一下这个文件中的类和函数, 以及各个函数的功能, 把握类的设计思路和整体结构.

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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
# ./detectron/modeling/detector.py

# 导入各种包及函数
from detectron.core.config import cfg
# ...

# 该类继承了cnn.CNNModelHelper, 用于表示各种检测模型. 最新的 caffe2 建议用 ModelHelper + brew 来替换 CNNModelHelper
class DetectionModelHelper(cnn.CNNModelHelper):
def __init__(self, **kwargs):
# 初始化函数, 处理各种配置选项
# ...

def TrainableParams(self, gpu_id=-1):
# 获取所有可训练参数的 blob 名字
# ..

def AffineChannel(self, blob_in, blob_out, dim, inplace=False):
# 当 BN 不能使用时(minibatch size 太小), 用仿射变换来替代BN
# ...

def GenerateProposals(self, blobs_in, blobs_out, anchors, spatial_scale):
# 用于生成 RPN 候选区域框的 Op
# ...

def GenerateProposalLabels(self, blobs_in):
# 用于生成 RPN 候选区域标签的 Op. 该函数只会在 e2e Faster/Mask 训练时才会被调用
# ...

def CollectAndDistributeFpnRpnProposals(self):
# 将多个 FPN levels 产生的 RPN proposals 融合, 然后再将这些候选区域分配到其自身适应的 FPN levels
# ...

def DropoutIfTraining(self, blob_in, dropout_rate):
# 添加 dropout


def RoIFeatureTransform(self, blobs_in, blob_out, blob_rois='rois', method='RoIPoolF', resolution=7, spatial_scale=1. / 16., sampling_ratio=0):
# 添加具体的 RoI pooling 方法.
# ...

def ConvShared(self, blob_in, blob_out, dim_in, dim_out, kernel, weight=None, bias=None, **kwargs):
# 添加与另一个 conv op 共享权重的 conv op
# ...

def BilinearInterpolation(self, blob_in, blob_out, dim_in, dim_out, up_scale):
# 双线性插值
# ...

def ConvAffine( # args in the same order of Conv()
self, blob_in, prefix, dim_in, dim_out, kernel, stride, pad,
group=1, dilation=1,
weight_init=None,
bias_init=None,
suffix='_bn',
inplace=False
):
# 在 AffineChannel Op 之后添加 Conv op
# ...

def ConvGN( # args in the same order of Conv()
self, blob_in, prefix, dim_in, dim_out, kernel, stride, pad,
group_gn, # gn 中 groups 的数量
group=1, dilation=1,
weight_init=None,
bias_init=None,
suffix='_gn',
no_conv_bias=1,
):
# 在 GroupNorm op 之后添加 Conv op
# ...

def DisableCudnn(self):
# ...

def RestorePreviousUseCudnn(self):
# ...

def UpdateWorkspaceLr(self, cur_iter, new_lr):
# 更新当前模型的学习率和workspace
# ...

def _SetNewLr(self, cur_lr, new_lr):
# 更新模型和 workspace blobs
# ...

def _CorrectMomentum(self, correction):
# MomentumSGDUpdate op 实现
# ...

def GetLossScale(self):
# 动态开支 loss 大小
# ...

def AddLosses(self, losses):
# ...

def AddMetrics(self, metrics):
# ...


def _get_lr_change_ratio(cur_lr, new_lr):
# ...

导入的包和函数

该文件导入的包和函数, 及其它们的功能如下所示:

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

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

import numpy as np
import logging

# caffe2 包
from caffe2.python import cnn # cnn.CNNModelHelper
from caffe2.python import core
from caffe2.python import workspace
from caffe2.python.modeling import initializers
from caffe2.python.modeling.parameter_info import ParameterTags

# detectron 包及函数
from detectron.core.config import cfg
from detectron.ops.collect_and_distribute_fpn_rpn_proposals \
import CollectAndDistributeFpnRpnProposalsOp
from detectron.ops.generate_proposal_labels import GenerateProposalLabelsOp
from detectron.ops.generate_proposals import GenerateProposalsOp
import detectron.roi_data.fast_rcnn as fast_rcnn_roi_data
import detectron.utils.c2 as c2_utils

初始化函数

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

# from caffe2.python import cnn
class DetectionModelHelper(cnn.CNNModelHelper):
# 该类是cnn.CNNModelHelper的一个子类, 这一点很重要
def __init__(self, **kwargs):
# 处理属于DetectionModelHelper的特定的参数, 其他的都会传到CNNModelHelper当中
self.train = kwargs.get('train', False) # train
self.num_classes = kwargs.get('num_classes', -1) # cfg.MODEL.NUM_CLASSES
assert self.num_classes > 0, 'num_classes must be > 0'
for k in ('train', 'num_classes'):
if k in kwargs:
del kwargs[k]# TODO 貌似是因为下面要再次调用构造函数的原因?
# 设置数据的排列顺序: batchsize, channels, heiht, width
kwargs['order'] = 'NCHW'

# TODO, 不懂这个选项的实际作用
kwargs['cudnn_exhaustive_search'] = False
super(DetectionModelHelper, self).__init__(**kwargs) # TODO ?? 这样不会造成递归调用吗? 为什么要递归调用构造函数
self.roi_data_loader = None
self.losses = []
self.metrics = []
self.do_not_update_params = [] # 位于该列表中的参数不会被更新
self.net.Proto().type = cfg.MODEL.EXECUTION_TYPE
self.net.Proto().num_workers = cfg.NUM_GPUS * 4
self.prev_use_cudnn = self.use_cudnn
self.gn_params = [] # 位于该列表中的元素是GroupNorm参数

损失函数相关

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# ./detectron/modeling/detector.py
class DetectionModelHelper(cnn.CNNModelHelper):
#...

def GetLossScale(self):
# 当多 GPU 分布式训练时起作用, 根据 GPU 的数量调整损失的大小
return 1.0 / cfg.NUM_GPUS

def AddLosses(self, losses):
# losses 是一个列表, eg: ['loss_cls', 'loss_bbox']
if not isinstance(losses, list): # 如果不是列表, 则将其转换成列表
losses = [losses]
# 对字符串进行转换, 使得 losses 包含相应blob的引用(将 'gpu_0/foo' 转换成 'foo')
losses = [c2_utils.UnscopeName(str(l)) for l in losses]
self.losses = list(set(self.losses + losses))

def AddMetrics(self, metrics):
if not isinstance(metrics, list):
metrics = [metrics]
self.metrics = list(set(self.metrics + metrics))