Caffe2 - (三十一) Detectron 之 modeling - FPN 与 optimizer

Caffe2 - (三十一) Detectron 之 modeling - FPN 与 optimizer

1. FPN.py

FPN 模块.

"""
Feature Pyramid Network (FPN) 使用的相关函数.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import collections
import numpy as np

from core.config import cfg
from modeling.generate_anchors import generate_anchors ## 
from utils.c2 import const_fill
from utils.c2 import gauss_fill
import modeling.ResNet as ResNet ## 
import utils.blob as blob_utils
import utils.boxes as box_utils

"""
Backbone 骨干网络的最低(lowest) 和最高(highest) 金字塔(pyramid) 层(level).
对于 FPN， 这里假设所有的网络都有 5 个空间维度的减少(spatial reductions),
每一个减少的因子都是 2.
Level 1 对应于输入图片，此时使用时是没有意义的.
"""
LOWEST_BACKBONE_LVL = 2   # E.g., "conv2"-like level
HIGHEST_BACKBONE_LVL = 5  # E.g., "conv5"-like level


# ---------------------------------------------------------------------------- #
# FPN with ResNet
# ---------------------------------------------------------------------------- #

def add_fpn_ResNet50_conv5_body(model):
    return add_fpn_onto_conv_body(model, 
                                  ResNet.add_ResNet50_conv5_body, 
                                  fpn_level_info_ResNet50_conv5 )


def add_fpn_ResNet50_conv5_P2only_body(model):
    return add_fpn_onto_conv_body(model,
                                  ResNet.add_ResNet50_conv5_body,
                                  fpn_level_info_ResNet50_conv5,
                                  P2only=True )


def add_fpn_ResNet101_conv5_body(model):
    return add_fpn_onto_conv_body(model, 
                                  ResNet.add_ResNet101_conv5_body, 
                                  fpn_level_info_ResNet101_conv5)


def add_fpn_ResNet101_conv5_P2only_body(model):
    return add_fpn_onto_conv_body(model, 
                                  ResNet.add_ResNet101_conv5_body,
                                  fpn_level_info_ResNet101_conv5,
                                  P2only=True )


def add_fpn_ResNet152_conv5_body(model):
    return add_fpn_onto_conv_body(model, 
                                  ResNet.add_ResNet152_conv5_body, 
                                  fpn_level_info_ResNet152_conv5 )


def add_fpn_ResNet152_conv5_P2only_body(model):
    return add_fpn_onto_conv_body(model, 
                                  ResNet.add_ResNet152_conv5_body,
                                  fpn_level_info_ResNet152_conv5,
                                  P2only=True )


# ---------------------------------------------------------------------------- #
# 添加 FPN 到骨干backbone 网络结构的相关函数
# ---------------------------------------------------------------------------- #

def add_fpn_onto_conv_body(model, conv_body_func, fpn_level_info_func, P2only=False):
    """
    添加指定的 conv body 到模型model， 然后再往模型添加 FPN levels.
    """
    """
    注意事项：
        blobs_conv 是倒序方式：[fpn5, fpn4, fpn3, fpn2]
        类似地，dims_conv：[2048, 1024, 512, 256]
        类似地，spatial_scales_fpn：[1/32, 1/16, 1/8, 1/4]
    """

    conv_body_func(model)
    blobs_fpn, dim_fpn, spatial_scales_fpn = add_fpn(model, fpn_level_info_func())

    if P2only:
        # 只使用最精细层次, finest level
        return blobs_fpn[-1], dim_fpn, spatial_scales_fpn[-1]
    else:
        # 使用所有的层次 all levels
        return blobs_fpn, dim_fpn, spatial_scales_fpn


def add_fpn(model, fpn_level_info):
    """
    基于 FPN 论文叙述的模型，添加 FPN 连接.
    """
    """
    FPN levels 是从骨干backbone 网络的 highest/coarest level(通常为 conv5) 开始构建的.
    首先向下，递归地(recursively)构建 lower/finer 分辨率的 FPN levels；
    然后向上，构建比起始 level higher/coarser 分辨率的 FPN levels.
    """
    fpn_dim = cfg.FPN.DIM
    min_level, max_level = get_min_max_levels()

    """
    计算从 coarest backbone 阶段stage(通常是 "conv5"-like level) 生成 FPN levels时，
    backbone 阶段stages 数.
    如，如果 backbone levels info 定义了 4 stages："conv5", "conv4", "conv3", "conv2"，且 min_level=2，
    那么，将会添加 FPN 上的 backbone stages 数为：4 - (2 - 2) = 4
    """
    num_backbone_stages = (len(fpn_level_info.blobs) - (min_level - LOWEST_BACKBONE_LVL) )

    lateral_input_blobs = fpn_level_info.blobs[:num_backbone_stages]
    output_blobs = ['fpn_inner_{}'.format(s) for s in fpn_level_info.blobs[:num_backbone_stages] ]
    fpn_dim_lateral = fpn_level_info.dims
    xavier_fill = ('XavierFill', {})

    # 对于 coarest backbone level: 1x1 conv only seeds recursion
    model.Conv(lateral_input_blobs[0],
               output_blobs[0],
               dim_in=fpn_dim_lateral[0],
               dim_out=fpn_dim,
               kernel=1,
               pad=0,
               stride=1,
               weight_init=xavier_fill,
               bias_init=const_fill(0.0) )

    #
    # Step 1: 从 coarest backbone level 开始，递归地向下构建 FPN levels
    #

    # 对于其它 levels，添加 top-down 和侧向连接(lateral connections)
    for i in range(num_backbone_stages - 1):
        add_topdown_lateral_module(
            model,
            output_blobs[i],             # top-down blob
            lateral_input_blobs[i + 1],  # lateral blob 侧向 blob
            output_blobs[i + 1],         # next output blob 下一个输出 blob
            fpn_dim,                     # output dimension 输出维度
            fpn_dim_lateral[i + 1]       # lateral input dimension 侧向输入维度
        )

    # Post-hoc scale-specific 3x3 convs
    blobs_fpn