SSD理解及代码解析

本文最后更新于：1 年前

SSD理解

SSD(Single Shot MultiBox Detector)是一个one-stage的目标检测模型,它的网络结构如下:

ssd-yolo-arch-contrast

根据上面的网络结构图,有:

SSD的backbone采用的是VGG的前五层,后面的maxpool以及fc层都被换成卷积层
SSD有6个feature map用于预测分类和回归,即多尺度检测
SSD是通过卷积操作来预测分类和定位回归,而YOLO则是通过全连接层变换到指定维度
SSD是基于anchor来预测的,而不同feature map的cell/grid是对应不同的anchor数目

SSD具体的操作,如default box(就是anchor)生成,正负样本匹配策略等会在下面结合代码一起理解

写于代码分析前

我们知道,目标检测由以下组件组成:backbone,neck,head

backbone: 骨干网络,用来进行特征提取.我们之前学习的分类任务中提及的网络都有良好的特征提取能力,因此用来充当backbone
neck: 用于特征融合,以获取不同尺度的感受野信息
head: 检测头,用来预测目标类别和定位回归

SSD代码分析

整体流程如下:

ResNet50基础网络
backbone用于提取特征
ssd_model完成整个训练及预测的流程
1. 构建多尺度特征图用于预测分类和回归
2. default_box生成
3. 正负样本匹配
- 训练
  - loss计算
    - 定位loss
    - 分类loss(hard negative mining)
- 预测
  - 后处理,如nms

因此本文将以上述顺序进行分析

注:本文分析的是输入为300*300的SSD

ResNet-50

由于这是16年的模型,所以当时作者用的是VGG,但在本文,我们用ResNet-50来作为backbone(ResNet具体信息查看image-classification的一些常用模型)

class Bottleneck(nn.Module):
    expansion = 4

    def __init__(self, in_channel, out_channel, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=in_channel, out_channels=out_channel, kernel_size=1, stride=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channel)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel, kernel_size=3, stride=stride,
                               bias=False, padding=1)
        self.bn2 = nn.BatchNorm2d(out_channel)
        self.conv3 = nn.Conv2d(in_channels=out_channel, out_channels=out_channel * self.expansion, kernel_size=1,
                               stride=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channel * self.expansion)
        self.downsample = downsample

    def forward(self, x):
        identity = x
        if self.downsample is not None:
            identity = self.downsample(x)
        residual = self.conv1(x)
        residual = self.bn1(residual)
        residual = self.relu(residual)

        residual = self.conv2(residual)
        residual = self.bn2(residual)
        residual = self.relu(residual)

        residual = self.conv3(residual)
        residual = self.bn3(residual)

        residual += identity
        residual = self.relu(residual)

        return residual

我们知道,ResNet-50相较于ResNet-18/34,其残差结构块是由三个卷积构成,分别:降维->卷积->升维这三个步骤,呈现两头大中间细,因此称为Bottleneck结构,而对于每一次的卷积操作,都是要经历卷积->Batch Normalization->Non Linearity(ReLU)这几步操作.以上便是Bottleneck的代码部分,其中expansion是用于标识第三层卷积的卷积核数目相较于第一层卷积的卷积核数目的膨胀系数;downsample则是用于identity mapping的维度变换(因为同样shape才可以和residual mapping进行相加)

下图中绿色框标识的则为每一个Bottleneck的out_channel,黄色框标识则为in_channel,而右边灰色框框住的则为每一个Bottleneck需要迭代的次数,因此我们需要定义一个方法用来逐层执行(一层有多个Bottleneck)

resnet-arch-for-code

def _make_layer(self, block, channel, block_num, stride=1):
    downsample = None
    if stride != 1 or self.in_channel != channel * block.expansion:
        downsample = nn.Sequential(
            nn.Conv2d(self.in_channel, channel * block.expansion, kernel_size=1, stride=stride, bias=False),
            nn.BatchNorm2d(channel * block.expansion)
        )
    layers = []
    layers.append(block(self.in_channel, channel, downsample=downsample, stride=stride))
    self.in_channel = channel * block.expansion

    for _ in range(1, block_num):
        layers.append(block(self.in_channel, channel))

    return nn.Sequential(*layers)

_make_layer便是逐层执行的方法,其中block是Bottleneck类,而block_num代表的是一个Bottleneck迭代的次数

由此,可以根据_make_layer方法构建出整个网络的结构,ResNet类的__init__和forward方法如下所示:

def __init__(self, block, blocks_num, num_classes=1000, include_top=True):
    super(ResNet, self).__init__()
    self.include_top = include_top
    self.in_channel = 64

    self.conv1 = nn.Conv2d(3, self.in_channel, kernel_size=7, stride=2, padding=3, bias=False)
    self.bn1 = nn.BatchNorm2d(self.in_channel)
    self.relu = nn.ReLU(inplace=True)

    self.maxpool = nn.MaxPool2d(3, stride=2, padding=1)

    self.layer1 = self._make_layer(block, 64, blocks_num[0])
    self.layer2 = self._make_layer(block, 128, blocks_num[1], stride=2)
    self.layer3 = self._make_layer(block, 256, blocks_num[2], stride=2)
    self.layer4 = self._make_layer(block, 512, blocks_num[3], stride=2)

    if self.include_top:
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512 * block.expansion, num_classes)

    for m in self.modules():
        if isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')

def forward(self, x):
    x = self.conv1(x)
    x = self.bn1(x)
    x = self.relu(x)
    x = self.maxpool(x)

    x = self.layer1(x)
    x = self.layer2(x)
    x = self.layer3(x)
    x = self.layer4(x)

    if self.include_top:
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
    return x

至此,便完成了ResNet50的代码部分

backbone

在开始介绍之前,我们先看下用Resnet-50作为backbone的SSD的网络结构图:(图源B站up:霹雳吧啦Wz)

res50-backbone-ssd-arch

由图中可知:并非整个ResNet-50都用做backbone,我们选取conv4这一层作为第一个做预测的特征图(conv5及后面的层全部抛弃).因此backbone便是截取到conv4,也即是对于nn.Module.children()而言,即是取到索引的第7位,然后根据SSD论文,需要对第四层的步距修改为1,因此也需要对identity mapping的stride也进行修改,代码如下:

class Backbone(nn.Module):
    def __init__(self, pretrain_path=None):
        super(Backbone, self).__init__()
        net = resnet50()
        # 用于后面的特征图的channel
        self.out_channels = [1024, 512, 512, 256, 256, 256]

        if pretrain_path is not None:
            net.load_state_dict(torch.load(pretrain_path))

        self.feature_extractor = nn.Sequential(*list(net.children())[:7])

        conv4_block1 = self.feature_extractor[-1][0]
		
        # 这一行代码不是必须的,是用于ResNet-18/34的
        conv4_block1.conv1.stride = (1, 1)
        conv4_block1.conv2.stride = (1, 1)
        conv4_block1.downsample[0].stride = (1, 1)

    def forward(self, x):
        x = self.feature_extractor(x)
        return x

由此,我们便完成了backbone的构造,当我们通过如下方式调用

1
2
3

image_test = torch.randn(1,3,300,300)
backbone = Backbone()
feature_extractor = backbone(image_test)

便可以得到1*1024*38*38这样大小的tensor

ssd_model

多尺度特征图的构建

根据上面的网络结构图,可知六个特征图分别为:[38*38*1024,19*19*512,10*10*512,5*5*256,3*3*256,1*1*256],在经过backbone后,我们即可以得到38*38*1024尺寸的特征图,而后的特征图中,均需要经过两个卷积层后得到其尺寸,因此以下的方法便是为获取后续的特征图做的操作,并将相应操作存于模型容器中

def _build_addtional_features(self, input_size):
    """
    :param input_size [1024,512,512,256,256,256]
    """
    additional_blocks = []
    middle_channels = [256, 256, 128, 128, 128]
    for i, (in_ch, out_ch, mid_ch) in enumerate(zip(input_size[:-1], input_size[1:], middle_channels)):
        padding, stride = (1, 2) if i < 3 else (0, 1)
        layer = nn.Sequential(
            nn.Conv2d(in_ch, mid_ch, kernel_size=1, bias=False),
            nn.BatchNorm2d(mid_ch),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid_ch, out_ch, kernel_size=3, padding=padding, stride=stride, bias=False),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(inplace=True)
        )
        additional_blocks.append(layer)
    self.additional_blocks = nn.ModuleList(additional_blocks)

多尺度特征图预测分类和回归

在提及预测分类和回归前,我们先需要知道每一个特征图对应的grid/cell有几个default box

在论文关于scale和aspect ratio中,作者所提及的scale的计算公式如下示:

s_{k}=s_{min}+\frac{s_{max}-s_{min}}{m-1}(k-1),k\in[1,m]

其中 $m$ 根据Suppose we want to use m feature maps for prediction原文这句话可以知道是 $6$ ,即多尺度特征图的个数,但由于下面的aspect ratio即default box的长宽比部分,提及了除了某些比例外,每个特征图还有一个 $1:1$ 的 $s^{\prime}_{k}=\sqrt{s_{k}s_{k+1}}$ ,显然至少需要 $7$ 个 $s_{k}$ ,而以下是作者的GitHub代码部分关于此处的处理:

min_dim = 300
# res3b3_relu ==> 38 x 38
# res5c_relu ==> 19 x 19
# res5c_relu/conv1_2 ==> 10 x 10
# res5c_relu/conv2_2 ==> 5 x 5
# res5c_relu/conv3_2 ==> 3 x 3
# pool6 ==> 1 x 1
mbox_source_layers = ['res3b3_relu', 'res5c_relu', 'res5c_relu/conv1_2', 'res5c_relu/conv2_2', 'res5c_relu/conv3_2', 'pool6']
# in percent %
min_ratio = 20
max_ratio = 95
step = int(math.floor((max_ratio - min_ratio) / (len(mbox_source_layers) - 2)))
min_sizes = []
max_sizes = []
for ratio in xrange(min_ratio, max_ratio + 1, step):
  min_sizes.append(min_dim * ratio / 100.)
  max_sizes.append(min_dim * (ratio + step) / 100.)
min_sizes = [min_dim * 10 / 100.] + min_sizes
max_sizes = [[]] + max_sizes

我们将它稍加修改用torch打印出来即是:

min\_sizes[30,60,114,168,222,276] \\ max\_sizes[60,114,168,222,276,330]

显然,跟上面论文给出的计算方法及结果均不一样,估计是后期又调了参,而在本文讲解中,采用的是如下这一组scale:

min\_sizes[21, 45, 99, 153, 207, 261]\\max\_sizes[45, 99, 153, 207, 261, 315]

来源amdegroot/ssd.pytorch

而关于aspect ratio,在Experiment中作者说到For conv4_3,conv10_2 and conv11_2, we only associate 4 default boxes at each feature map location–omitting aspect ratios of $\frac{1}{3}$ and 3

也即是1,5,6这三个特征图中只有对应 $s_k$ 的 $[1,2,\frac{1}{2}]$ 这几个宽高比例以及刚刚提及的 $1:1$ 的 $s^{\prime}_{k}$ ,即1,5,6这三个特征图的grid上只有4个default box,而2,3,4则是6个default box

而在我们预测分类和回归的时候,我们需要将多尺度特征图的channel转换为k*(4+num_classes),其中k表示的是该特征图grid对应的default box数目,4是用于分类的四个回归参数(偏移值),num_classes则是用于分类的类别数目

因此,构建用于分类和定位的module如下代码所示:

self.num_defaults = [4, 6, 6, 6, 4, 4]
   location_extractors = []
   confidence_extractors = []
   for nd, oc in zip(self.num_defaults, self.feature_extractor.out_channels):
       location_extractors.append(nn.Conv2d(oc, nd * 4, kernel_size=3, padding=1))
       confidence_extractors.append(nn.Conv2d(oc, nd * num_classes, kernel_size=3, padding=1))
# 用于定位回归的位置提取器
self.loc = nn.ModuleList(location_extractors)
# 用于预测分类的分类提取器
self.conf = nn.ModuleList(confidence_extractors)

我们结合构建好的多尺度特征图以及上述用于预测分类和回归的module,则有下述代码:

def bbox_view(self, features, loc_extractor, conf_extractor):
    """
    :param features: 多尺度特征图
    :param loc_extractor: self.loc
    :param conf_extractor: self.conf
    :return: locs:Tensor[N,4,8732] confs:Tensor[N,num_classes,8732] N->batch_size
    """
    locs = []
    confs = []
    for f, l, c in zip(features, loc_extractor, conf_extractor):
        # view->给tensor做维度变换,不改变数据
        # [batch_size,n*4,feature_size,feature_size] -> [batch_size,4,-1]
        locs.append(l(f).view(f.size(0), 4, -1))
        confs.append(c(f).view(f.size(0), self.num_classes, -1))
    # 将每一张图的不同尺度的矩阵给拼接起来,然后改变其内存的存储结构 [batch_size,4,8732], [batch_size,num_classes,8732]
    locs, confs = torch.cat(locs, 2).contiguous(), torch.cat(confs, 2).contiguous()
    return locs, confs

其中代码中注解的 $8732$ 是所有尺寸的特征图的default box的总数目

38*38*4+19*19*6+10*10*6+5*5*6+3*3*4+1*1*4=8732

以上便是SSD利用多尺度特征图预测分类和定位回归的过程

default box生成

在训练和预测前,①需要利用匹配策略将default box与ground truth box进行匹配,从而得到正负样本,②而得到了正负样本我们才可以执行论文里的hard negative mining(即使得正负样本比例 $1:3$ ),用以处理密集检测中的正负样本极不平衡的情况,③之后才可以计算训练需要用的loss;④而之于预测阶段,我们需要利用回归参数对default box移动到对应的位置,得到预测框,⑤而后再进行NMS去除冗余框,得到最终的结果

因此在开展后续的代码讲解前,需要先对default box的生成做一个讲解

在上面的介绍中,我们知道了每个feature map的grid对应不同数目的default box,而每个default box又由scale和aspect ratio决定,这里补充一下如何利用scale和aspect ratio计算对应default box的宽高

w_{k}=s_{k}\sqrt{a_{r}} \\ h_{k}=\frac{s_{k}}{\sqrt{a_{r}}}

其中 $s_{k}$ 表示的是该feature map的scale, $a_{r}$ 表示的是对应的宽高比,至此,我们便可以计算出不同特征图的default box不同宽高比的尺寸大小,而每个default box都是在其所属的grid的中心,因此我们可以计算出每一个default box的 ${cx,cy,w,h}$ 坐标

由此可以得出生成default box的代码,如下示:

class DefaultBoxes(object):
    def __init__(self, fig_size, feature_size, steps, scales, aspect_ratios, scale_xy=.1, scale_wh=.2):
        # 输入图像的大小300
        self.fig_size = fig_size
        # 多尺度特征图的大小[38,19,10,5,3,1]
        self.feature_size = feature_size
		
        # 用于后期加速loss收敛
        self.scale_xy_ = scale_xy
        self.scale_wh_ = scale_wh

        # 特征图在原图上的跨度 [8,16,32,64,100,300]
        self.steps = steps
        # 每一个特征图对应的default_box的大小 [21,45,99,153,207,261,315]
        self.scales = scales
        # 每一个特征图对应的default_box的宽高比 [[2],[2,3],[2,3],[2,3],[2],[2]]
        self.aspect_ratios = aspect_ratios
        # 每一个特征图的fk,按原论文应是featur_size,但是这里用了别的方法
        fk = fig_size / np.array(steps)

        self.default_boxes = []
        # 计算每一个特征图的default box
        for idx, single_feature in enumerate(self.feature_size):
            # 计算相对坐标
            sk1 = scales[idx] / fig_size
            sk2 = scales[idx + 1] / fig_size
            sk3 = sqrt(sk1, sk2)
            # 添加两个宽高比为1的default box的宽高
            all_sizes = [(sk1, sk1), (sk3, sk3)]

            # 处理feature_map不同宽高比
            for alpha in aspect_ratios[idx]:
                w, h = sk1 * sqrt(alpha), sk1 / sqrt(alpha)
                all_sizes.append((w, h))
                all_sizes.append((h, w))

            # 计算default_box对应的坐标
            for w, h in all_sizes:
                for i, j in itertools.product(range(single_feature), repeat=2):
                    # 计算每一个小格子对应的default box的中心坐标
                    cx, cy = (j + 0.5) / fk[idx], (i + 0.5) / fk[idx]
                    self.default_boxes.append((cx, cy, w, h))
        # 转换数据类型为Tensor
        self.dboxes = torch.as_tensor(self.default_boxes, dtype=torch.float32)
        # 把范围夹紧在[0,1],超出1的划为1,小于0的划为0
        self.dboxes.clamp_(min=0, max=1)

        # 将(cx,cy,w,h)转换为(xmin,ymin,xmax,ymax) left top right bottom,便于IoU计算
        self.dboxes_ltrb = self.dboxes.clone()
        self.dboxes_ltrb[:, 0] = self.dboxes[:, 0] - 0.5 * self.dboxes[:, 2]
        self.dboxes_ltrb[:, 1] = self.dboxes[:, 1] - 0.5 * self.dboxes[:, 3]
        self.dboxes_ltrb[:, 2] = self.dboxes[:, 0] + 0.5 * self.dboxes[:, 2]
        self.dboxes_ltrb[:, 3] = self.dboxes[:, 1] + 0.5 * self.dboxes[:, 3]

    @property
    def scale_xy(self):
        return self.scale_xy_

    @property
    def scale_wh(self):
        return self.scale_wh_
	
    # 要用哪个形式的default box,任君选择
    def __call__(self, order='ltrb'):
        if order == "ltrb":
            return self.dboxes_ltrb

        if order == "xywh":
            return self.dboxes

从上面我们知道default box有两种格式,其中一种是xywh,即是位于grid的中心,因为我们预测的是回归参数,即偏移值,是相对于这一格式而言的,因此想要得到最终的预测框,需要在这种格式上对其进行变换得到

另一种格式为ltrb,此种格式便于我们计算IoU,同时也是VOC数据集提供的GTbox的格式.而计算IoU这一metric,即存在于正负样本匹配,也存在于最后的非极大值抑制中

因此这两种格式在求解目标检测问题时是时常相互转换的,二者均不可或缺

正负样本匹配

在完成了default box的生成后,我们要开始训练前,需要根据匹配策略,区分什么样的default box为正样本,什么样的default box为负样本

在论文中,作者提出了两步的匹配策略:

ssd-matching-strategy

即是按照以下策略匹配:

每个GTbox匹配到的最大IoU的default box为正样本
每个default box匹配到的最大IoU的GTbox,其IoU>阈值(0.5)则为正样本

即一张图片比如有2个GTbox,有4个default box,其IoU如下所示:

	dbox1	dbox2	dbox3	dbox4
GTbox1	0.4	0.3	0	0
GTbox2	0.2	0.8	0.6	0.3

根据第一条匹配策略,匹配到的:(GTbox1,dbox1)=0.4和(GTbox2,dbox2)=0.8,这俩为正样本

根据第二条匹配策略,匹配到的:(GTbox2,dbox3)=0.6>阈值(0.5),也为正样本

若是没匹配到,则为负样本,则其匹配到的是背景类,而该default box的位置信息则无意义(定位loss只算正样本的)

因此,有以下代码,注意,传入encode的bboxes_in和labels_in是通过对数据集标注文件进行解析得到的GTbox的真实位置和真实标签(如对VOC数据集进行解析)

class Encoder(object):
    def __init__(self, dboxes):
        # [8732,4]
        self.dboxes = dboxes(order="ltrb")
        # [1,8732]
        self.dboxes_xywh = dboxes(order="xywh").unsqueeze(dim=0)
        # 8732
        self.nboxes = self.dboxes.size(0)
        self.scale_xy = dboxes.scale_xy
        self.scale_wh = dboxes.scale_wh

    def encode(self, bboxes_in, labels_in, criteria=0.5):
        """
        正负样本匹配,其中nboxes是ground truth boxes
        :param bboxes_in: [nboxes*4]
        :param labels_in: [nboxes]
        :param criteria: IoU阈值,用于判断正样本
        :return:
        """
        # [nboxes,8732],这里的dboxes是ltrb格式
        ious = calc_iou_tensor(bboxes_in, self.dboxes)
        # 每个default box对应GTbox的最大值 [8732]
        best_dbox_ious, best_dbox_idx = ious.max(dim=0)
        # 每个GTbox对应default box的最大值 [nboxes]
        best_gtbox_ious, best_gtbox_idx = ious.max(dim=1)

        # 匹配策略一:GTbox最匹配的default box为正样本
        # 因此,将对应的default box匹配GTbox的IoU更改为>0.5的值,使后面置之为正样本 [8732]
        best_dbox_ious.index_fill_(0, best_gtbox_idx, 2.0)
        # 更换满足匹配一的default box的索引,使之与GTbox匹配,因为后面要用GTbox的坐标和标签
        idx = torch.arange(0, best_gtbox_idx.size(0), dtype=torch.int64)
        best_dbox_idx[best_gtbox_idx[idx]] = idx

        # 匹配策略二:与任意GTbox IoU>0.5的default box即为正样本
        mask = best_dbox_ious > criteria

        # [8732],默认填充为0 0匹配的是背景
        labels_out = torch.zeros(self.nboxes, dtype=torch.int64)
        # 给正样本打上与对应GTbox同样的标签
        labels_out[mask] = labels_in[best_dbox_idx[mask]]

        # 给正样本赋予对应GTbox同样的坐标
        bboxes_out = self.dboxes.clone()
        bboxes_out[mask, :] = bboxes_in[best_dbox_idx[mask], :]

        # 转换为xywh格式
        x = 0.5 * (bboxes_out[:, 0] + bboxes_out[:, 2])
        y = 0.5 * (bboxes_out[:, 1] + bboxes_out[:, 3])
        w = bboxes_out[:, 2] - bboxes_out[:, 0]
        h = bboxes_out[:, 3] - bboxes_out[:, 1]
        bboxes_out[:, 0] = x
        bboxes_out[:, 1] = y
        bboxes_out[:, 2] = w
        bboxes_out[:, 3] = h
        return bboxes_out, labels_out

训练

我们完成了上述工作后,便可以开始训练.而训练过程中,我们会通过一些处理工作,获取到一批(上面是单张)数据的正负样本,即:location:Tensor[N,4,8732],confidence:Tensor[N,8732],因此接下来,我们将对这批数据进行loss计算

loss计算

我们对default box预测的是它相对于中心位置的回归参数,也就是说,是相对于中心位置的偏移值,而在论文中,偏移值的计算定义如下:

\hat{g}^{cx}_{j} = \frac{g^{cx}_{j}-d^{cx}_{i}}{d^{w}_{i}} \quad \hat{g}^{cy}_{j} = \frac{g^{cy}_{j}-d^{cy}_{i}}{d^{h}_{i}}\\ \hat{g}^{w}_{j} = \log{\frac{g^{w}_{j}}{d^{w}_{i}}} \quad \hat{g}^{h}_{j} = \log{\frac{g^{h}_{j}}{d^{h}_{i}}}

其中 $j$ 表示的是第 $j$ 个GTbox, $i$ 表示的是第 $i$ 个default box,上标的 $cx,cy,w,h$ 表示计算的对象(中心坐标x等), $g$ 表示的是GTbox, $d$ 表示的是default box

我们通过等式右边的计算便可以得到左侧的偏移值,而我们的目标是使得网络学习的偏移值 $l$ ,尽量地与真实的偏移值相近,即误差越小越好,本文中采用的定位回归loss函数是SmoothL1Loss

而对于分类损失,本文则采取了交叉熵损失函数,形式如下(拆成了正负样本来写):

L_{conf} = - (\sum_{i \in Pos}^{N}{x^{p}_{ij}\log{\hat{c}^{p}_{i}}} + \sum_{i \in Neg}{\log{\hat{c}^{0}_{i}}})

其中 $x^{p}_{ij}$ 是一个指示器,它的取值是 $\{0,1\}$ ,表示的是匹配第 $i$ 个default box和类别为 $p$ 的第 $j$ 个GTbox,对数函数里面的 $\hat{c}^{p}_{i}$ 表示的是第 $i$ 个default box预测属于 $p$ 这一类别的概率值,之所以有 $\hat{}$ 是因为这是经过softmax后的结果

以下是loss计算的__init__方法,定义了上述所说的损失函数:

class Loss(nn.Module):
    def __init__(self, dboxes):
        super(Loss, self).__init__()
        self.scale_xy = 1.0 / dboxes.scale_xy  # 10
        self.scale_wh = 1.0 / dboxes.scale_wh  # 5

        self.location_loss = nn.SmoothL1Loss(reduction='None')
        # [8732,4] -> [4,8732] -> [1,4,8732] 8732 -> number of anchors
        self.dboxes = nn.Parameter(dboxes(order="xywh").transpose(0, 1).unsqueeze(dim=0),
                                   requires_grad=False)
        self.confidence_loss = nn.CrossEntropyLoss(reduction='None')

我们知道,对于定位问题,我们只对正样本计算loss(负样本的loss没有意义),而正样本的位置信息就是GTbox的位置信息(xywh这一类型的),而预测的结果是default box的偏移值,因而需要对传入的正样本信息计算出偏移值,才可以计算定位损失,以下是代码部分:

def _location_vec(self, loc):
    """
    计算gtbox相较于default box的回归参数
    :param loc: gtbox
    :return:
    """
    # [batch_size,2,8732]
    gxy = self.scale_xy * (loc[:, :2, :] - self.dboxes[:, :2, :]) / self.dboxes[:, 2:, :]
    # [batch_size,2,8732]
    gwh = self.scale_wh * (loc[:, 2:, :] / self.dboxes[:, 2:, :]).log()
    # [batch_size,4,8732]
    return torch.cat((gxy, gwh), dim=1).contiguous()

将以上的返回值和预测的回归信息交由前面定义的SmoothL1Loss便可以计算出定位损失

而分类损失则是对正负样本均要计算,但在计算之前,我们需要知道,一张图片,目标可能只有几个,但我们的default box有8k多个,如果正负样本全拿来计算loss,那正样本的loss其实对网络的学习起的作用微乎其微,因为正负样本极不平衡

因此本文提出了hard negative mininng的思想,即正负样本保持 $1:3$ 的比例,负样本是根据置信度损失大小降序排序而被挑选出来的(就是挑置信度损失高的,也就是难以区分为背景的)

以下是Loss类的forward的代码部分,包括了上述说的定位损失以及分类损失的计算,也包括了hard negative mining的处理部分,其中对两种损失的处理是直接相加得到的:

def forward(self, ploc, plabel, gloc, glabel):
    """
    :param ploc: predicted location [N,4,8732]
    :param plabel: predicted label [N,num_classes,8732]
    :param gloc: gt location [N,4,8732]
    :param glabel: gt label [N,8732]
    :return:
    """
    # 获取正样本的mask Tensor [N,8732]
    mask = torch.gt(glabel, 0)
    # 计算每一张图片的正样本数目 [N]
    pos_num = mask.sum(dim=1)

    # 计算ground truth的回归参数 [N,4,8732]
    vec_gd = self._location_vec(gloc)

    # 计算定位损失 [N,4,8732] -> [N,8732]
    loc_loss = self.location_loss(ploc, vec_gd).sum(dim=1)
    # 只计算正样本的loss [N,8732] -> [N] (即每张图片(N的其中一份子)的定位loss组成的)
    loc_loss = (mask.float() * loc_loss).sum(dim=1)

    # 以下是hard negative mining的过程
    # out:[N,8732]
    conf = self.confidence_loss(plabel, glabel)

    conf_neg = conf.clone()
    # 将正样本的分类loss置零,便于后面对置信度loss高的负样本进行获取
    conf_neg[mask] = 0.0
    # 按置信度loss降序排序,其中获取到的sorted_矩阵表示的是排序后的值,
    # 而conf_idx则是排序后的值对应原conf_neg矩阵对应值的索引矩阵
    sorted_, conf_idx = conf_neg.sort(dim=1, descending=True)
    # 升序排一次,使得conf_rank矩阵为原conf_neg矩阵的位次阵,即对应元素代表了其大小(0最大)
    sorted__, conf_rank = conf_idx.sort(dim=1)
    # 负样本是正样本的三倍,总数要小于8732 out:[N,1]
    neg_num = torch.clamp(3 * pos_num, max=mask.size(1)).unsqueeze(-1)
    # 保留置信度loss高的负样本
    neg_mask = torch.lt(conf_rank, neg_num)

    # confidence计算 [N,8732] -> [N] 对每一幅图片的正负样本计算loss,并合在一起(单张图的正负样本的分类loss) [N]
    conf_loss = (conf * (mask.float() + neg_mask.float())).sum(dim=1)

    # 总的loss
    total_loss = conf_loss + loc_loss
    # 避免某一些图片的gt为0
    num_mask = torch.gt(pos_num, 0).float()
    # 避免gt为0的图片作为分母时无法被除
    pos_num = pos_num.float().clamp(min=1e-6)
    # 计算gt非0的图片的loss(每张图片平均gt的loss,再对这一批图片loss取均值)
    ret = (total_loss * num_mask / pos_num).mean(dim=0)
    return ret

预测

后处理

在验证阶段或者预测阶段,我们需要计算出最终预测框的位置(default box加上偏移值),计算出类别概率,并且对冗余的预测框进行去除,最终用一些metric验证当下模型的性能或者输出预测后的图像

以下是后处理部分的__init__:

class PostProcess(nn.Module):
    def __init__(self, dboxes):
        super(PostProcess, self).__init__()
        # [8732,4] -> [1,8732,4]
        self.dboxes_xywh = nn.Parameter(dboxes(order="xywh").unsqueeze(dim=0), requires_grad=False)
        # 0.1
        self.scale_xy = dboxes.scale_xy
        # 0.2
        self.scale_wh = dboxes.scale_wh
        # NMS的阈值
        self.criteria = .5
        # 一张图片允许的最多框数
        self.max_output = 100

最终预测框

要得到最终预测框,就要将模型预测的回归参数作用到default box上,同时也需要对模型预测的类别进行softmax处理,以转换为概率值,以下是获得最终预测框的代码:

def scale_back_batch(self, bboxes_in, scores_in):
    """
    传入的bboxes_in和scores_in,是
    经过预测后的多张图片多个default box的回归参数以及多个类别的置信度
    本方法目的:
    ①得到最终预测框;
    ②对预测框变换坐标格式,便于后期nms;
    ③对多个类别置信度用softmax获得类别概率
    :param bboxes_in: [N,4,8732]
    :param scores_in: [N,num_classes,8732]
    :return: (bboxes_in,scores_in)
    """
    # [N,8732,4]
    bboxes_in = bboxes_in.permute(0, 2, 1)
    # [N,8732,num_classes]
    scores_in = scores_in.permute(0, 2, 1)

    # loss为了加速收敛分别对xy和wh除以了0.1和0.2,这里要乘回来,得到预测的回归参数
    bboxes_in[:, :, :2] = self.scale_xy * bboxes_in[:, :, :2]
    bboxes_in[:, :, 2:] = self.scale_wh * bboxes_in[:, :, 2:]

    # 最终预测框的位置
    bboxes_in[:, :, :2] = bboxes_in[:, :, :2] * self.dboxes_xywh[:, :, 2:] + self.dboxes_xywh[:, :, :2]
    bboxes_in[:, :, 2:] = bboxes_in[:, :, 2:].exp() * self.dboxes_xywh[:, :, 2:]

    # 转换为ltrb格式,便于之后计算nms
    l = bboxes_in[:, :, 0] - 0.5 * bboxes_in[:, :, 2]
    r = bboxes_in[:, :, 0] + 0.5 * bboxes_in[:, :, 2]
    t = bboxes_in[:, :, 1] - 0.5 * bboxes_in[:, :, 3]
    b = bboxes_in[:, :, 1] + 0.5 * bboxes_in[:, :, 3]

    bboxes_in[:, :, 0] = l
    bboxes_in[:, :, 1] = t
    bboxes_in[:, :, 2] = r
    bboxes_in[:, :, 3] = b

    # 对置信度用softmax转换为概率值 [N,8732,num_classes]
    return bboxes_in, F.softmax(scores_in, dim=-1)

nms

当我们得到最终的预测框以及概率值,我们便要进入非极大值抑制(nms),这一步的目的是去除冗余框

比如人脸检测的时候,在一个人的周围会生成很多框,但是只应留下一个最合适的,很明显即应该留下预测为人的概率值最大的

若是在一幅图片中有两张人脸需要检测,也依旧如此:①保留概率最大的预测框②与概率最大的预测框进行IoU计算,若大于阈值则认为检测的是同一张人脸,应该删去③若小于等于阈值则认为可能是别的人脸,则此时在剩下的预测框中(原先①中最大的已被剔除,②中大于阈值的也被剔除了)重复以上过程

最终即可以得到我们所期望的对应的预测框

以下是nms代码(官方给我们调用的以及自己写的),其中pytorch官方有提供相应函数:

def nms(boxes, scores, iou_threshold):
    """
    非极大值抑制,去除冗余框
    :param boxes: tensor[N,4]
    :param scores: tensor[N]
    :param iou_threshold: float be used to remove the similarity region boxes
    :return: tensor,one dimension,refer to the useful boxes indexes
    """
    return torch.ops.torchvision.nms(boxes, scores, iou_threshold)

def nms_(boxes: Tensor, scores: Tensor, iou_threshold: float) -> Tensor:
    """
    根据scores排序,将最大分值的box的索引加入keep,然后用最大分值的box与其他boxes相比(IoU),大于阈值则滤去,
    然后在剩下的scores又拿最大的进行比较,直至全部比较完毕
    :param boxes: Tensor[N,4]
    :param scores: Tensor[N]
    :param iou_threshold:
    :return: keep:Tensor
    """
    # 最终留下来的boxes的索引用keep存储
    keep = []
    # 对scores升序排序,取索引
    idxs = scores.argsort()
    # 当idxs还有元素,就继续循环,除非元素都没了(或移除,或留下至keep)
    while idxs.numel() > 0:
        max_score_idx = idxs[-1]
        # [1,4]
        max_score_box = boxes[max_score_idx][None, :]
        keep.append(max_score_idx)
        # 表示当前只剩下最后一个元素,且已经加入keep中,直接退出循环即可
        if idxs.size(0) == 1:
            break
        # 去除得分最大框
        idxs = idxs[:-1]
        # [M,4],其中M表示是idxs的size,指取出还没处理的框
        other_boxes = boxes[idxs]
        # [1,M]
        ious = calc_iou_tensor(max_score_box, other_boxes)
        # 去除冗余框!
        idxs = idxs[ious[0] <= iou_threshold]

    keep = idxs.new(keep)
    return keep

而从上述代码,我们也不难看出来nms是只对同一类别进行处理,若是多个类别呢?我们当然可以一个类别逐次进行nms,但是这样效率较慢

我们知道IoU计算实际上就是计算两个框的重叠性,那也即是说IoU本质上也是受到空间距离因素影响的,我们只要将不同类别的框的空间距离拉开不同的距离,而同类别框的空间距离拉开相同的距离,就可以使得nms起作用

以下是代码部分:

def batched_nms(bboxes, scores, idxs, iou_threshold):
    """
    :param bboxes: 位置信息
    :param scores: 分类概率
    :param idxs: 标签
    :param iou_threshold: IoU阈值
    :return: keep:Tensor 留下来的框的索引
    """
    if bboxes.numel() == 0:
        return torch.empty(0, dtype=torch.int64, device=bboxes.device)

    # 在所有default box的所有坐标中取最大值,作为一个离散因子,离散化不同类别,因为nms是针对同一类别所做的
    max_factor = bboxes.max()
    # 偏移值:按照类别离散化
    offsets = idxs.to(bboxes) * (max_factor + 1)
    # 对bboxes加上偏移值
    bboxes_nms = bboxes + offsets[:, None]
    # nms
    keep = nms(bboxes_nms, scores, iou_threshold)
    return keep

其实在进行nms之前,我们还需要对一些低概率box,面积过小box进行移除,之后再进入nms,以减少不必要的计算

以下是包含了上述说的以及nms过程的代码,其中该方法是对单张图片进行处理,需要注意,SSD中预测定位与分类是两个不大相干的任务,即预测定位无需关心它是哪一类别的目标,而在这个方法中,我们采取了Faster RCNN的方法,来将一个预测框与多个类别相关联,以便于后面的nms:

def decode_singe_image(self, bboxes_in, scores_in, criteria, max_output):
    """
    :param bboxes_in: [8732,4]
    :param scores_in: [8732,num_classes]
    :param criteria: nms criteria
    :param max_output: the maximal number of output boxes
    :return:
    """
    device = bboxes_in.device
    num_classes = scores_in.shape[-1]

    # 裁剪边界
    bboxes_in = bboxes_in.clamp(min=0, max=1)

    # default box只关注回归问题 -> default box的回归问题与object类别相关(用于简化nms)
    # [8732,4] -> [8732,21,4]
    bboxes_in = bboxes_in.repeat(1, num_classes).reshape(scores_in.shape[0], -1, 4)

    # 创建标签,用以与default box们对应
    # [num_classes]
    labels = torch.arange(num_classes, device=device)
    # [num_classes] -> [1,num_classes] -> [8732,num_classes]
    labels = labels.view(1, -1).expand_as(scores_in)

    # 移除背景类的信息 num_classes -> num_classes - 1
    bboxes_in = bboxes_in[:, 1:, :]
    scores_in = scores_in[:, 1:]
    labels = labels[:, 1:]

    # 合并default box的维度(之前扩了num_classes倍)
    # [8732*20,4] 20是用的voc数据集,有21个类别(含背景)
    bboxes_in = bboxes_in.reshape(-1, 4)
    # [8732*20]
    scores_in = scores_in.reshape(-1)
    # [8732*20]
    labels = labels.reshape(-1)

    # 移除低概率目标
    idxs = torch.where(torch.gt(scores_in, 0.05))[0]
    bboxes_in, scores_in, labels = bboxes_in[idxs, :], scores_in[idxs], labels[idxs]

    # 移除面积很小很小的box
    widths, heights = (bboxes_in[:, 2] - bboxes_in[:, 0]), (bboxes_in[:, 3] - bboxes_in[:1])
    keep = (widths >= 1 / 300) & (heights >= 1 / 300)
    keep = torch.where(keep)[0]
    bboxes_in, scores_in, labels = bboxes_in[keep, :], scores_in[keep], labels[keep]

    # NMS
    keep = batched_nms(bboxes_in, scores_in, labels, iou_threshold=criteria)

    # 按iou高的取预测框
    keep = keep[:max_output]
    bboxes_out = bboxes_in[keep, :]
    scores_out = scores_in[keep]
    labels_out = labels[keep]

    return bboxes_out, scores_out, labels_out

以下是PostProcess这个类的forward方法:

def forward(self, bboxes_in, scores_in):
    """
    :param bboxes_in: [N,4,8732]
    :param scores_in: [N,num_classes,8732]
    """
    # 获得最终预测框,并且对置信度做softmax运算获得每个类别的概率
    # [N,8732,4],[N,8732,num_classes]
    bboxes, probs = self.scale_back_batch(bboxes_in, scores_in)

    # 用于TorchScript的
    outputs = torch.jit.annotate(List[Tuple[Tensor, Tensor, Tensor]], [])

    # 逐张图片做出预测结果
    for bbox, prob in zip(bboxes.split(1, 0), (probs.split(1, 0))):
        # [1,8732,4],[1,8732,num_classes] -> [8732,4],[8732,num_classes]
        bbox = bbox.squeeze(0)
        prob = prob.squeeze(0)
        outputs.append(self.decode_singe_image(bbox, prob, self.criteria, self.max_output))
    return outputs

至此,本文已经讲完了SSD进行目标检测任务的整体模型搭建过程,而在实际应用中,还需要涉及到如数据库数据解析(比如VOC或COCO数据集的解析),数据预处理,优化器选择,学习率调整等部分

最后给出一张利用SSD进行预测的图像:

ssd_experiment

SSD Object Detection

本博客所有文章除特别声明外，均采用 CC BY-SA 4.0 协议，转载请注明出处！

Pipeline Parallelism GPipe And PipeDream 上一篇

Some lib of Python 下一篇