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

DetectionModelHelper 类概览

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

# ./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):
    # ...

导入的包和函数

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

# ./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

初始化函数

# 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参数

损失函数相关

# ./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))