【PyTorch】SSD六个stage的边框生成源码debug解析

最新推荐文章于 2022-11-07 21:31:45 发布

原创最新推荐文章于 2022-11-07 21:31:45 发布 · 472 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#PyTorch #SSD #PriorBox

卷积神经网络专栏收录该内容

7 篇文章

订阅专栏

本文介绍了一个PyTorch实现的PriorBox类，用于计算SSD目标检测模型中的先验框坐标。通过配置参数，如特征图尺寸、先验框尺寸范围等，PriorBox能够为每个特征图位置生成一组不同比例和大小的边界框。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

环境：

pytorch1.1.0

源码：

from __future__ import division
from math import sqrt as sqrt
from itertools import product as product
import torch

voc = {
    'num_classes': 21,
    'lr_steps': (80000, 100000, 120000),
    'max_iter': 120000,
    'feature_maps': [38, 19, 10, 5, 3, 1],#每个stage特征点个数
    'min_dim': 300, #输入图像的最小维度【300*300】
    'steps': [8, 16, 32, 64, 100, 300],#根据本层特征图与输入图片之间点得映射关系得到
    'min_sizes': [30, 60, 111, 162, 213, 264],
    'max_sizes': [60, 111, 162, 213, 264, 315],
    'aspect_ratios': [[2], [2, 3], [2, 3], [2, 3], [2], [2]],
    'variance': [0.1, 0.2],
    'clip': True,       #clip修剪
    'name': 'VOC',
}

#PriorBox返回得是所有default box得四个参数归一化后的值，即得到所有先验框的位置
class PriorBox(object):
    """Compute priorbox coordinates in center-offset form for each source
    feature map.
    """
    def __init__(self, cfg):#cfg当前传入的是voc的信息
        super(PriorBox, self).__init__()
        self.image_size = cfg['min_dim'] #获取输入图像的尺寸
        # number of priors for feature map location (either 4 or 6)
        self.num_priors = len(cfg['aspect_ratios'])
        self.variance = cfg['variance'] or [0.1]
        self.feature_maps = cfg['feature_maps']
        self.min_sizes = cfg['min_sizes']
        self.max_sizes = cfg['max_sizes']
        self.steps = cfg['steps']
        self.aspect_ratios = cfg['aspect_ratios']
        self.clip = cfg['clip']
        self.version = cfg['name']
        for v in self.variance:
            if v <= 0:
                raise ValueError('Variances must be greater than 0')

    def forward(self):
        mean = []  #feaure_maps = [38, 19, 10, 5, 3, 1]每一层的特征点个数
        for k, f in enumerate(self.feature_maps):
            #product表示可迭代对象的笛卡尔内积，repeate=2表示两个range(f)的列表进行内积
            for i, j in product(range(f), repeat=2):
                f_k = self.image_size / self.steps[k] #图像尺度变换的幅度
                # unit center x,y
                cx = (j + 0.5) / f_k
                cy = (i + 0.5) / f_k

                # aspect_ratio: 1
                # rel size: min_size
                s_k = self.min_sizes[k]/self.image_size
                mean += [cx, cy, s_k, s_k]

                # aspect_ratio: 1
                # rel size: sqrt(s_k * s_(k+1))
                s_k_prime = sqrt(s_k * (self.max_sizes[k]/self.image_size))
                mean += [cx, cy, s_k_prime, s_k_prime]
                #[[2], [2, 3], [2, 3], [2, 3], [2], [2]]
                # rest of aspect ratios
                for ar in self.aspect_ratios[k]:
                    mean += [cx, cy, s_k*sqrt(ar), s_k/sqrt(ar)]
                    mean += [cx, cy, s_k/sqrt(ar), s_k*sqrt(ar)]
        # back to torch land
        output = torch.Tensor(mean).view(-1, 4) #表示每个feature_maps边框个数
        if self.clip:#clamp_函数返回的是最大值和最小值之间的数，防止边框越界
            output.clamp_(max=1, min=0)
        return output

def main():
    priorbox = PriorBox(voc)
    output=priorbox.forward()
    print(output)
    print('output.shape:',output.shape)

if __name__ == '__main__':
    main()

结果：

tensor([[0.0133, 0.0133, 0.1000, 0.1000],
        [0.0133, 0.0133, 0.1414, 0.1414],
        [0.0133, 0.0133, 0.1414, 0.0707],
        ...,
        [0.5000, 0.5000, 0.9612, 0.9612],
        [0.5000, 0.5000, 1.0000, 0.6223],
        [0.5000, 0.5000, 0.6223, 1.0000]])
output.shape: torch.Size([8732, 4])