pytorch-ssd源码解读（二）------------detection（预测层）

「爱情、让人受尽委屈。」 2022-02-13 10:15 145阅读 0赞

# 一、定义 #

detection层是ssd预测阶段的最后一层。它接收底层网络输出的位置偏移量（loc\_data）、各个框置信度（conf\_data）以及默认框（prior\_data）。该层的作用是整合各层的预测结果，过滤置信度太低的预测框，通过类内nms抑制大量相同的预测框。

# 二、代码解读 #

### 1.输入 ###

*  loc\_data    网络六个层预测的坐标偏移。Shape:\[batch,num\_priors\*4\]
 *  conf\_data     各预测框对应各类的置信度（每个预测框针对每一类都预测一个得分，因此每个预测框对应num\_classes个得分）                            Shape: \[batch\*num\_priors,num\_classes\]
 *  prior\_data   默认框（上篇博客中介绍过），网络的预测框其实就是针对这个默认框的偏移量。 Shape: \[num\_priors,4\]

### 2.参数配置 ###

*  top\_k                          一张图片中，每一类保存top\_k个预测框
 *  conf\_thresh                置信度阈值，置信度低于该阈值的预测框会被抛弃
 *  nms\_thresh                nms阈值

### 3.输出格式 ###

output的Shape为 \[batch, num\_classes, top\_k, 5\]

取其中一个输出output\[i,j,k,:\]表示在当前mini\_batch中的第i张图片的第j类的第k个框对应的预测结果。

最后一维的五个数依次为\[score,xmin,ymin,xmax,ymax\]

###  4.重点代码解读 ###

NO.1

"""
    conf_data  :  [batch*num_priors,num_classes]
    conf_preds :  [batch,num_classes,num_priors]
    """
    conf_preds = conf_data.view(num, num_priors,
                                  self.num_classes).transpose(2, 1)

conf\_data是预测框对应各类的置信度，为一个二维Tensor，Shape: \[batch\*num\_priors,num\_classes\]。

conf\_data.view(num, num\_priors,self.num\_classes)先将Tensor由\[batch\*num\_priors,num\_classes\]转变为\[batch,num\_priors,num\_classes\]。然后通过transpose(2,1)调换第1和第2维度,最终变为\[batch,num\_classes,num\_priors\]形式.

该步骤的目的是方便下面按照batch,class的顺序进行处理。

NO.2

decoded_boxes = decode(loc_data[i], prior_data, self.variance)

由预测的偏移值和默认框生成最终的预测框 。

偏移值和默认框为\[x,y,w,h\]形式，解码的预测框为\[xmin,ymin,xmax,ymax\]形式。

最终decode\_boxes的shape为\[num\_priors,4\]，关于decode函数文末会详细讲一下。

NO.3

c_mask = conf_scores[cl].gt(self.conf_thresh)  #选择大于设置的阈值的得分掩码
    scores = conf_scores[cl][c_mask]               #筛选大于阈值的得分
    l_mask = c_mask.unsqueeze(1).expand_as(decoded_boxes)   #c_mask是一维的，将其扩展为[num_priors,4]
    boxes = decoded_boxes[l_mask].view(-1, 4)

这段代码先筛选出大于阈值的置信度，然后取出对应的预测框。

个人感觉不如直接写为下面的代码简洁：

c_mask = conf_scores[cl].gt(self.conf_thresh)  
    scores = conf_scores[cl][c_mask]              
    boxes = decoded_boxes[c_mask]

NO.5

ids, count = nms(boxes, scores, self.nms_thresh, self.top_k)

该步骤对上一步经过阈值筛选的boxes和scores进行nms极大值抑制，取出符合条件的预测框。

输出ids为符合条件的预测框的index是一个列表，count为符合条件的框的个数。

# 附件： #

def forward(self, loc_data, conf_data, prior_data):
            """
            Args:
                loc_data: (tensor) Loc preds from loc layers           预测框的偏移量
                    Shape: [batch,num_priors*4]
                conf_data: (tensor) Shape: Conf preds from conf layers   每个预测框中各类别的得分
                    Shape: [batch*num_priors,num_classes]
                prior_data: (tensor) Prior boxes and variances from priorbox layers       默认框
                    Shape: [1,num_priors,4]
    
            该函数根据默认框和预测框计算出最终的检测框
            """
            num = loc_data.size(0)  # batch size
            num_priors = prior_data.size(0)  #默认框的数目
            output = torch.zeros(num, self.num_classes, self.top_k, 5)
            conf_preds = conf_data.view(num, num_priors,
                                        self.num_classes).transpose(2, 1)
    
            # Decode predictions into bboxes.
            for i in range(num):
                """
                由预测的偏移值和默认框生成最终的预测框
                偏移值和默认框为[x,y,w,h]形式，解码的预测框为[xmin,ymin,xmax,ymax]形式
                decode_boxes的shape为[num_priors,4]
                """
                decoded_boxes = decode(loc_data[i], prior_data, self.variance)
                # For each class, perform nms
                #一张图片中所有的预测框得分，shape为[num_classes,num_priors]
                conf_scores = conf_preds[i].clone()
                #类内做nms
                for cl in range(1, self.num_classes):
                    c_mask = conf_scores[cl].gt(self.conf_thresh)  #选择大于设置的阈值的得分掩码
                    scores = conf_scores[cl][c_mask]               #筛选大于阈值的得分
                    if scores.dim() == 0:                          #如果当前类没有符合条件的预测框，继续下一个类的循环
                        continue
                    """
                    decoded_boxes[l_mask]其实是一维的,排列方式形如[x1min,y1min,x1max,y1max,x2min,y2min,x2max,y2max...]
                    因此decoded_boxes[l_mask].view(-1, 4) 才会转变为[num_priors,4]的形状，使得每一行对应一个bbox
                    个人感觉大可不必使用l_mask,可将下面代码直接替代为
                    boxes = decode_boxes[c_mask]
                
                    """
                    l_mask = c_mask.unsqueeze(1).expand_as(decoded_boxes)   #c_mask是一维的，将其扩展为[num_priors,4]
                    boxes = decoded_boxes[l_mask].view(-1, 4)               
                    # idx of highest scoring and non-overlapping boxes per class
                    ids, count = nms(boxes, scores, self.nms_thresh, self.top_k)
                    """
                    构造输出，最后一个output为(batchsize, num_classes, self.top_k, 5)形状
                    最后一纬的5个数为[score,xmin,ymin,xmax,ymax]
                    """
                    output[i, cl, :count] = \
                        torch.cat((scores[ids[:count]].unsqueeze(1),
                                   boxes[ids[:count]]), 1)
            flt = output.contiguous().view(num, -1, 5)
            _, idx = flt[:, :, 0].sort(1, descending=True)
            _, rank = idx.sort(1)
            flt[(rank < self.top_k).unsqueeze(-1).expand_as(flt)].fill_(0)
            return output