第二部分生成RPN和ROI网络的输入部分真让我研究了好久,有些细节部分也是反复琢磨才明白,感觉代码还是要多写多看。接下来这部分是模型的原理,就是网络部分,分为基网络提取特征、RPN和ROI网络。感觉会有点难,终于涉及pytorch了,加油加油!这部分参考了这篇博客。ROI pooling部分参考博客。
模型这块真的是挺难的,特别散,我觉得还是应该先看trainer,py,对训练有个整体框架看起来才能轻松点,这一块我差点儿就要放弃了。我先写训练代码部分,再回头写的这部分。
首先看的是model/region_proposal_network.py函数
def __init__(
self, in_channels=512, mid_channels=512, ratios=[0.5, 1, 2],
anchor_scales=[8, 16, 32], feat_stride=16,
proposal_creator_params=dict(),
):
super(RegionProposalNetwork, self).__init__()
self.anchor_base = generate_anchor_base(
anchor_scales=anchor_scales, ratios=ratios) #调用generate_anchor_base()函数,生成左上角9个anchor_base
self.feat_stride = feat_stride
self.proposal_layer = ProposalCreator(self, **proposal_creator_params)
n_anchor = self.anchor_base.shape[0] #9
self.conv1 = nn.Conv2d(in_channels, mid_channels, 3, 1, 1)
self.score = nn.Conv2d(mid_channels, n_anchor * 2, 1, 1, 0)
self.loc = nn.Conv2d(mid_channels, n_anchor * 4, 1, 1, 0)
normal_init(self.conv1, 0, 0.01) 归一化
normal_init(self.score, 0, 0.01)
normal_init(self.loc, 0, 0.01)
def forward(self, x, img_size, scale=1.):
n, _, hh, ww = x.shape #(batch_size,512,H/16,W/16),其中H,W分别为原图的高和宽
anchor = _enumerate_shifted_anchor(
np.array(self.anchor_base),
self.feat_stride, hh, ww) #在9个base_anchor基础上生成hh*ww*9个anchor,对应到原图坐标
n_anchor = anchor.shape[0] // (hh * ww) #hh*ww*9/hh*ww=9
h = F.relu(self.conv1(x)) #512个3x3卷积(512, H/16,W/16),后面都不写batch_size了
rpn_locs = self.loc(h) #n_anchor(9)*4个1x1卷积,回归坐标偏移量。(9*4,hh,ww)
rpn_locs = rpn_locs.permute(0, 2, 3, 1).contiguous().view(n, -1, 4) #转换为(n,hh,ww,9*4)后变为(n,hh*ww*9,4)
rpn_scores = self.score(h) #n_anchor(9)*2个1x1卷积,回归类别。(9*2,hh,ww)
rpn_scores = rpn_scores.permute(0, 2, 3, 1).contiguous()#转换为(n,hh,ww,9*2)
rpn_softmax_scores = F.softmax(rpn_scores.view(n, hh, ww, n_anchor, 2), dim=4) #计算{Softmax}(x_{i}) = \{exp(x_i)}{\sum_j exp(x_j)}
rpn_fg_scores = rpn_softmax_scores[:, :, :, :, 1].contiguous()#得到前景的分类概率
rpn_fg_scores = rpn_fg_scores.view(n, -1)#得到所有anchor的前景分类概率
rpn_scores = rpn_scores.view(n, -1, 2)#得到每一张feature map上所有anchor的网络输出值
rois = list()
roi_indices = list()
for i in range(n): #n为batch_size数
roi = self.proposal_layer(
rpn_locs[i].cpu().data.numpy(),
rpn_fg_scores[i].cpu().data.numpy(),
anchor, img_size,
scale=scale) # 调用ProposalCreator函数, rpn_locs维度(hh*ww*9,4),rpn_fg_scores维度为(hh*ww*9),anchor的维度为(hh*ww*9,4), img_size的维度为(3,H,W),H和W是经过数据预处理后的。计算(H/16)x(W/16)x9(大概20000)个anchor属于前景的概率,取前12000个并经过NMS得到2000个近似目标框G^的坐标。roi的维度为(2000,4)
batch_index = i * np.ones((len(roi),), dtype=np.int32)
rois.append(roi) #rois为所有batch_size的roi
roi_indices.append(batch_index)
rois = np.concatenate(rois, axis=0)#按行拼接(即没有batch_size的区分,每一个[]里都是一个anchor的四个坐标)
roi_indices = np.concatenate(roi_indices, axis=0)#这个 roi_indices在此代码中是多余的,因为我们实现的是batch_siae=1的网络,一个batch只会输入一张图象。如果多张图象的话就需要存储索引以找到对应图像的roi
return rpn_locs, rpn_scores, rois, roi_indices, anchor #rpn_locs的维度(hh*ww*9,4),rpn_scores维度为(hh*ww*9,2), rois的维度为(2000,4),roi_indices用不到,anchor的维度为(hh*ww*9,4)
```
写完上面这块我就真的不知道从哪里开始写了,找了半天,理清了思路才开始继续,
下面写的是model/faster_rcnn_vgg16,py
def decom_vgg16():
if opt.caffe_pretrain: #这里我运行的时候设置为True,使用下载下来的caffe预训练模型而不是torchvision的
model = vgg16(pretrained=False)
if not opt.load_path:
model.load_state_dict(t.load(opt.caffe_pretrain_path))#加载参数信息
else:
model = vgg16(not opt.load_path)
features = list(model.features)[:30] #加载预训练模型vgg16的conv5_3之前的部分
classifier = model.classifier
classifier = list(classifier)
del classifier[6] #这一块看不懂,先放着
if not opt.use_drop: #删除两个dropout
del classifier[5]
del classifier[2]
classifier = nn.Sequential(*classifier)
for layer in features[:10]: #冻结vgg16前2个stage,不进行反向传播
for p in layer.parameters():
p.requires_grad = False
return nn.Sequential(*features), classifier #拆分为特征提取网络和分类网络
class FasterRCNNVGG16(FasterRCNN):#分别对特征VGG16的特征提取部分、分类部分、RPN网络、VGG16RoIHead网络进行了实例化
feat_stride = 16 #vgg16通过5个stage下采样16倍
def __init__(self,
n_fg_class=20,
ratios=[0.5, 1, 2],
anchor_scales=[8, 16, 32]
): #总类别数为20类,三种尺度三种比例的anchor
extractor, classifier = decom_vgg16() #conv5_3及之前的部分,分类器(这个暂时不知道什么鬼)
rpn = RegionProposalNetwork(
512, 512,
ratios=ratios,
anchor_scales=anchor_scales,
feat_stride=self.feat_stride,
) #返回rpn_locs, rpn_scores, rois, roi_indices, anchor
head = VGG16RoIHead(
n_class=n_fg_class + 1,
roi_size=7,
spatial_scale=(1. / self.feat_stride),
classifier=classifier
)#接着的下面分析VGG16RoIHead(),n_class = 21(加上背景)
super(FasterRCNNVGG16, self).__init__(
extractor,
rpn,
head,
) #相当于给faster_rcnn传入参数extractor, rpn, head
class VGG16RoIHead(nn.Module):
def __init__(self, n_class, roi_size, spatial_scale, classifier):
# n_class includes the background
super(VGG16RoIHead, self).__init__()
self.classifier = classifier #vgg16中的最后两个全连接层
self.cls_loc = nn.Linear(4096, n_class * 4)
self.score = nn.Linear(4096, n_class)
normal_init(self.cls_loc, 0, 0.001)
normal_init(self.score, 0, 0.01) #全连接层权重初始化
self.n_class = n_class #加上背景21类
self.roi_size = roi_size #7
self.spatial_scale = spatial_scale # 1/16
self.roi = RoIPooling2D(self.roi_size, self.roi_size, self.spatial_scale) #将大小不同的roi变成大小一致,得到pooling后的特征,大小为[300, 512, 7, 7]。.利用Cupy实现在线编译的
def forward(self, x, rois, roi_indices):
roi_indices = at.totensor(roi_indices).float() #ndarray->tensor
rois = at.totensor(rois).float()
indices_and_rois = t.cat([roi_indices[:, None], rois], dim=1)
# NOTE: important: yx->xy
xy_indices_and_rois = indices_and_rois[:, [0, 2, 1, 4, 3]]
indices_and_rois = xy_indices_and_rois.contiguous() #把tensor变成在内存中连续分布的形式
pool = self.roi(x, indices_and_rois) #接下来分析roi_module.py中的RoI()
pool = pool.view(pool.size(0), -1) #flat操作
fc7 = self.classifier(pool) #decom_vgg16()得到的calssifier,得到4096
roi_cls_locs = self.cls_loc(fc7) #(4096->84)
roi_scores = self.score(fc7) #(4096->21)
return roi_cls_locs, roi_scores #roi回归输出的是128*84,然而真实位置参数是128*4和真实标签128*1
class RoI(Function): #将大小不同的roi变成大小一致,得到pooling后的特征,大小为[300, 512, 7, 7]。反正意思为将每个feature map 变成统一大小为7x7的。
def forward(self, x, rois):
x = x.contiguous() #变成在内存中连续分布的形式
rois = rois.contiguous()
self.in_size = B, C, H, W = x.size()
self.N = N = rois.size(0) #每张图所有的anchors数
output = t.zeros(N, C, self.outh, self.outw).cuda()
self.argmax_data = t.zeros(N, C, self.outh, self.outw).int().cuda()
self.rois = rois
args = [x.data_ptr(), rois.data_ptr(),
output.data_ptr(),
self.argmax_data.data_ptr(),
self.spatial_scale, C, H, W,
self.outh, self.outw,
output.numel()] #data_ptr()返回一个时间戳,numel()返回一个tensor变量内所有元素
stream = Stream(ptr=torch.cuda.current_stream().cuda_stream)
self.forward_fn(args=args,
block=(CUDA_NUM_THREADS, 1, 1),
grid=(GET_BLOCKS(output.numel()), 1, 1),
stream=stream) #这一步是实现RoI pooling的关键,通过Cupy实现在线编译,调用roi_cupy代码。
return output
下面是model/faster_rcnn.py
def forward(self, x, scale=1.):
img_size = x.shape[2:] (H,W)
h = self.extractor(x) #输入一张图片得到其特征图feature map
rpn_locs, rpn_scores, rois, roi_indices, anchor = \
self.rpn(h, img_size, scale) #给定特征图后产生一系列RoIs
roi_cls_locs, roi_scores = self.head(
h, rois, roi_indices) #利用这些RoIs对应的特征图对这些RoIs中的类别进行分类,并提升定位精度
return roi_cls_locs, roi_scores, rois, roi_indices
在类FasterRCNN中便初始化了这三个重要步骤:
self.extractor
self.rpn
self.head
函数forward实现前向传播。