SSD(Single Shot MultiBox Detector)是一种实时对象检测算法。
它的主要特点是:
使用一个卷积神经网络在一次前馈中对多尺度的特征图进行预测。
默认框定义在不同尺度和不同方块上的,然后预测这些默认框中是否包含物品及其类别。
通过匹配和非极大抑制获得最终的检测结果。
它使用了VGG-16作为骨干模型,然后在conv4_3,conv7,conv8_2,conv9_2和conv10_2这几个层的特征图上进行预测。
- conv4_3用于检测小尺度物体,conv10_2用于检测大尺度物体。
- 每个特征图上的预测结果都包含4个部分:是否包含物体、物体类预测、框的位置预测和框的尺度预测。
SSD使用了几个超参数来控制用于预测的特征图尺度,以及每个特征图上的默认框(也称先验框)尺度。主要有:
- min_dim: 输入图像的最小尺寸
- max_dim: 输入图像的最大尺寸
- scales:用于预测的特征图的尺度,对应不同的层
- aspect_ratios: 每个特征图上的默认框的宽高比
SSD的代码实现如下:
python
import torch
import torch.nn as nn
import torchvision.models as models
class SSD(nn.Module):
def __init__(self, num_classes):
super(SSD, self).__init__()
# 获得VGG-16模型的各个部分
self.vgg = models.vgg16(pretrained=True).features
# 在conv4_3之后添加额外的卷积层
self.conv4_3_norm = nn.Conv2d(512, 1024, kernel_size=3, padding=1)
self.conv4_3_norm_relu = nn.ReLU(inplace=True)
# 在其它层之后也添加卷积层进行预测
self.conv7_norm = nn.Conv2d(1024,256, kernel_size=3, padding=1)
self.conv7_norm_relu = nn.ReLU(inplace=True)
self.conv8_2_norm = nn.Conv2d(512,256, kernel_size=3, padding=1)
self.conv8_2_norm_relu = nn.ReLU(inplace=True)
self.conv9_2_norm = nn.Conv2d(256,256, kernel_size=3, padding=1)
self.conv9_2_norm_relu = nn.ReLU(inplace=True)
self.conv10_2_norm = nn.Conv2d(256,256, kernel_size=3, padding=1)
self.conv10_2_norm_relu = nn.ReLU(inplace=True)
# 分类器和边界框回归器
self.num_classes = num_classes
self.loc_preds = nn.ModuleList()
self.cl_preds = nn.ModuleList()
def forward(self, x):
x = self.vgg(x)
# conv4_3预测小物体
conv4_3_norm = self.conv4_3_norm(x)
conv4_3_norm_relu = self.conv4_3_norm_relu(conv4_3_norm)
loc_pred_conv4_3, cl_pred_conv4_3 = self.get_preds(conv4_3_norm_relu, 4)
# conv7预测中等大小物体
conv7_norm = self.conv7_norm(x)
conv7_norm_relu = self.conv7_norm_relu(conv7_norm)
loc_pred_conv7, cl_pred_conv7 = self.get_preds(conv7_norm_relu, 6)
# conv8_2, conv9_2, conv10_2 一一对应的预测
loc_pred_conv8_2, cl_pred_conv8_2 = self.get_preds(self.conv8_2_norm_relu
# conv10_2预测大物体
conv10_2_norm = self.conv10_2_norm(x)
conv10_2_norm_relu = self.conv10_2_norm_relu(conv10_2_norm)
loc_pred_conv10_2, cl_pred_conv10_2 = self.get_preds(conv10_2_norm_relu, 6)
# 合并各个层的预测结果
loc_preds = torch.cat([loc_pred_conv4_3, loc_pred_conv7,
loc_pred_conv8_2, loc_pred_conv9_2,
loc_pred_conv10_2], 1)
cl_preds = torch.cat([cl_pred_conv4_3, cl_pred_conv7,
cl_pred_conv8_2, cl_pred_conv9_2,
cl_pred_conv10_2], 1)
return loc_preds, cl_preds
get_preds方法用于在每个特征图上得到定位和分类的预测结果:
python
def get_preds(self, x, num_anchors):
batch_size = x.size(0)
# 创建预测结果的初始化梯度
loc_preds = torch.zeros(batch_size, num_anchors * 4,
x.size(2), x.size(3))
cl_preds = torch.zeros(batch_size, num_anchors * self.num_classes,
x.size(2), x.size(3))
for ndx in range(num_anchors):
loc_preds[..., ndx * 4:(ndx + 1) * 4] = self.loc_preds[ndx](x)
cl_preds[..., ndx * self.num_classes : (ndx + 1) * self.num_classes] = self.cl_preds[ndx](x)
return loc_preds, cl_preds在forward方法中,SSD利用不同尺度的特征图获得检测框和类预测,这些结果被拼接成最终的预测结果。get_preds方法在每个特征图上为不同尺度的先验框获得定位和分类预测。
SSD通过一张输入图像的一次前向传播就完成了检测预测,这也是其名称的由来。但它也面临匹配先验框与真实框的难度,以及特征融合的挑战。理解SSD的结构与实现可以帮助我们设计高效的检测模型。
SSD属于一类以检测为目标的单阶段检测方法,其与Faster R-CNN等两阶段方法相比,存在速度和精度的权衡。