Faster RCNN代码整体框架
前言
Faster RCNN作为经典的双阶段目标检测算法,掌握其中的思想和代码实现的方法,对于我们实现单阶段目标检测或者双阶段目标检测都是很有帮助的。
相较于单阶段目标检测,双阶段目标检测主要多了一步生成proposal,也就是候选框的生成。在Faster RCNN中,对于图像中的生成的每一个anchor而言,首先要经过RPN(在这里只区分前景或者背景)做第一次筛选,选出概率大的一些anchor(train时采用预测目标概率前2000个anchor,test时采用1000个)作为proposal,proposal用于后续Faster RCNN的计算,这样的方式也就称为双阶段目标检测。
我以源码为基础,将代码更精简了一点,实现起来更加容易,并且细化了代码量并在每部分做了详细的注释以及shape的变化,以便更好的理解代码的原理和实现的过程。
整体组成
整体Faster RCNN主要包括transform、backbone、RegionProposalNetwork(RPN)、RoIHead、postprocess组成。
其中:transform的作用是将输入图像,缩放到一个给定的大小中,组成一个batch数据输入给网络。因为对于输入图像,其都是一张张大小不一尺寸的图像,无法输入到网络做平行计算,因此transform就是为了将输入图像缩放到指定大小,组成一个batch,然后输入到网络中。
backbone:对于输入图像进行特征提取,得到特征图。在这里采用的是ResNet50+FPN作为backbone。特征图包括ResNet50中layer1、layer2、layer3、layer4的输出并做FPN,将输出的通道数(256, 512, 1024, 2048)都调整为256,并且在layer4生成的特征图后使用一个maxpool,得到’pool’。一共生成5个预测特征图,分别为layer1, layer2, layer3, layer4, maxpool。对应的键为0, 1, 2, 3, pool.
RegionProposalNetwork: 将图像中生成的anchors进行筛选,生成proposal。
RoIHead: 由RoI pooling(RoIAlign pooling) + MLP(FastRCNN中ROI-pooling后的展平操作+两个全连接层) + FastRCNNPredictor(输出的预测部分,包括bbox回归参数和类别概率)组成。对RPN生成的proposal进行bbox回归参数和类别参数的预测。
postprocess: 将网络的预测结果还原到原图像尺寸上。
Faster RCNN Model整体代码
import torch
from torch import nn, Tensor
import torch.nn.functional as F
from torchvision.ops import MultiScaleRoIAlign
from transforms import RcnnTransforms
from RPN import AnchorsGenerator, RPNHead, RegionProposalNetwork
from RoIHead import RoIHead
class FasterRCNNBase(nn.Module):
def __init__(self, backbone, rpn, roi_heads, transform):
super(FasterRCNNBase, self).__init__()
self.transform = transform
self.backbone = backbone
self.rpn = rpn
# roi pooling + MLP + FastRCNNPredictor + postprocess detections
self.roi_heads = roi_heads
def forward(self, images, targets=None):
# 这里输入的images大小都是不同的,后面进行transform后修改成为同样大小的tensor打包成一个batch
# images为dataloader的返回结果,list中每一个图像都是一个tensor类型
# targets也是一个list,其里面每个元素是一个字典类型,包含每一个图像中的标注信息。
original_image_sizes = []
for img in images:
val = img.shape[-2:] # B C H W
assert len(val) == 2 # 防止输入的是个一维向量
# 记录原始图像的size 因为transform后图像size发生变化,在最终的输出会映射回原图大小。
original_image_sizes.append((val[0], val[1]))
# 对图像进行预处理 得到新的 images targets。
# 在输入transform之前,图像都是一张张的图像,图像大小不一。通过transform,将图像统一放到一个给定大小的tensor中,得到了一个batch的数据。
images, targets = self.transform(images, targets)
# 将图像输入backbone得到特征图
features = self.backbone(images.tensors)
# 将特征层以及标注target信息传入rpn中 -> 区域建议框proposal和RPN的损失loss
# proposals shape: [num_proposals, 4] 4: xmin ymin xmax ymax
proposals, proposal_losses = self.rpn(images, features, targets)
# 将rpn生成的数据以及标注target信息传入fastRCNN后半部分
detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)
# 对网络的预测结果进行后处理(主要将bboxes还原到原图像尺度上)
# 得到最终检测的一系列目标和FastRCNN的损失值
# images.image_sizes 通过transform后的图像尺寸 original_image_sizes预处理前的图像尺寸。
# 将预测结果还原到最初的图像尺寸上
detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)
losses = {}
losses.update(detector_losses)
losses.update(proposal_losses)
if self.training:
return losses
return detections
class TwoMLPHead(nn.Module):
def __init__(self, in_channels, representation_size):
super(TwoMLPHead, self).__init__()
self.fc6 = nn.Linear(in_channels, representation_size)
self.fc7 = nn.Linear(representation_size, representation_size)
def forward(self, x):
# x为经过roiAlign后
# x shape: [batch_size*RPN输出的正负样本总数(512), out_channels, 7, 7]
# flatten: -> [batch_size*RPN输出的正负样本总数(512), out_channels * 7 * 7]
x = x.flatten(start_dim=1)
x = F.relu(self.fc6(x))
x = F.relu(self.fc7(x))
return x
class FastRCNNPredictor(nn.Module):
def __init__(self, in_channels, num_classes):
super(FastRCNNPredictor, self).__init__()
self.cls_score = nn.Linear(in_channels, num_classes)
self.bbox_pred = nn.Linear(in_channels, num_classes * 4)
def forward(self, x):
# x shape: [batch_size*RPN输出的正负样本总数(512), 1024] 1024: 对应TwoMLPHead最终的输出维度
# x = x.flatten(start_dim=1)
scores = self.cls_score(x)
bbox_deltas = self.bbox_pred(x)
return scores, bbox_deltas
class FasterRCNN(FasterRCNNBase):
# backbone 自行定义的特征提取网络
# num_classes 检测的类别个数(需要加入背景类)eg: VOC: num_classes=21
def __init__(self, backbone, num_classes=None,
# transform parameter
min_size=800, max_size=1333, # 预处理resize时限制的最小尺寸与最大尺寸
image_mean=None, image_std=None, # 预处理normalize时使用的均值和方差(ImageNet)
# RPN parameters
# rpn_anchor_generator用于生成anchor的生成器。
# rpn_head由一个3*3conv、分类层和边界框回归层组成
rpn_anchor_generator=None, rpn_head=None,
# NMS前后相同主要是针对带有FPN的网络。由于FPN存在多个预测的特征图,每层在NMS前proposal数都保留2000个,总共加起来就上万了,然后再通过NMS保留2000个。
# 通过预测信息和anchor生成器,可以得到一系列的proposal。根据预测的score,在进行NMS之前先滤除一部分proposal,将剩余的proposal输入到NMS。
rpn_pre_nms_top_n_train=2000, rpn_pre_nms_top_n_test=1000, # rpn中在nms处理前保留的proposal数(根据score)
rpn_post_nms_top_n_train=2000, rpn_post_nms_top_n_test=1000, # rpn中在nms处理后保留的proposal数
rpn_nms_thresh=0.7, # rpn中进行nms处理时使用的iou阈值
# rpn计算损失时,采集正负样本设置的阈值。anchor与GT的IOU大于0.7标记为正样本,anchor与任何一个GT的IOU均小于0.3标记为负样本
# fg: foreground bg: background
rpn_fg_iou_thresh=0.7, rpn_bg_iou_thresh=0.3,
# 在正样本和负样本中进行随机采样 计算RPN损失时。总共采样256个样本。rpn_positive_fraction为正样本占全部样本的比例
rpn_batch_size_per_image=256, rpn_positive_fraction=0.5, # rpn计算损失时采样的样本数,以及正样本占总样本的比例
# Box parameters(ROI Head中的参数)
# box_roi_pool对应ROI Pooling
# box_head 对应MLPHead
# box_predictor对应两个fc层,一个预测类别概率、另一个用于预测边界回归框参数。
box_roi_pool=None, box_head=None, box_predictor=None,
# 移除小概率目标的阈值 fastRCNN中进行nms处理的阈值 对预测结果根据score排序取前100个目标
box_score_thresh=0.05, box_nms_thresh=0.5, box_detections_per_img=100,
# fastRCNN计算误差时,采集正负样本设置的阈值.proposal与GT的IOU大于0.5,定义为正样本。proposal与所有的GT的IOU均小于0.5,定义为负样本。
box_fg_iou_thresh=0.5, box_bg_iou_thresh=0.5,
# 共采样512个样本,正样本占全部样本的0.25.
box_batch_size_per_image=512, box_positive_fraction=0.25, # fast rcnn计算误差时采样的样本数,以及正样本占所有样本的比例
bbox_reg_weights=None):
# 预测特征层的channels resnet50+FPN out_channels=256
out_channels = backbone.out_channels
# 对数据进行标准化,缩放,打包成batch等处理部分(预处理)
# 预处理的图像均值和方差
if image_mean is None:
image_mean = [0.485, 0.456, 0.406]
if image_std is None:
image_std = [0.229, 0.224, 0.225]
transform = RcnnTransforms(min_size, max_size, image_mean, image_std)
# 若anchor生成器为空,则自动生成针对resnet50_fpn的anchor生成器
if rpn_anchor_generator is None:
anchor_sizes = ((32,), (64,), (128,), (256,), (512,))
aspect_ratios = ((0.5, 1.0, 2.0),) * len(anchor_sizes) # aspect_ratios将(0.5, 1.0, 2.0)重复五次,每一个元素对应一个特征层上的尺度
rpn_anchor_generator = AnchorsGenerator(
anchor_sizes, aspect_ratios
)
# 生成RPN通过滑动窗口预测网络部分
if rpn_head is None:
# rpn_anchor_generator.num_anchors_per_location()[0]: 对应每一层预测特征图上生成anchor的数量 由size和ratios决定
rpn_head = RPNHead(
out_channels, rpn_anchor_generator.num_anchors_per_location()[0]
)
# 默认rpn_pre_nms_top_n_train = 2000, rpn_pre_nms_top_n_test = 1000,
# 默认rpn_post_nms_top_n_train = 2000, rpn_post_nms_top_n_test = 1000,
rpn_pre_nms_top_n = dict(training=rpn_pre_nms_top_n_train, testing=rpn_pre_nms_top_n_test)
rpn_post_nms_top_n = dict(training=rpn_post_nms_top_n_train, testing=rpn_post_nms_top_n_test)
# 定义整个RPN框架
# rpn_batch_size_per_image 为RPN计算损失时采用的正负样本的总个数
# rpn_positive_fraction 在计算损失时 正样本占全部样本的比例
# rpn_pre_nms_top_n 在进入NMS之前 对于每一个预测特征图所保存的目标个数
rpn = RegionProposalNetwork(
rpn_anchor_generator, rpn_head,
rpn_fg_iou_thresh, rpn_bg_iou_thresh,
rpn_batch_size_per_image, rpn_positive_fraction,
rpn_pre_nms_top_n, rpn_post_nms_top_n, rpn_nms_thresh)
# ROIAlign Pooling
# 经过roiAlign后,每一个proposal的特征矩阵为7*7
if box_roi_pool is None:
box_roi_pool = MultiScaleRoIAlign(
featmap_names=['0', '1', '2', '3'], # 在哪些特征层进行roi pooling
output_size=[7, 7], # 指定输出大小
sampling_ratio=2)
# FastRCNN中ROI-pooling后的展平处理两个全连接层部分(MLPHead)
if box_head is None:
resolution = box_roi_pool.output_size[0] # 默认等于7
representation_size = 1024
box_head = TwoMLPHead(
out_channels * resolution ** 2,
representation_size
)
# 在box_head的输出上预测部分
if box_predictor is None:
representation_size = 1024
box_predictor = FastRCNNPredictor(
representation_size,
num_classes)
# 将roi pooling, box_head以及box_predictor结合在一起,组成RoIHead
roi_heads = RoIHead(
# box
box_roi_pool, box_head, box_predictor, # 分别对应上面三个部分
box_fg_iou_thresh, box_bg_iou_thresh, # 0.5 0.5
box_batch_size_per_image, box_positive_fraction, # box_batch_size_per_image每张图像中选择的proposal的数量(512) 0.25
bbox_reg_weights, # 默认(10 10 5 5 )
# 对得到的最终结果进行后处理时使用 0.05 0.5 100
box_score_thresh, box_nms_thresh, box_detections_per_img)
super(FasterRCNN, self).__init__(backbone, rpn, roi_heads, transform)
接下来的几篇文章,依次实现transform、backbone、RPN、RoIHead部分。