之前，对SSD的论文进行了解读，可以回顾之前的博客：https://www.cnblogs.com/dengshunge/p/11665929.html。
为了加深对SSD的理解，因此对SSD的源码进行了复现，主要参考的github项目是ssd.pytorch。
搭建SSD的项目，可以分成以下三个部分：
- 数据读取；
- 网络搭建；
- 损失函数的构建。
接下来，本篇博客重点分析网络搭建。
该部分整体比较简单，思路也很清晰。
首先，在train.py中，网络搭建的函数入口是函数build_ssd()，该函数需要传入以下几个参数："train"或者"test"字符串、图片尺寸、类别数。其中，"train"或者"test"字符串用于区分该网络是用于训练还是测试，这两个阶段的网络有些许不同，本文主要将训练阶段的网络；而类别数需要加上背景，对于VOC而言，有20个类别，加上1个背景，即类别数是21。
ssd_net = build_ssd('train', voc['min_dim'], voc['num_classes'])
这里，先放一张SSD的网络结构图，可以看出，SSD网络是有3部分组成的，vgg主干网络，新增网络（Conv6之后的层）和用于检测的头部网络（Extra Feature Layers）。
接着，在ssd.py中，首先定了一个参数，如下所示。这里主要以SSD300为例。这些参数有什么用呢？字典base的参数指的是用于搭建VGG主干网络输出通道数，其中“M”表示需要进行maxpooling；字典extras的参数同样表示新增层的输出通道数，其中“S”表示需要stride=2的降采样；字典mbo的参数表示用于特征融合的层中，每个层对应未知(x,y)的锚点框数量，在SSD300中，使用了6个层进行特征融合，如Conv_4层中，每个位置使用4个锚点框进行预测。
base = {
'300': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M',
512, 512, 512], # M表示maxpolling
'512': [],
}
extras = {
'300': [256, 'S', 512, 128, 'S', 256, 128, 256, 128, 256], # S表示stride=2
'512': [],
}
mbox = {
'300': [4, 6, 6, 6, 4, 4], # 每个特征图的每个点，对应锚点框的数量
'512': [],
}
当定义完需要使用到的参数后，可以进行如具体搭建的环节。函数build_ssd()的定义如下所示。利用函数multibox()来构建SSD网络的各个部分，分别是VGG主干网络，新增层和用于检测的头部网络（或许可以理解为分类头和回归头）。而VGG主干网络是通过函数vgg()来实现，新增层是通过函数add_extras()来实现，而函数multibox()则搭建用于检测的头部网络。最后用这些层来初始化类SSD。
def build_ssd(phase, size=300, num_class=21):
if phase != 'test' and phase != 'train':
raise ("ERROR: Phase: " + phase + " not recognized")
base_, extras_, head_ = multibox(vgg(base[str(size)]),
add_extras(extras[str(size)], in_channels=1024),
mbox[str(size)],
num_class)
return SSD(phase, size, base_, extras_, head_, num_class)
我们来看一下VGG主干网络是如何搭建的。函数vgg()需要将上述的base字典传入进去，根据base字典，来搭建卷积层和池化层。作者对vgg网络进行了改进，即将fc6和fc7更改成conv6和conv7。值得留意的是，在conv6中，使用了空点卷积，dilation=6，增大感受野。在SSD论文的最后，也讨论了空洞卷积对结果有好的影响。最后，将这些卷积层和池化层放入list中，并返回这个list。
def vgg(cfg=base['300'], batch_norm=False):
'''
该函数来源于torchvision.models.vgg19()中的make_layers()
'''
layers = []
in_channels = 3
# vgg主体部分,到论文的conv6之前
for v in cfg:
if v == 'M':
# ceil_mode是向上取整
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, ceil_mode=True)
conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6)
conv7 = nn.Conv2d(1024, 1024, kernel_size=1)
layers += [pool5, conv6, nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)]
return layers
接下来，我们来了解一下SSD在vgg中新增层，即conv7之后的网络层。同样，函数add_extras()需要传入字典extras，来构建网络层。这里可以留意一下，kernel_size的写法，(1,3)为一个元祖tuple，flag来控制取哪个值，即可变换使用3*3或者1*1的卷积核，减少代码的冗余。最后将新构建的层存入list中，并返回这个list。
def add_extras(cfg=extras['300'], in_channels=1024):
'''
完成SSD后半部分的网络构建，即作者新加上去的网络,从conv7之后到conv11_2
'''
layers = []
flag = False # 交替控制卷积核,使用1*1或者使用3*3
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
当有了vgg的主干网络和新增层后，可以将某些层进行特征融合和预测了。这里，就需要使用到函数multibox()。需要将vgg主干网络和新增层的list、字典mbox和类别数传入函数中。首先，函数multibox()会创建两个list，用于保存位置回归的层和置信度的层。对于每个用于融合的特征层，会分成两部分，一个用于回归，使用3*3的卷积，输出通道数是cfg[k] * 4，其中cfg[k]表示每个位置上锚点框的数量，4表示[x_min,y_min,x_max,y_max]；另外一个用于类别的判断，也是使用3*3的卷积，输出通道数是cfg[k] * num_class，表示每个锚点框判断其属于哪一个类别，在voc中，num_class=21（包含背景）。可以理解成将此特征层分成了分类头和回归头，每个锚点框会输出4个坐标和21个类别置信度。最后将vgg主干网络、新增层、分类头和回归头返回。
def multibox(vgg, extra_layers, cfg, num_class):
'''
返回vgg网络,新增网络,位置网络和置信度网络
'''
loc_layers = [] # 判断位置
conf_layers = [] # 判断置信度
vgg_source = [21, -2] # 21表示conv4_3的序号,-2表示conv7的序号
for k, v in enumerate(vgg_source):
# vgg[v]表示需要提取的特征图
# cfg[k]代表该特征图下每个点对应的锚点框数量
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_class, kernel_size=3, padding=1)]
for k, v in enumerate(extra_layers[1::2], 2):
# [1::2]表示,从第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_class, kernel_size=3, padding=1)]
return vgg, extra_layers, (loc_layers, conf_layers)
函数multibox()返回的各个层，用于初始化类SSD。首先，由于因为“train”阶段和“test”阶段是有点区别的，本节依然主要将“train”阶段，因此，需要传入phase参数，参数只能是两个值(train,test)。函数PriorBox()的作用是来创建先验锚点框，返回的shape为[8732,4]，其中具有8732个锚点框，4表示每个锚点框的坐标[中心点x，中心点y，宽，高]，这里的坐标值有点不太一样。由于传入的网络层是以list列表的形式，因此，用nn.ModuleList()将其转换为pytorch的网络结构。
接下来看类SSD中的函数forward()，用于前向推理。按顺序对输入图片进行处理，在conv4中，需要对特征图进行L2正则化。并将用于特征融合的特征图存在放sources中。在得到5个用于融合的特征图后，将这些特征图输入到分类头和回归头中，每个特征图对应各自的分类头和回归头。这里注意一下，分类头或者回归头卷积后，使用了permute()函数。该函数的作用是交换维度，原本的维度是[batch_size,channel,height,weight]，交换维度后变成了[batch_size,height,weight,channel]，这样做的目的是方便后续的处理。将处理后的结果保存在loc和conf这两个List中。后续接着对loc和conf进行变换，利用view()函数，最终，loc的shape为[batch_size,8732*4]，conf的shape为[batch_size,8732*21]。
最后，将loc和conf这两个List又变换维度，返回出去，用于计算loss损失函数（感觉这么多变换，有点重复呀，应该可以省略一部分）。"train"阶段和"test"阶段返回的结果类似，其中不同点是，在test阶段，置信度需要经过softmax。
class SSD(nn.Module):
'''
构建SSD的主函数，将base(vgg)、新增网络和位置网络与置信度网络组合起来
'''
def __init__(self, phase, size, base, extras, head, num_classes):
super(SSD, self).__init__()
self.phase = phase
self.num_classes = num_classes
self.priors = torch.Tensor(PriorBox(voc))
self.size = size
# SSD网络
self.vgg = nn.ModuleList(base)
# 对conv4_3的特征图进行L2正则化
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=self.num_classes, top_k=200,
conf_thresh=0.01, nms_thresh=0.45)
def forward(self, x):
sources = [] # 保存特征图
loc = [] # 保存每个特征图进行位置网络后的信息
conf = [] # 保存每个特征图进行置信度网络后的信息
# 处理输入至conv4_3
for k in range(23):
x = self.vgg[k](x)
# 对conv4_3进行L2正则化
s = self.L2Norm(x)
sources.append(s)
# 完成vgg后续的处理
for k in range(23, len(self.vgg)):
x = self.vgg[k](x)
sources.append(x)
# 使用新增网络进行处理
for k, v in enumerate(self.extras):
x = F.relu(v(x), inplace=True)
if k % 2 == 1:
sources.append(x)
# 将特征图送入位置网络和置信度网络
# l(x)或者c(x)的shape为[batch_size,channel,height,weight],使用了permute后,变成[batch_size,height,weight,channel]
# 这样做应该是为了方便后续处理
for (x, l, c) in zip(sources, self.loc, self.conf):
loc.append(l(x).permute(0, 2, 3, 1).contiguous())
conf.append(c(x).permute(0, 2, 3, 1).contiguous())
# 进行格式变换
loc = torch.cat([o.view(o.size(0), -1) for o in loc], 1) # [batch_size,34928],锚点框的数量8732*4=34928
conf = torch.cat([o.view(o.size(0), -1) for o in conf], 1)
if self.phase == 'train':
output = (loc.view(loc.size(0), -1, 4), # [batch_size,num_priors,4]
conf.view(conf.size(0), -1, self.num_classes), # [batch_size,num_priors,21]
self.priors) # [num_priors,4]
else: # Test
output = self.detect(
loc.view(loc.size(0), -1, 4), # 位置预测
self.softmax(conf.view((conf.size(0), -1, self.num_classes))), # 置信度预测
self.priors.cuda() # 先验锚点框
)
return output
在上面类SSD中，提及到了先验锚点框的构建函数PriorBox()，这个函数在models/prior_box.py中。首先，根据用于融合的特征图尺寸和product()函数，生成一系列的点，如(0,0),(0,1),(0,2)等。然后根据这些像素点位置，偏移0.5作为锚点框的中心点，即cx和cy，并将其归一化。然后计算论文中的$s_k$和${s_k}'$，对应s_k和s_k_prime，先计算$a_r=1$的情况，再计算其余$a_r$的情况。此时，mean的shape为[1,34928]，因此，需要使用view()函数，将其切割出来，变成[8732,4]。记得，这里的锚点框的坐标是[中心点x，中心点y，宽，高]。
def PriorBox(cfg):
'''
为所有特征图生成预设的锚点框，返回所有生成的锚点框，尺寸为[8732,4]，
每行表示[中心点x，中心点y，宽，高]
'''
image_size = cfg['min_dim'] # 300
feature_maps = cfg['feature_maps'] # [38, 19, 10, 5, 3, 1],特征图尺寸
steps = cfg['steps'] # [8, 16, 32, 64, 100, 300]
min_sizes = cfg['min_sizes'] # [30, 60, 111, 162, 213, 264]
max_sizes = cfg['max_sizes'] # [60, 111, 162, 213, 264, 315]
aspect_ratios = cfg['aspect_ratios'] # [[2], [2, 3], [2, 3], [2, 3], [2], [2]]
mean = []
# 为所有特征图生成锚点框
for k, f in enumerate(feature_maps):
# product(list1,list2)的作用是依次取出list1中的每1个元素，与list2中的每1个元素，
# 组成元组，然后，将所有的元组组成一个列表，返回
# 而这里使用了repeat,说明1个list重复2次
for i, j in product(range(f), repeat=2):
f_k = image_size / steps[k]
# 计算中心点,这里的j是沿x方向变化的
cx = (j + 0.5) / f_k
cy = (i + 0.5) / f_k
# aspect_ratio=1有两种情况,s_k=s_k,s_k=sqrt(s_k*s_(k+1))
s_k = min_sizes[k] / image_size
mean += [cx, cy, s_k, s_k]
s_k_prime = sqrt(s_k * (max_sizes[k] / image_size))
mean += [cx, cy, s_k_prime, s_k_prime]
# 剩余的aspect_ratio
for ar in 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)]
# 此时的mean是1*34928的list，要4个数就分割出来，所以需要用view,从而变成[8732,4],即有8732个锚点框
output = torch.Tensor(mean).view(-1, 4)
if cfg['clip']:
# 对每个元素进行截断限制,限制为[0,1]之间
output.clamp_(min=0, max=1)
return output
最后，类SSD中还对conv4的特征层使用了L2正则化，该函数在models/l2norm.py中。在函数forwand()中，按每个通道对其值进行L2正则化，即除以通道的平方根来实现归一化。
class L2Norm(nn.Module):
'''
对conv4_3进行l2归一化
'''
def __init__(self, n_channels, scale):
super(L2Norm, self).__init__()
self.n_channels = n_channels
self.gamma = scale
self.eps = 1e-10
self.weight = nn.Parameter(torch.Tensor(self.n_channels)) # n_channels个随机数
self.reset_parameters()
def reset_parameters(self):
# 使用gamma来填充weight的每个值
nn.init.constant_(self.weight, self.gamma)
def forward(self, x):
# 按通道进行求值
norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps # [1,1,38,38]
x = torch.div(x, norm)
# 将weight通过3个unsqueeze展开成[1,512,1,1],然后通过expand_as进行扩展,形成[1,512,38,38]
out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(x) * x
return out
至此，SSD的网络搭建过程已经完成了，通过类SSD的forward()函数，即能返回预测框的坐标和类别置信度，以此可以计算损失函数。