Mask RCNN网络结构以及整体流程的详细解读

1、概述

• Mask RCNN是在Faster RCNN的基础上增加了mask head用于实例分割的模型。
• 总体来说，Mask RCNN网络结构可以分为：BackBone(ResNet+FPN) —> RPN网络(Region Proposal Network) —> ROI Head(ROIAlign + cls head + bbox head + mask head)
• 整体网络结构如下（来自原论文https://arxiv.org/pdf/1703.06870.pdf）：
  

2、Backbone

• Backbone主要由ResNet和FPN组成，如下图
  

1）C1-C5是指ResNet网络5个stage输出的特征图，下采样率分别为2, 4, 8, 16, 32，通道数分别为64, 256, 512, 1024, 2048

2）FPN网络类似于UNet网络中的跳级连接，都是为了增强多尺度表征能力的：分别对C2-C5四个特征图通过1 $\times$ 1卷积改变通道数为256，再与上采样路径中对应尺度的特征图相加。得到的P2-P6就是FPN网络的五个特征图，下采样率为4, 8, 16, 32, 64，通道数都为256
3）需要注意的是，P2-P6这五个特征图将分别作为RPN网络的输入和ROI Head中的ROI Align的输入：

RPN网络的输入：用于在这五个特征图上生成先验的anchor，并对这些anchor进行类别预测和bbox回归，以生成最终的anchor输入到ROI Head中
ROI Align的输入：根据RPN网络生成的anchors，提取这五个特征图中对应的ROI区域特征图，输入到cls、bbox和maskhead中

3、RPN网络

• RPN网络为Region Proposal Network，主要作用：用于生成先验的anchor box/proposals；并将P2-P6这五个特征图作为网络输入学习得到这些先验的anchor box/proposal的类别(前景or背景)以及通过bbox 回归得到偏移量；最终经过RPN网络预测得到的具有更高质量的proposal送入ROI Head中
• 主要流程为：1）anchor的生成；2）anchor的标注；3）P2-P6五个特征图进行二分类的预测和bbox回归；4）通过NMS后处理得到最终的anchor

3.1、anchor的生成

• 在P2-P6的五个特征图上分别对应设置5个不同的anchor size(32, 64, 128, 256, 512)并设置3种长宽比(0.5, 1.0, 2.0)，也就是每个特征图的每个像素点生成3个anchor(x, y, w, h)
• 具体来说，特征图中的每个点都会先映射到原始图像中，并以该点为中心，以对应的anchor size和3种长宽比的设置，生成3个anchor(x, y, w, h)
• 例如，输入图像为512$\times$ 512，那么五个特征图的尺寸分别为128, 64, 32, 16, 8，那么生成的anchors的数量为(128$\times$ 128 + 64$\times$ 64 + 32$\times$ 32 + 16$\times$ 16 + 8$\times$ 8)$\times$ 3 = 21824$\times$ 3 = 65472个anchors

3.2、anchor的标注/分配

• 该过程就是确定生成的所有anchor为正样本or负样本
• 标注流程如下：

1）先排除掉超过原图边界的anchors；
2）计算其余的每个anchor与 所有ground truth bbox的IoU，取与每个bbox的的最大IoU作为判断：大于0.7为正样本，小于0.3为负样本；
3）计算每个bbox与那个anchor的IoU最大，把该anchor也标记为正样本

• 最后随机选择128个正样本和128个负样本用于RPN网络的训练，对于分类：128个正样本的label为1,128个负样本的label和剩下的anchor的label为0；对于bbox回归：超过边界的anchor的label为(0, 0, 0, 0)，其余anchor的label是与它对应的具有最大IoU的bbox的实际偏移量

3.3、分类预测和bbox回归

• 首先对5个特征图进行3$\times$ 3卷积，然后分为2个分支
• 分类分支为1$\times$ 1卷积，输出通道数为num_anchors=3，表示的是特征图的每个像素点对应的3个anchors为前景的概率值。

需要注意的是，五个特征图分别作为3$\times$ 3卷积和该分支的1$\times$ 1卷积的输入计算得到概率值，再计算128个正样本和128个负样本的分类损失，也就是说只有一组3$\times$ 3卷积和1$\times$ 1卷积

• bbox回归分支为1$\times$ 1卷积，输出通道数为num_anchors$\times$ 4 = 12，表示3个anchor的4个预测位置参数

以下为bbox回归的详细介绍，预测得到的(dx, dy, dw, dh)与ground truth之间计算Smooth L1损失

• 训练RPN网络：对128个正样本和128个负样本计算分类损失：二分类交叉熵损失函数，以及回归损失：Smooth L1损失

3.4、NMS生成最终的anchor

• 根据RPN网络生成的所有anchor的score和4个位置参数(tx, ty, tw, th)进行如下操作：

1）根据偏移量对anchor的位置进行微调，得到(x, y, w, h)
2）删除掉超过边界的anchor
3）根据score对anchor从大到小进行排序
4）对于每个尺度上的特征图（就是P2-P6的五个特征图），选择前2000个anchor，在经过NMS之后选择最高1000个anchors作为最终的anchor输入到ROI Head中(这些数量属于超参数，可调整)

• 例如，5个特征图在NMS之前选择的前2000个anchor如下，3和4没到2000是因为他们的特征图P5和P6的尺寸分别为16和8，计算得到的anchors数量为768和192。
  

• 训练阶段和推理阶段有一些不同：

训练阶段：计算loss并生成anchor
推理阶段：直接生成anchor

• 生成anchor的NMS的配置可能不同，例如：

训练阶段

推理阶段

4、ROI Head

• 主要包括ROI ALign和3个head: class head, bbox head, mask head

4.1、ROI Align

• RPN网络最终生成的anchors将作为ROI Align的输入，也就是说根据这些anchors(x, y, w, h) 来提取P2-P6特征图中的anchors对应的特征作为3个head的特征图输入。
• 我们首先要知道anchors的坐标位置(x, y, w, h)表示的是原图像中的区域，可以将其映射到特征图中，Faster RCNN正是采用了ROI pooling层映射为固定尺寸的ROI区域：
  
• 但是ROI pooling层由于存在取整操作，对于一些小目标很容易导致区域不匹配问题，因此Mask RCNN采用了ROI ALign操作

1）将anchor映射到对应的特征图中，w和h在被除时不去整，ROI pooling取整了。对于五个特征图的选择方法如下：

2）将映射后的ROI区域均匀分为K$\times$K的bin，每个bin的大小不取整，ROI pooling取整了
3）每个bin的每个像素值计算为特征图中与它邻近的4个值的双线性插值

4）使用max pooling或者average pooling得到K$\times$K特征图(7$\times$ 7或者14$\times$ 14)作为3个head的输入

4.2、cls head和bbox head

• 首先是样本的标注与分配

1）计算由RPN网络最终生成的所有anchors与每个GT bbox的IoU值，取每个anchor的最大IoU值进行判断：IoU>0.5的样本中选择128个正样本，小于0.5的样本中选择384个负样本

2）对于类别标签：正样本为对应的GT bbox的类别标签，负样本为0；对于bbox回归的标签：与该样本最大IoU值的GT bbox的实际偏移量

• 将这512个样本anchors通过ROI Align获取得到512个7$\times$ 7$\times$ 256的特征图作为class head和bbox head的输入，进行如下的前向传播：一般为2个3*3卷积+2个fc层
  
• 最后就是损失函数的计算：分类损失为多类别交叉熵损失，回归损失也是Smooth L1损失

4.3、mask head

• mask标签和输入特征图的获取

1）计算每个anchor在原始图像中的大mask，就是根据位置参数在原图像的二值mask框出来的
2）根据计算得到的每个anchor对应尺度的特征图（P2-P6中的一个），将anchor位置参数除以缩放倍数/下采样率再进行ROI Align得到512个14$\times$ 14$\times$ 256的特征图作为mask head的输入；
3）而对于大mask也进行倍数缩放，就可以得到每个anchor在对应特征图上的mask(软掩码)—>二值化—>硬掩码

• 对输入特征图再通过FCN层进行预测得到28$\times$ 28的预测mask结果

• 最后进行损失函数的计算，为交叉熵损失
• 需要注意的是，推理阶段不一样，有先后顺序：