ssd.pytorch源码分析（一）— 网络结构
一、总览
在ssd的原论文中，采用的backbone为VGG16。VGG16定义如下图C这一列。从上到下依次为：conv1_2、conv2_2、conv3_3、conv4_3、conv5_3、fc6、fc7、fc8。其中_n指有n层。可以看到参数层为2+2+3+3+3+1+1+1=16层。
SSD作者在原VGG16的基础上进行了改进：将原来的FC7改为Conv7，并增加卷积层深度，即继续添加Conv8_2, Conv9_2, Conv10_2, Conv11_2, 如下图所示。其中s指的是stride。
原文中提到，SSD将bounding box的输出空间离散化成一系列的default boxes（可以理解为faster rcnn中的anchor）。这些default boxes由不同的层输出，最终汇总在一起进行nms。这样结合不同尺度的特征图的boxes，可以有效处理检测目标的多尺度问题。具体来说，就是选取上图中标黄部分的特征图，对每个特征图再进行特定的卷积操作，使得特征图每个default box输出（num_classes+4）个值，分别代表了分类和定位预测。
二、vgg16结构
"""
input-> conv1_2:3x3x64,relu,3x3x64,relu -> pool1 ->
conv2_2:3x3x128,relu,3x3x128,relu -> pool2 ->
conv3_3:3x3x256,relu,3x3x256,relu,3x3x256,relu -> pool3 ->
conv4_3:3x3x512,relu,3x3x512,relu,3x3x512,relu -> pool4 ->
conv5_3:3x3x512,relu,3x3x512,relu,3x3x512,relu -> pool5 ->
conv6:3x3x1024 atrous，relu->
conv7:1x1x1024，relu
"""
def vgg(cfg, i, batch_norm=False):
"""
cfg = [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M',
512, 512, 512] #各卷积层通道数
i = 3 #输入图像通道数
"""
layers = []
in_channels = i
for v in cfg:
if v == 'M':
layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
elif v == 'C':
layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
else:
conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
if batch_norm:
layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
else:
layers += [conv2d, nn.ReLU(inplace=True)]
in_channels = v
pool5 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1) # change pool5 from 2 × 2 − s2 to 3 × 3 − s1
conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6) # atrous algorithm
conv7 = nn.Conv2d(1024, 1024, kernel_size=1)
layers += [pool5, conv6,
nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)]
return layers #分析后可知len(layers)=35
注意此处vgg函数定义的不是原始vgg的结构，而是如上文所述进行适当修改，定义了conv1_2、conv2_2、conv3_3、conv4_3、conv5_3、conv6、conv7。可以看到，函数的返回值layers为一个list，用于后续提取特征图之用。
三、额外层定义
vgg函数中只定义至conv7，add_extras函数将定义后面的conv8_2, conv9_2, conv10_2, conv11_2。
"""
feature -> conv8_2:1x1x256,3x3x512s2 -> conv9_2:1x1x128,3x3x256s2 ->
conv10_2:1x1x128,3x3x256s1 -> conv11_2:1x1x128,3x3x256s1
"""
def add_extras(cfg, i):
# Extra layers added to VGG for feature scaling
# add_extras函数中没有池化操作而是改用了stride=2来减少分辨率
"""
cfg = [256, 'S', 512, 128, 'S', 256, 128, 256, 128, 256]
i = 1024
"""
layers = []
in_channels = i
flag = False
for k, v in enumerate(cfg):
if in_channels != 'S':
if v == 'S':
layers += [nn.Conv2d(in_channels, cfg[k + 1],
kernel_size=(1, 3)[flag], stride=2, padding=1)]
else:
layers += [nn.Conv2d(in_channels, v, kernel_size=(1, 3)[flag])]
flag = not flag
in_channels = v
return layers #分析后可知len(layers)=8
同论文中的网络结构比对，可以知道后面的extra部分已被正确定义了。其实很多代码都是串行地定义网络结构，不过这里用的是循环，每次读入cfg不同的值，这样的定义方式一定程度上减少了代码量，对于大型深度网络值得借鉴。同样地，函数的输出layers也为一个list，用于后续提取特征图之用。
四、multibox定义
multibox即将选中的各个层的特征图提取出来（选中的层都在前面用黄色标出），然后对每个特征图的每个像素点预测目标box相对于default box的offsets和每个类的置信度。具体来说，如果某一层每个像素点定义k个default boxes，有num_classes个类，则该像素点总的输出个数为k(num_classes+4)。不过在代码中，网络分为2个head（loc_layers和conf_layers）分别预测位置信息和类别信息。
def multibox(vgg, extra_layers, cfg, num_classes):
"""
选取的特征图在输入的两个list(vgg和add_extras)中的索引：
vgg：21，-2
conv4_3去掉relu的末端；conv7relu之前的1x1卷积；
add_extras：1，3，5，7
conv8_2末端；conv9_2末端；conv10_2末端；conv11_2末端
cfg = [4, 6, 6, 6, 4, 4] # 每个特征图中各个像素定义的default boxes数量
"""
loc_layers = []
conf_layers = []
vgg_source = [21, -2] #21：conv4_3中最后一个3x3x512 -2：conv7relu之前的1x1卷积
for k, v in enumerate(vgg_source):
loc_layers += [nn.Conv2d(vgg[v].out_channels,
cfg[k] * 4, kernel_size=3, padding=1)]
conf_layers += [nn.Conv2d(vgg[v].out_channels,
cfg[k] * num_classes, kernel_size=3, padding=1)]
for k, v in enumerate(extra_layers[1::2], 2):
loc_layers += [nn.Conv2d(v.out_channels, cfg[k]
* 4, kernel_size=3, padding=1)]
conf_layers += [nn.Conv2d(v.out_channels, cfg[k]
* num_classes, kernel_size=3, padding=1)]
#输出维度结果：
#令C=[4,6,6,6,4,4]
#loc_layers:6x[batch, C_i*4, H_i, W_i]
#conf_layers:6x[batch, C_i, H_i, W_i]
return vgg, extra_layers, (loc_layers, conf_layers)
五、前向传播
通过使用上面三个部分，就可以定义出完整的SSD网络结构了。
class SSD(nn.Module):
def __init__(self, phase, size, base, extras, head, num_classes):
#phase:"train"/"test"; size:输入图像尺寸，300;
#base, extras, head:分别为上文中三个函数的输出
super(SSD, self).__init__()
self.phase = phase
self.num_classes = num_classes
self.cfg = (coco, voc)[num_classes == 21]
self.priorbox = PriorBox(self.cfg) #默认框的获取,将在其他博客中分析
self.priors = Variable(self.priorbox.forward(), volatile=True)
self.size = size
# SSD network
self.vgg = nn.ModuleList(base)
# Layer learns to scale the l2 normalized features from conv4_3
self.L2Norm = L2Norm(512, 20)
self.extras = nn.ModuleList(extras)
self.loc = nn.ModuleList(head[0])
self.conf = nn.ModuleList(head[1])
if phase == 'test':
self.softmax = nn.Softmax(dim=-1)
self.detect = Detect(num_classes, 0, 200, 0.01, 0.45)
forward函数：
输入值：
——x: 输入图. Shape: [batch,3,300,300]；
返回值：
——根据phase输出不同的结果
(num_priors=所有层金字塔的像素点之和)
def forward(self, x):
sources = list() #6张特征图
loc = list() #所有默认框的位置预测结果，列表中一个元素对应一张特征图
conf = list() #所有默认框的分类预测结果，列表中一个元素对应一张特征图
# 前向传播vgg至conv4_3 relu 得到第1个特征图
for k in range(23):
x = self.vgg[k](x)
s = self.L2Norm(x)
sources.append(s)
# 继续前向传播vgg至fc7得到第2个特征图
for k in range(23, len(self.vgg)):
x = self.vgg[k](x)
sources.append(x)
# 在extra layers中前向传播得到另外4个特征图
for k, v in enumerate(self.extras):
x = F.relu(v(x), inplace=True)
if k % 2 == 1:
sources.append(x)
# 将各个特征图中的定位和分类预测结果append进列表中
for (x, l, c) in zip(sources, self.loc, self.conf):
loc.append(l(x).permute(0, 2, 3, 1).contiguous()) #6*(N,C,H,W)->6*(N,H,W,C) C=k*4
conf.append(c(x).permute(0, 2, 3, 1).contiguous()) #6*(N,C,H,W)->6*(N,H,W,C) C=k*num_class
loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1) #[N,-1]
conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1) #[N,-1]
if self.phase == "test":
#如果是测试阶段需要对定位和分类的预测结果进行分析得到最终的预测框
output = self.detect(
loc.view(loc.size(0), -1, 4), # loc preds ->[N,num_priors,4]
self.softmax(conf.view(conf.size(0), -1,
self.num_classes)), # conf preds [N,num_priors,num_classes] 最后一维softmax
self.priors.type(type(x.data)) # default boxes [num_priors,4] 4:[cx,cy,w,h]
) #output: [N,num_classes,num_remain*5]
else:
#如果是训练阶段则直接输出定位和分类预测结果以计算损失函数
output = (
loc.view(loc.size(0), -1, 4), #[N,num_priors,4]
conf.view(conf.size(0), -1, self.num_classes), #[N,num_priors,num_classes]
self.priors #[num_priors,4]
)
return output