读Mask R-CNN源码备忘录（训练部分）

此文为读Mask RCNN源码过程中的随笔，很“流水账”，我想价值在于对照着源码把每个步骤的“输入”、“输出”张量的维度标注了一下，会有助于对整体代码的理解。可能有些错误或遗漏，希望发现者指正，以期共同进步。

源码：https://github.com/matterport/Mask_RCNN

训练部分

模型输入：

input_image (batch_size, height, width, channels) #默认(2, 1024, 1024, 3)

input_image_meta (batch_size, 1 + 3 + 3 + 4 + 1 + config.NUM_CLASSES) #默认(2, 93)

input_rpn_match (batch_size, num_anchors, 1) #默认(2, 261888, 1)

# num_anchors计算
import numpy as np
BACKBONE_STRIDES = [4,8,16,32,64] #基础cnn网络（resnet101）输出的五层特征图对应输入图像的缩放比例
IMAGE_SHAPE = (1024, 1024) #输入图像尺寸
RPN_ANCHOR_RATIOS = [0.5, 1, 2] #每个像素取三种width/height比例的anchor

# 每张特征图所有像素取三个尺寸的anchor
num_anchors = sum([x[0][0] * x[0][1] / np.square(x[1]) * len(RPN_ANCHOR_RATIOS) for x in zip([IMAGE_SHAPE] * len(BACKBONE_STRIDES), BACKBONE_STRIDES)])
print(num_anchors)

input_rpn_bbox (batch_size, config.RPN_TRAIN_ANCHORS_PER_IMAGE, 4) #默认(2, 256, 4)

input_gt_class_ids (batch_size, config.MAX_GT_INSTANCES) #默认(2, 100)

 

* 注意：这里用norm_boxes_graph函数将原始坐标做了归一化处理

input_gt_boxes (batch_size, config.MAX_GT_INSTANCES, 4) -> gt_boxes (batch_size, config.MAX_GT_INSTANCES, 4) #默认(2, 100, 4)

* 注意：这里需要判断config.USE_MINI_MASK是True或者False

input_gt_masks (batch_size, config.MINI_MASK_SHAPE[0], config.MINI_MASK_SHAPE[1], config.MAX_GT_INSTANCES) #默认(2, 56, 56, 100)

input_gt_masks (batch_size, config.IMAGE_SHAPE[0], config.IMAGE_SHAPE[1], config.MAX_GT_INSTANCES) #默认(2, 1024, 1024, 100)

第一步：resnet_graph网络

C2, C3, C4, C5 为resnet_graph的四个stage输出，输出尺寸依次为：

C2: (batch_size, config.IMAGE_SHAPE[0] / config.BACKBONE_STRIDES[0], config.IMAGE_SHAPE[1] / config.BACKBONE_STRIDES[0], 256) #默认(2, 256, 256, 256)

C3: (batch_size, config.IMAGE_SHAPE[0] / config.BACKBONE_STRIDES[1], config.IMAGE_SHAPE[1] / config.BACKBONE_STRIDES[1], 512) #默认(2, 128, 128, 512)

C4: (batch_size, config.IMAGE_SHAPE[0] / config.BACKBONE_STRIDES[2], config.IMAGE_SHAPE[1] / config.BACKBONE_STRIDES[2], 1024) #默认(2, 64, 64, 1024)

C5: (batch_size, config.IMAGE_SHAPE[0] / config.BACKBONE_STRIDES[3], config.IMAGE_SHAPE[1] / config.BACKBONE_STRIDES[3], 2048) #默认(2, 32, 32, 2048)

P5: 对C5做(1, 1)卷积filters=256，效果就是在维度不变的情况下，将特征图数量由2048降为256 #默认(2, 32, 32, 256)

P4: 对P5做(2, 2)上采样，并且将C4做(1, 1)卷积filters=256，然后将两者相加，效果就是P5尺寸加倍与C4卷积后尺寸相等，将两者相加作为P4 #默认(2, 64, 64, 256)

P3: 对P4做(2, 2)上采样，并且将C3做(1, 1)卷积filters=256，然后将两者相加，效果就是P4尺寸加倍与C3卷积后尺寸相等，将两者相加作为P3 #默认(2, 128, 128, 256)

P2: 同P3，P4 #默认(2, 256, 256, 256)

对P2, P3, P4, P5做(3, 3)卷积filters=256，padding="SAME"，所以尺寸不变，只是对特征抽象级别提升了一下。

这里多出来一个P6，P6是对P5做了(1, 1)池化stride=2，所以尺寸减半，效果是对原图像做了隔像素采样。#默认(2, 16, 16, 256)

rpn_feature_maps: [P2, P3, P4, P5, P6]用于RPN网络

mrcnn_feature_maps: [P2, P3, P4, P5]用于classifier heads(FPN网络)

第二步：生成ANCHORS

对[P2, P3, P4, P5, P6]各层anchor数量可以这样计算：

[int(x[0][0] * x[0][1] / np.square(x[1]) * config.IMAGE_CHANNEL_COUNT / config.RPN_ANCHOR_STRIDE) for x in zip([config.IMAGE_SHAPE] * len(config.BACKBONE_STRIDES), config.BACKBONE_STRIDES)]

每一层anchor数量是，(原图像尺寸高/本层stride) * (原图像尺寸宽/本层stride) * 每个像素3个长宽比 / anchor_stride，以P2层为例：(1024/4)*(1024/4)*3/1 = 196608 #此处默认anchor_stride=1

anchors: 最终[P2, P3, P4, P5, P6]对应的anchors为[196608, 49152, 12288, 3072, 768]，总共anchors数量为261888 #默认(2, 261888, 4)

第三步：创建RPN模型，对每一层特征图做预测

输入：

rpn_feature_maps #默认[(2, 256, 256, 256), (2, 128, 128, 256), (2, 64, 64, 256), (2, 32, 32, 256), (2, 16, 16, 256)]

输出：

rpn_class_logits #默认(2, 261888, 2)

rpn_class #默认(2, 261888, 2)

rpn_bbox #默认(2, 261888, 4)

思路：通过对每个特征图做卷积操作，分别得到class和bbox偏移量的回归网络分支

shared：对特征图做(3, 3)卷积，filters=512, strides=anchor_stride，padding='same' #默认（以P2特征图为例）(2, 256, 256, 512)

x：对shared做(1, 1)卷积，filters=2 * anchors_per_location，padding='valid' #默认（以P2特征图为例）(2, 256, 256, 6)

rpn_class_logits：对上步x做reshape(batch_size, -1, 2)  #默认（以P2特征图为例）(2, 196608, 2)

rpn_probs：对上步rpn_class_logits做softmax #默认（以P2特征图为例）(2, 196608, 2)

x: 对shared做(1, 1)卷积，filters=4 * anchor_per_location， padding='valid' #默认（以P2特征图为例）(2, 256, 256, 12)

rpn_bbox：对上步x做reshape(batch_size, -1, 4) #默认（以P2特征图为例）(2, 196608, 4)

至此rpn网络生成：输入每层特征图input_feature_map，输出每层特征图预测到的rpn_class_logits, rpn_probs, rpn_bbox

对每层输出做KL.Concatenate操作后最终得到的输出：

rpn_class_logits #默认(2, 261888, 2)

rpn_class #默认(2, 261888, 2)

rpn_bbox #默认(2, 261888, 4)

第四步：创建ProposalLayer层，对上一步结果通过NMS获取rpn_rois

计算proposal_count #默认2000

ProposalLayer层：

输入:

rpn_class #默认(2, 261888, 2)

rpn_bbox #默认(2, 261888, 4)

anchors #默认(2, 261888, 4)

输出:

rpn_rois #默认(2, 2000, 4)

思路：

scores: rpn_class代表此实例是前景（fg）或者背景（bg）的概率预测，[batch, num_anchors, (bg prob, fg prob)]，所以第三维的第二项可以看成这个实例检测到物体的分数 scores = rpn_class[:, :, 1] #默认(2, 261888)

deltas: rpn_bbox代表原anchor与真实实例位置的偏移量，deltas = rpn_bbox * np.reshape(self.config.RPN_BBOX_STD_DEV, [1, 1, 4]) #默认(2, 261888, 4)

anchors即上边通过feature_map尺寸得到的所有anchor #默认(2, 261888, 4)

首先通过scores倒排序，取前pre_nms_limit个（去除分数低的实例），得到scores，deltas，pre_nms_anchors。#默认pre_nms_limit为6000

scores #默认(2, 6000, 1)

deltas #默认(2, 6000, 4)

pre_nms_anchors #默认(2, 6000, 4)

boxes: 然后用deltas调整pre_nms_anchors位置，并且剪切超出边界的anchor，得到boxes #默认(2, 6000, 4)

proposals：对boxes做NMS操作，并且根据scores取其中前proposal_count个，不足用0padding #默认(2, 2000, 4)

target_rois: rpn_rois 等于proposals #默认(2, 2000, 4)

第五步：创建DetectionTargetLayer，生成检测目标

DetectionTargetLayer：

输入：

proposals: target_rois #默认(2, 2000, 4)

gt_class_ids: input_gt_class_ids #默认(2, 100) 

gt_boxes #默认(2, 100, 4)

gt_masks: input_gt_masks #默认(2, 1024, 1024, 100)

输出:

rois #默认(2, 200, 4)

target_class_ids #默认(2, 200)

target_deltas: target_bbox #默认(2, 200, 4)

target_mask #默认(2, 200, 28, 28)

*注意，detection_targets_graph是对单个特征图做操作的，所以不包含batch_size

detection_targets_graph

输入：

proposals: target_rois #默认(2000, 4)

gt_class_ids: input_gt_class_ids #默认(100, ) 

gt_boxes #默认(100, 4)

gt_masks: input_gt_masks #默认(1024, 1024, 100)

输出：

rois #默认(200, 4)

target_class_ids #默认(200, )

target_deltas: target_bbox #默认(200, 4)

target_mask #默认(200, 28, 28)

思路：

首先去除0padding后

proposals #默认(<=2000, 4)

gt_boxes #默认(<=100, 4)

gt_class_ids #默认(<=100)

gt_masks #默认(<=100, 1024, 1024) 或者 (<=100, 56, 56)

将gt_boxes分为crowd实例和非crowd实例

将proposals中与非crowd gt_boxes iou >= 0.5的实例作为正实例

将proposals中与非crowd gt_boxes iou < 0.5并且与crowd gt_boxes iou < 0.001的实例作为负实例

提取positive_count=int(config.TRAIN_ROIS_PER_IMAGE * config.ROI_POSITIVE_RATIO)个正实例positive_rois，和对应比例的负实例negative_rois。 #默认postitive_rois(<=int(200 * 0.33), 4), negative_rois(<=int(200 * 0.67), 4)

获取与positive_rois对应的roi_gt_boxes, roi_gt_class_ids, roi_masks, 然后用positive_rois与roi_gt_boxes计算proposal的偏移量deltas #默认roi_gt_boxes(<=66, 4), roi_gt_class_ids(<=66), roi_masks(<=66, 1024, 1024)或者(<=66, 56, 56),deltas(<=66, 4)

boxes=positive_rois #默认(<=66, 4)

masks 是roi_masks通过boxes裁剪和缩放得到的 #默认(<=66, 28, 28)

rois：postitive_rois与negative_rois合并在一起 #默认(<=200, 4)

最后将rois, roi_gt_class_ids, deltas, masks用0padding补齐第一维200

第六步：创建fpn_classifier_graph （FPN）

fpn_classifier_graph：

输入rois, feature_maps, image_meta, pool_size, num_classes, train_bn=True, fc_layers_size=1024

rois  #默认(2, 200, 4)

feature_maps: mrcnn_feature_maps #默认[P2, P3, P4, P5] [(2, 256, 256, 256), (2, 128, 128, 256), (2, 64, 64, 256), (2, 32, 32, 256)]

image_meta #默认(2, 93)

pool_size #默认7

num_classes #默认81

输出mrcnn_class_logits, mrcnn_class, mrcnn_bbox

mrcnn_class_logits #默认(2, 200, 81)

mrcnn_class #默认(2, 200, 81)

mrcnn_bbox #默认(2, 200, 81, 4)

思路：

首先通过PyramidROIAlign 层对每个feature_map用roi提取区域特征：

***PyramidROIAlign 开始***

输入:

boxes: rois #默认(2, 200, 4)

image_meta #默认(2, 93)

feature_maps #默认[P2, P3, P4, P5] [(2, 256, 256, 256), (2, 128, 128, 256), (2, 64, 64, 256), (2, 32, 32, 256)]

输出:

pooled regions，即roi坐标对应的feature_map上的一小块区域的集合 #默认(2, 200, 7, 7, 3)

思路：

首先根据每个roi对应的面积计算roi_level #默认(2, 200)

循环依次处理2～5层特征图

处理方法：以P2层为例，先在roi_level中过滤出level=2的索引ix, 然后将这些roi从boxes中取出，得到level_boxes，获取level_boxes中每个roi对应的batch索引box_indices，用tf.image.crop_and_resize获得level_boxes在feature_map上的区域特征并存入pooled，以此类推获取每层中提取的区域特征。

ix #默认(level=2的roi个数, 2)

level_boxes #默认(level=2的roi个数, 4)

box_indices #默认(level=2的roi个数)

理解上边的处理思路可以参考这几句代码：

>>> import tensorflow as tf
>>> import numpy as np
>>> roi_level = np.array([[1, 0, 3], [2, 0, 0]])
>>> boxes = np.array([[[1,2,3,4],[5,6,7,8],[9,10,11,12]], [[13,14,15,16],[17,18,19,20],[21,22,23,24]]])
>>> ix = tf.where(tf.equal(roi_level, 0))
>>> ix
<tf.Tensor: shape=(3, 2), dtype=int64, numpy=
array([[0, 1],
       [1, 1],
       [1, 2]])>
>>> level_boxes = tf.gather_nd(boxes, ix)
>>> level_boxes
<tf.Tensor: shape=(3, 4), dtype=int64, numpy=
array([[ 5,  6,  7,  8],
       [17, 18, 19, 20],
       [21, 22, 23, 24]])>
>>> box_indices = tf.cast(ix[:, 0], tf.int32)
>>> box_indices
<tf.Tensor: shape=(3,), dtype=int32, numpy=array([0, 1, 1], dtype=int32)>

pooled #默认(2 * 200, 7, 7, 3)

box_range #默认(2 * 200, 1)

box_to_level #默认(2 * 200, (batch_index, roi_index, box_range_value))

ix：对box_to_level根据batch_index * 10000 + roi_index倒序排序以保持原boxes倒序排序 #默认(2 * 200, 1)

pooled 根据ix重排序 #默认(2 * 200, 7, 7, 3)

对pooled reshape(2, 200, 7, 7, 3)作为输出，至此PyramidROIAlign完成，输出x #默认(2, 200, 7, 7, 3)

***PyramidROIAlign 结束***

x: 用TimeDistributed对batch中每个sample(特征图)的每个roi做(7, 7)卷积，filters=1024，padding=valid，然后做batchnorm操作 #默认(2, 200, 1, 1, 1024)

x: 用TimeDistributed对batch中每个sample(特征图)的每个roi做(1, 1)卷积，filters=1024，然后做batchnorm操作 #默认(2, 200, 1, 1, 1024)

shared: 对x做squeeze操作  #默认(2, 200, 1024)

mrcnn_class_logits: 对shared每个roi做dense(81)操作 #默认(2, 200, 81)

mrcnn_probs: 对mrcnn_class_logits中每个roi做softmax操作 #默认(2, 200, 81)

x: 对shared中每个roi做dense(4 * 81)操作 #默认(2, 200, 4 * 81)

mrcnn_bbox: 对x做reshape #默认(2, 200, 81, 4)

第七步：创建build_fpn_mask_graph

输入：

rois #默认(2, 200, 4)

feature_maps: mrcnn_feature_maps #默认[P2, P3, P4, P5] [(2, 256, 256, 256), (2, 128, 128, 256), (2, 64, 64, 256), (2, 32, 32, 256)]

image_meta: input_image_meta #默认(2, 93)

pool_size #默认14

num_classes #默认81

train_bn #默认False

输出：

mrcnn_mask #默认(2, 200, 28, 28, 81)

思路：

x: 首先通过PyramidROIAlign层对每个feature_map用roi提取区域特征，具体过程参照第六步中相同操作 #默认(2, 200, 14, 14, 3)

x: 对x中每个roi做(3, 3)卷积，filters=256，padding=same，然后batchnorm，重复4次 #默认(2, 200, 14, 14, 256)

x: 对x中每个roi做(2, 2)反卷积，filters=256，strides=2 #默认(2, 200, 28, 28, 256)

#计算反卷积的尺寸公式
new_rows = (rows - 1) * strides[0] + kernel_size[0] - 2 * padding[0] + output_padding[0]
new_cols = (cols - 1) * strides[1] + kernel_size[1] - 2 * padding[1] + output_padding[1]

#默认
new_rows = (14 - 1) * 2 + 2 - 2 * 0 + 0 = 28
new_cols = (14 - 1) * 2 + 2 - 2 * 0 + 0 = 28

x: mrcnn_mask 对x每个roi做(1, 1)卷积，filters=81，strides=1，以此作为mrcnn_mask输出 #默认(2, 200, 28, 28, 81)

第八步：计算loss

rpn_class_loss: 通过rpn_class_loss_graph获取

输入：

rpn_match: input_rpn_match #默认(2, 261888, 1) target

rpn_class_logits: rpn_class_logits #默认(2, 261888, 2)

输出：

loss #默认(1)

思路：过滤掉neutral anchors，得到只有正负样本的anchors，然后用交叉熵损失函数求loss，过程中的中间量维度为，rpn_match(2, 261888)，anchor_class(2, 261888)，indices(正负样本个数, 1)，rpn_class_logits(正负样本个数, 2)，anchor_class(正负样本个数, 1)

rpn_bbox_loss: 通过rpn_bbox_loss_graph获取

输入：

target_bbox: input_rpn_bbox (batch_size, config.RPN_TRAIN_ANCHORS_PER_IMAGE, (dy, dx, log(dh), log(dw))) #默认(2, 256, 4) target

rpn_match: input_rpn_match #默认(2, 261888, 1) 辅助target 

rpn_bbox #默认(2, 261888, 4)

输出：

loss #默认(1)

思路：分别从target_bbox和rpn_bbox中过滤出正样本，然后用smooth_l1_loss求loss，过程中间变量维度为，rpn_match(2, 261888)，indices(正例样本个数, 1)，rpn_bbox(正例样本个数, 4)，batch_counts(2, 1)，target_bbox(正例样本个数, 4)

*注意：这里batch_pack_graph中counts取正样例的方式，是因为target_bbox中正例样本都在最前边

class_loss: 通过mrcnn_class_loss_graph获取

输入：

target_class_ids #默认(2, 200) target

pred_class_logits: mrcnn_class_logits #默认(2, 200, 81)

active_class_ids #默认(2, 81) 辅助target

输出：

loss #默认(1)

思路：先对pred_class_logits利用argmax求pred_class_ids，然后用交叉熵损失求损失值，最后过滤掉dataset中不存在的类别，中间变量维度为，pred_class_ids(2, 200)，pred_active(2, 200)

bbox_loss: 通过mrcnn_bbox_loss_graph获取

输入：

target_bbox #默认(2, 200, 4) target

target_class_ids #默认(2, 200) 辅助target

pred_bbox: mrcnn_bbox #默认(2, 200, 81, 4)

输出：

loss #默认(1)

思路：先把前两维合并(2, 200, ...) -> (2 * 200, ...)，然后从target_bbox和pred_bbox中分别过滤出正例样本，用smooth_l1_loss求loss，中间变量维度为，target_class_ids(2 * 200, )，target_bbox(2 * 200, 4)，pred_bbox(2 * 200, 81, 4)，positive_roi_ix(正例样本个数, 1)，positive_roi_class_ids(正例样本个数, 1)，indices(正例样本个数, 2)，target_bbox(正例样本个数, 4)，pred_bbox(正例样本个数, 4)

mask_loss: 通过mrcnn_mask_loss_graph获取

输入：

target_masks: target_mask #默认(2, 200, 28, 28)

target_class_ids #默认(2, 200) 

pred_masks: mrcnn_mask #默认(2, 200, 28, 28, 81)

输出：

loss #默认(1)

思路：与mrcnn_bbox_loss_graph基本一样

第九步：创建model

inputs = [input_image, input_image_meta, input_rpn_match, input_rpn_bbox, input_gt_class_ids, input_gt_boxes, input_gt_masks]

outputs = [rpn_class_logits, rpn_class, rpn_bbox, mrcnn_class_logits, mrcnn_class, mrcnn_bbox, mrcnn_mask, rpn_rois, output_rois, rpn_class_loss, rpn_bbox_loss, class_loss, bbox_loss, mask_loss]

第十步：data_generator

anchors #默认(261888, 4)

image #默认(1024, 1024, 3)

image_meta #默认(1 + 3 + 3 + 4 + 1 + 81)

gt_class_ids #默认(instance_count, ) 注意：instance_count是根据图片中实际标注的样本个数一致的

gt_boxes #默认(instance_count, 4)

gt_masks #默认(1024, 1024, instance_count) 或者 (56, 56, instance_count)

build_rpn_targets:

输入：

image_shape: image.shape #默认(1024, 1024, 3)

anchors #默认(261888, 4)

gt_class_ids #默认(instance_count, )

gt_boxes #默认(instance_count, 4)

输出：

rpn_match #默认(216888, )

rpn_bbox: (config.RPN_TRAIN_ANCHORS_PER_IMAGE, 4) #默认(256, 4)

思路：

anchors与crowd gt_bbox的iou<0.001的作为no_crowd_bool，anchors与非crowd gt_bbox的iou<0.3和no_crowd_bool的交集作为负样本，anchors与非crowd gt_bbox的iou>0.7的为正样本，对于每个非crowd gt_bbox，与之iou最大的那个anchor为正样本。使正样本数量不大于config.RPN_TRAIN_ANCHORS_PER_IMAGE / 2，并且正负样本数量之和为config.RPN_TRAIN_ANCHORS_PER_IMAGE，最终得到rpn_match，其中包含config.RPN_TRAIN_ANCHORS_PER_IMAGE个正负样本，正样本不超过config.RPN_TRAIN_ANCHORS_PER_IMAGE/2。用anchors正样本与gt_boxs计算出偏移rpn_bbox。config.RPN_TRAIN_ANCHORS_PER_IMAGE默认值为256

循环batch_size次，组成一个batch后返回，batch_image_meta(2, 93)，batch_rpn_match(2, 216888, 1)，batch_rpn_bbox(2, 256, 4)，batch_images(2, 1024, 1024, 3)，batch_gt_class_ids(2, 100)，batch_gt_boxes(2, 100, 4)，batch_gt_masks(2, 1024, 1024, 100)，inputs=[batch_images, batch_image_meta, batch_rpn_match, batch_rpn_bbox, batch_gt_class_ids, batch_gt_boxes, batch_gt_masks]

*注意：

要区分各个阶段的bbox指的是坐标还是偏移量，rpn_bbox指的偏移量，gt_bbox指的坐标，target_bbox是偏移量，input_rpn_bbox是偏移量，mrcnn_bbox是偏移量

所有的roi都是坐标