逐字理解目标检测simple-faster-rcnn-pytorch-master代码(三)

第二部分生成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实现前向传播。


你可能感兴趣的:(目标检测)