【代碼研讀】Mask RCNN代碼閱讀筆記(二)骨架網絡backbone和入口

前言

筆記分爲三個部分,backbone,rpn,roi_head三個部分,之前的該項目總覽見【鏈接】。本文主要是講解backbone部分的文件,通過在總體把握和實現細節兩個方面對其進行記錄。

檢測的入口

│   │   ├── detector
│   │   │   ├── detectors.py  #檢測的代碼入口
│   │   │   ├── generalized_rcnn.py #生成各種組合的檢測模型

進行檢測的過程中,首先的代碼入口爲detectors.py。
detectors.py作用: 根據給定的配置信息實例化一個 class GeneralizedRCNN 的對象。GeneralizedRCNN的構建在文件 generalized_rcnn.py中

generalized_rcnn.py作用:定義了 MaskrcnnBenchmark 的 GeneralizedRCNN 類, 用於表示各種組合後的目標檢測模型

  • 該類是 MaskrcnnBenchmark 中所有模型的共同抽象, 目前支持 boxes 和 masks 兩種形式的標籤
  • 該類主要包含以下三個部分:
    • backbone
    • rpn(option)
    • heads: 利用前面網絡輸出的 features 和 proposals 來計算 detections / masks.

骨架網絡總體:

│   │   ├── backbone 
│   │   │   ├── backbone.py
│   │   │   ├── fpn.py
│   │   │   └── resnet.py

由總體概覽知,骨架網絡主要分爲backbone部分,resnet部分,以及由resnet組成的rpn部分三個主要部分。下面分部分來進行筆記。

backbone

主要作用:將resent和rpn進行調用,並創建對應的組合特徵提取網絡。
backbone.py 文件中的兩個函數 build_resnet_backbone() 和 build_resnet_fpn_backbone() 都使用了 body = resnet.ResNet(cfg) 來創建網絡的主體, 這部分的代碼定義位於 ./maskrcnn_benchmark/modeling/backbone/resnet.py 文件

resnet

這裏實現了50,101層的resnet的代碼。

在這裏插入圖片描述

  1. stem作爲stage1被進行創建
  2. resnet2~5 階段的整體結構是非常相似的, 都是有最基礎的 resnet bottleneck block 堆疊形成的。

在maskrcnn benchmark中,對上面提到的這兩種block結構進行的衍生和封裝,Bottleneck和Stem分別衍生出帶有Batch Normalization 和 Group Normalizetion的封裝類,分別爲:

  • BottleneckWithFixedBatchNorm,
  • StemWithFixedBatchNorm
  • BottleneckWithGN,
  • StemWithGN.

搭建模塊層級的核心代碼

def _make_stage(
    transformation_module,
    in_channels,
    bottleneck_channels,
    out_channels,
    block_count,
    num_groups,
    stride_in_1x1,
    first_stride,
    dilation=1,
    dcn_config={}
):
    blocks = []
    stride = first_stride
    # 根據不同的配置,構造不同的卷基層
    for _ in range(block_count):
        blocks.append(
            transformation_module(
                in_channels,
                bottleneck_channels,
                out_channels,
                num_groups,
                stride_in_1x1,
                stride,
                dilation=dilation,
                dcn_config=dcn_config
            )
        )
        stride = 1
        in_channels = out_channels
    return nn.Sequential(*blocks)

resnet核心實現代碼

# ./maskrcnn_benchmark/modeling/backbone/resnet.py

class ResNet(nn.Module):
    def __init__(self, cfg):
        super(ResNet, self).__init__()

        # 如果我們希望在 forward 函數中使用 cfg, 那麼我們就應該創建一個副本以供其使用
        # self.cfg = cfg.clone()

        # 將配置文件中的字符串轉化成具體的實現, 下面三個分別使用了對應的註冊模塊, 定義在文件的最後

        # 這裏是 stem 的實現, 也就是 resnet 的第一階段 conv1
        # cfg.MODEL.RESNETS.STEM_FUNC = "StemWithFixedBatchNorm"
        stem_module = _STEM_MODULES[cfg.MODEL.RESNETS.STEM_FUNC]

        # resnet conv2_x~conv5_x 的實現
        # eg: cfg.MODEL.CONV_BODY="R-50-FPN"
        stage_specs = _STAGE_SPECS[cfg.MODEL.CONV_BODY]

        # residual transformation function
        # cfg.MODEL.RESNETS.TRANS_FUNC="BottleneckWithFixedBatchNorm"
        transformation_module = _TRANSFORMATION_MODULES[cfg.MODEL.RESNETS.TRANS_FUNC]

        # 獲取上面各個組成部分的實現以後, 就可以利用這些實現來構建模型了

        # 構建 stem module(也就是 resnet 的stage1, 或者 conv1)
        self.stem = stem_module(cfg)

        # 獲取相應的信息來構建 resnet 的其他 stages 的卷積層

        # 當 num_groups=1 時爲 ResNet, >1 時 爲 ResNeXt
        num_groups = cfg.MODEL.RESNETS.NUM_GROUPS

        #
        width_per_group = cfg.MODEL.RESNETS.WIDTH_PER_GROUP

        # in_channels 指的是向後面的第二階段輸入時特徵圖譜的通道數,
        # 也就是 stem 的輸出通道數, 默認爲 64
        in_channels = cfg.MODEL.RESNETS.STEM_OUT_CHANNELS

        # 第二階段輸入的特別圖譜的通道數
        stage2_bottleneck_channels = num_groups * width_per_group

        # 第二階段的輸出, resnet 系列標準模型可從 resnet 第二階段的輸出通道數判斷後續的通道數
        # 默認爲256, 則後續分別爲512, 1024, 2048, 若爲64, 則後續分別爲128, 256, 512
        stage2_out_channels = cfg.MODEL.RESNETS.RES2_OUT_CHANNELS

        # 創建一個空的 stages 列表和對應的特徵圖譜字典
        self.stages = []
        self.return_features = {}

        for stage_spec in stage_specs: # 關於 stage_specs 的定義可以看上一節
            name = "layer" + str(stage_spec.index)

            # 計算每個stage的輸出通道數, 每經過一個stage, 通道數都會加倍
            stage2_relative_factor = 2 ** (stage_spec.index - 1)

            # 計算輸入圖譜的通道數
            bottleneck_channels = stage2_bottleneck_channels * stage2_relative_factor

            # 計算輸出圖譜的通道數
            out_channels = stage2_out_channels * stage2_relative_factor

            # 當獲取到所有需要的參數以後, 調用本文件的 `_make_stage` 函數,
            # 該函數可以根據傳入的參數創建對應 stage 的模塊(注意是module而不是model)
            module = _make_stage(
                transformation_module,
                in_channels, # 輸入的通道數
                bottleneck_channels, # 壓縮後的通道數
                out_channels, # 輸出的通道數
                stage_spec.block_count, #當前stage的卷積層數量
                num_groups, # ResNet時爲1, ResNeXt時>1
                cfg.MODEL.RESNETS.STRIDE_IN_1X1,
                # 當處於 stage3~5時, 需要在開始的時候使用 stride=2 來downsize
                first_stride=int(stage_spec.index > 1) + 1,
            )

            # 下一個 stage 的輸入通道數即爲當前 stage 的輸出通道數
            in_channels = out_channels

            # 將當前stage模塊添加到模型中
            self.add_module(name, module)

            # 將stage的名稱添加到列表中
            self.stages.append(name)

            # 將stage的布爾值添加到字典中
            self.return_features[name] = stage_spec.return_features

        # 根據配置文件的參數選擇性的凍結某些層(requires_grad=False)
        self._freeze_backbone(cfg.MODEL.BACKBONE.FREEZE_CONV_BODY_AT)

    def _freeze_backbone(self, freeze_at):
        # 根據給定的參數凍結某些層的參數更新
        for stage_index in range(freeze_at):
            if stage_index == 0:
                m = self.stem # resnet 的第一階段, 即爲 stem
            else:
                m = getattr(self, "layer" + str(stage_index))
            # 將 m 中的所有參數置爲不更新狀態.
            for p in m.parameters():
                p.requires_grad = False

    # 定義 ResNet 的前行傳播過程
    def forward(self, x):
        outputs = []
        x = self.stem(x) # 先經過 stem(stage 1)

        # 再依次計算 stage2~5的結果
        for stage_name in self.stages:
            x = getattr(self, stage_name)(x)
            if self.return_features[stage_name]:
                # 將stage2~5的所有計算結果(也就是特徵圖譜)以列表形式保存
                outputs.append(x)

        # 將結果返回, outputs爲列表形式, 元素爲各個stage的特徵圖譜, 剛好作爲 FPN 的輸入
        return outputs

FPN

我們將通過resnet搭建對應的FPN網路模型,FPN網絡主要應用於多層特徵提取,使用多尺度的特徵層來進行目標檢測,可以利用不同的特徵層對於不同大小特徵的敏感度不同,將他們充分利用起來,以更有利於目標檢測
這裏直接利用了上面resnet的類,並在此基礎上進行搭建,主要流程爲定義11卷積改變通道數,定義33卷積進行特徵圖提取。在進行前向計算階段,進行縮小2倍的特徵並進行直接的相加。將這個作爲下一個階段的輸入。


resnet+FPN的實例圖

核心實現代碼

# ./maskrcnn_benchmark/modeling/backbone/fpn.py

import torch
import torch.nn.functional as F
from torch import nn

class FPN(nn.Module):
    # 在一系列的 feature map (實際上就是stage2~5的最後一層輸出)添加 FPN
    # 這些 feature maps 的 depth 假定是不斷遞增的, 並且 feature maps 必須是連續的(從stage角度)

    def __init__(self, in_channels_list, out_channels, top_blocks=None):
        # in_channels_list (list[int]): 指示了送入 fpn 的每個 feature map 的通道數
        # out_channels (int): FPN表徵的通道數, 所有的特徵圖譜最終都會轉換成這個通道數大小
        # top_blocks (nn.Module or None): 當提供了 top_blocks 時, 就會在 FPN 的最後
        # 的輸出上進行一個額外的 op, 然後 result 會擴展成 result list 返回
        super(FPN, self).__init__()

        # 創建兩個空列表
        self.inner_blocks = []
        self.layer_blocks = []

        # 假設我們使用的是 ResNet-50-FPN 和配置, 則 in_channels_list 的值爲:
        # [256, 512, 1024, 2048]
        for idx, in_channels in enumerate(in_channels_list, 1): # 下標從1開始
            # 用下表起名: fpn_inner1, fpn_inner2, fpn_inner3, fpn_inner4
            inner_block = "fpn_inner{}".format(idx)

            # fpn_layer1, fpn_layer2, fpn_layer3, fpn_layer4
            layer_block = "fpn_layer{}".format(idx)

            # 創建 inner_block 模塊, 這裏 in_channels 爲各個stage輸出的通道數
            # out_channels 爲 256, 定義在用戶配置文件中
            # 這裏的卷積核大小爲1, 該卷積層主要作用爲改變通道數到 out_channels(降維)
            inner_block_module = nn.Conv2d(in_channels, out_channels, 1)

            # 改變 channels 後, 在每一個 stage 的特徵圖譜上再進行 3×3 的卷積計算, 通道數不變
            layer_block_module = nn.Conv2d(out_channels, out_channels, 3, 1, 1)

            for module in [inner_block_module, layer_block_module]:
                # Caffe2 的實現使用了 XavierFill,
                # 實際上相當於 PyTorch 中的 kaiming_uniform_
                nn.init.kaiming_uniform_(module.weight, a=1)
                nn.init.constant_(module.bias, 0)

            # 在當前的特徵圖譜上添加 FPN
            self.add_module(inner_block, inner_block_module) #name, module
            self.add_module(layer_block, layer_block_module)

            # 將當前 stage 的 fpn 模塊的名字添加到對應的列表當中
            self.inner_blocks.append(inner_block)
            self.layer_blocks.append(layer_block)

        # 將top_blocks作爲 FPN 類的成員變量
        self.top_blocks = top_blocks

    def forward(self, x):
        # x (list[Tensor]): 每個 feature level 的 feature maps,
        # ResNet的計算結果正好滿足 FPN 的輸入要求, 也因此可以使用 nn.Sequential 將二者直接結合
        # results (tuple[Tensor]): 經過FPN後的特徵圖譜組成的列表, 排列順序是高分辨率的在前

        # 先計算最後一層(分辨率最低)特徵圖譜的fpn結果.
        last_inner = getattr(self, self.inner_blocks[-1])(x[-1])

        # 創建一個空的結果列表
        results=[]

        # 將最後一層的計算結果添加到 results 中
        results.append(getattr(self, self.layer_blocks[-1])(last_inner))


        # [:-1] 獲取了前三項, [::-1] 代表從頭到尾切片, 步長爲-1, 效果爲列表逆置
        # 舉例來說, zip裏的操作 self.inner_block[:-1][::-1] 的運行結果爲
        # [fpn_inner3, fpn_inner2, fpn_inner1], 相當於對列表進行了逆置
        for feature, inner_block, layer_block in zip(
            x[:-1][::-1], self.inner_block[:-1][::-1], self.layer_blocks[:-1][::-1]
        ):
            # 根據給定的scale參數對特徵圖譜進行放大/縮小, 這裏scale=2, 所以是放大
            inner_top_down = F.interpolate(last_inner, scale_factor=2, mode="nearest")

            # 獲取 inner_block 的計算結果
            inner_lateral = getattr(self, inner_block)(feature)

            # 將二者疊加, 作爲當前stage的輸出 同時作爲下一個stage的輸入
            last_inner = inner_lateral + inner_top_down

            # 將當前stage輸出添加到結果列表中, 注意還要用 layer_block 執行卷積計算
            # 同時爲了使得分辨率最大的在前, 我們需要將結果插入到0位置
            results.insert(0, getattr(self, layer_block)(last_inner))

        # 如果 top_blocks 不爲空, 則執行這些額外op
        if self.top_blocks is not None:
            last_results = self.top_blocks(results[-1])
            results.extend(last_results) # 將新計算的結果追加到列表中

        # 以元組(只讀)形式返回
        return tuple(results)

# 最後一級的 max pool 層
class LastLevelMaxPool(nn.Module):
    def forward(self, x):
        return [F.max_pool2d(x, 1, 2, 0)]

參考鏈接:
【地址1 resnet.py】
【地址2 rpn.py】
【地址3 backbone.py】

【地址2個人主頁】
(文中的ROI_HEAD部分並沒有進行註解)

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章