ssd网络结构_SSD 源码实现详细解析 (PyTorch)-模型定义篇

最新推荐文章于 2024-05-09 11:59:43 发布

weixin_39989862

最新推荐文章于 2024-05-09 11:59:43 发布

阅读量1.7k

点赞数 2

文章标签： ssd网络结构

概览

SSD 和 YOLO 都是非常主流的 one-stage 目标检测模型, 并且相对于 two-stage 的 RCNN 系列来说, SSD 的实现更加的简明易懂, 接下来我将从以下几个方面展开对 SSD 模型的源码实现讲解: - 模型结构定义 - DefaultBox 生成候选框 - 解析预测结果 - MultiBox 损失函数 - Augmentations Trick - 模型训练 - 模型预测 - 模型验证 - 其他辅助代码

可以看出, 虽然 SSD 模型本身并不复杂, 但是也正是由于 one-stage 模型较简单的原因, 其检测的准确率相对于 two-stage 模型较低, 因此, 通常需要借助许多训练和检测时的 Tricks 来提升模型的精确度, 这些代码我们会放在第三部分讲解. 下面, 我们按照顺序首先对 SSD 模型结构定义的源码进行解析.(项目地址: https://github.com/amdegroot/ssd.pytorch)

模型结构定义

本部分代码主要位于 ssd.py 文件里面, 在本文件中, 定义了SSD的模型结构. 主要包含以下类和函数, 整体概览如下:

# ssd.py
class SSD(nn.Module): # 自定义SSD网络
    def __init__(self, phase, size, base, extras, head, num_classes):
        # ... SSD 模型初始化
    def forward(self, x):
        # ... 定义forward函数, 将设计好的layers和ops应用到输入图片 x 上

    def load_weights(self, base_file):
        # ... 加载参数权重值
def vgg(cfg, i, batch_norm=False):
    # ... 搭建vgg网络
def add_extras(cfg, i, batch_norm=False):
    # ... 向VGG网络中添加额外的层用于feature scaling
def multibox(vgg, extra_layers, cfg, num_classes):
    # ... 构建multibox结构
base = {
    ...} # vgg 网络结构参数
extras = {
    ...} # extras 层参数
mbox = {
    ...} # multibox 相关参数
def build_ssd(phase, size=300, num_classes=21):
    # ... 构建模型函数, 调用上面的函数进行构建

为了方便理解, 我们不按照文件中的定义顺序解析, 而是根据文件中函数的调用关系来从外而内, 从上而下的进行解析, 解析顺序如下: - build_ssd(...) 函数 - vgg(...) 函数 - add_extras(...) 函数 - multibox(...) 函数 - SSD(nn.Module) 类

build_ssd(...) 函数

在其他文件通常利用build_ssd(phase, size=300, num_classes=21)函数来创建模型, 下面先看看该函数的具体实现:

# ssd.py
class SSD(nn.Module): # 自定义SSD网络
    def __init__(self, phase, size, base, extras, head, num_classes):
        # ...
    def forward(self, x):
        # ...
    def load_weights(self, base_file):
        # ...
def vgg(cfg, i, batch_norm=False):
    # ... 搭建vgg网络
def add_extras(cfg, i, batch_norm=False):
    # ... 向VGG网络中添加额外的层用于feature scaling
def multibox(vgg, extra_layers, cfg, num_classes):
    # ... 构建multibox结构
base = {
     # vgg 网络结构参数
    '300': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M', 512, 512, 512],
    '500': []
}
extras = {
     # extras 层参数
    '300': [256, 'S', 512, 128, 'S', 256, 128, 256, 128, 256],
    '500': []
}
mbox = {
     # multibox 相关参数
    '300': [4, 6, 6, 6, 4, 4],
    '500': []
}
def build_ssd(phase, size=300, num_classes=21):
    # 构建模型函数, 调用上面的函数进行构建
    if phase != "test" and phase != "train": # 只能是训练或者预测阶段
        print("ERROR: Phase: " + phase + " not recognized")
        return
    if size != 300:
        print("ERROR: You specified size " + repr(size) + ". However, "+
                "currently only SSD300 is supported!") # 仅仅支持300size的SSD
        return
    base_, extras_, head_ = multibox(vgg(base[str(size)], 3),
            add_extras(extras[str(size), 1024),
            mbox[str(size)], num_classes )
    return SSD(phase, size, base_, extras_, head_, num_classes)

可以看到, build_ssd(...)函数主要使用了multibox(...)函数来获取base_, extras_, head_, 在调用multibox(...)函数的同时, 还分别调用了vgg(...)函数, add_extras(...)函数, 并将其返回值作为参数. 之后, 利用这些信息初始化了SSD网络. 那么下面, 我们就先查看一下这些函数定义和作用

vgg(...) 函数

我们以调用顺序为依据, 先对multibox(...)函数的内部实现进行解析, 但是在查看multibox(...)函数之前, 我们首先需要看看其参数的由来, 首先是vgg(...)函数, 因为 SSD 是以 VGG 网络作为 backbone 的, 因此该函数主要定义了 VGG 网络的结果, 根据调用语句vgg(base[str(size)], 3)可以看出, 调用vgg时向其传入了两个参数, 分别为base[str(size)] 和3, 对应的就是base['300']和3.

# ssd.py

def vgg(cfg, i, batch_norm = False):
    # cfg = base['300'] = [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)]
        if v == 'C':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
        else:
            conv2d = nn.Conv2d(in_channels=in_channels, out_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)
        conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6)
        conv7 = nn.Con2d(1024, 1024, kernel_size=1)
        layers += [pool5, conv6, nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)]
        return layers

上面的写法是 ssd.pytorch 代码中的原始写法, 代码风格体现了 PyTorch 灵活的编程特性, 但是这种写法不是那么直观, 需要很详细的解读才能看出来这个网络的整个结构是什么样的. 建议大家结合 VGG 网络的整个结构来解读这部分代码, 核心思想就是通过预定义的 cfg=base={...} 里面的参数来设置 vgg 网络卷积层和池化层的参数设置, 由于 vgg 网络的模型结构很经典, 有很多文章都写的很详细, 这里就不再啰嗦了, 我们主要来看一下 SSD 网络中比较重要的点, 也就是下面的 extras_layers.

add_extras(...) 函数

想必了解 SSD 模型的朋友都知道, SSD 模型中是利用多个不同层级上的 feature map 来进行同时进行边框回归和物体分类任务的, 除了使用 vgg 最深层的卷积层以外, SSD 还添加了几个卷积层, 专门用于执行回归和分类任务(如文章开头图2所示), 因此, 我们在定义完 VGG 网络以后, 需要额外定义这些新添加的卷积层. 接下来, 我们根据论文中的参数设置, 来看一下 add_extras(...) 的内部实现, 根据调用语句add_extras(extras[str(size)], 1024) 可知, 该函数中参数cfg = extras['300'], i=1024.

# ssd.py
def add_extras(cfg, i, batch_norm=False):
    # 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': # (1,3)[True] = 3, (1,3)[False] = 1
                layers += [nn.Conv2d(in_channels=in_channels, out_channels=cfg[k+1],
                                    kernel_size=(1, 3)[flag], stride=2, padding=1)]
            else:
                layers += [nn.Conv2d(in_channels=in_channels, out_channels=v,
                                    kernel_size=(1, 3)[flag])]
            flag = not flag
        in_channels = v
    return layers

注意, 在extras中, 卷积层之间并没有使用 BatchNorm 和 ReLU, 实际上, ReLU 的使用放在了forward函数中

同样的问题, 上面的定义不是很直观, 因此我将上面的代码用 PyTorch 重写了, 重写后的代码更容易看出网络的结构信息, 同时可读性也较强, 代码如下所示(与上面的代码完全等价):

def add_extras():
    exts1_1 = nn.Conv2d(in_channels=1024, out_channels=256, kernel_size=1)
    exts1_2 = nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, stride=2, padding=1)
    exts2_1 = nn.Conv2d(512, 128, 1, 1, 0)
    exts2_2 = nn.Conv2d(128, 256, 3, 2, 1)
    exts3_1 = nn.Conv2d(256, 128, 1, 1, 0)
    exts3_2 = nn.Conv2d(128, 256, 3, 1, 0)
    exts4_1 = nn.Conv2d(256, 128, 1, 1, 0)
    exts4_2 = nn.Conv2d(128, 256, 3, 1, 0)

    return [exts1_1, exts1_2, exts2_1, exts2_2, exts3_1, exts3_2, exts4_1, exts4_2]

在定义完整个的网络结构以后, 我们就需要定义最后的 head 层, 也就是特定的任务层, 因为 SSD 是 one-stage 模型, 因此它是同时在特征图谱上产生预测边框和预测分类的, 我们根据类别的数量来设置相应的网络预测层参数, 注意需要用到多个特征图谱, 也就是说要有多个预测层(原文中用了6个卷积特征图谱, 其中2个来自于 vgg 网络, 4个来自于 extras 层), 代码实现如下:

multibox(...) 函数

multibox(...) 总共有4个参数, 现在我们已经得到了两个参数, 分别是vgg(...)函数返回的layers, 以及add_extras(...)函数返回的layers, 后面两个参数根据调用语句可知分别为mbox[str(size)](mbox['300'])和num_classes(默认为21). 下面, 看一下multibox(...)函数的具体内部实现:

# ssd.py
def multibox(vgg, extra_layers, cfg, num_classes):
    # cfg = [4, 6, 6, 6, 4, 4]
    # num_classes = 21
    # ssd总共会选择6个卷积特征图谱进行预测, 分别为, vggnet的conv4_3, 以及extras_layers的5段卷积的输出(每段由两个卷积层组成, 具体可看extras_layers的实现).
    # 也就是说, loc_layers 和 conf_layers 分别具有6个预测层.
    loc_layers = []
    conf_layers = []
    vgg_source = [21, -2]
    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)]
    return vgg, extra_layers, (loc_layers, conf_layers)

同样, 我们可以将上面的代码写成可读性更强的形式:

# ssd.py
def multibox(vgg, extras, num_classes):
    loc_layers = []
    conf_layers = []
    #vgg_source=[21, -2] # 21 denote conv4_3, -2 denote conv7

    # 定义6个坐标预测层, 输出的通道数就是每个像素点上会产生的 default box 的数量
    loc1 = nn.Conv2d(vgg[21].out_channels, 4*4, 3, 1, 1) # 利用conv4_3的特征图谱, 也就是 vgg 网络 List 中的第 21 个元素的输出(注意不是第21层, 因为这中间还包含了不带参数的池化层).
    loc2 = nn.Conv2d(vgg[-2].out_channels, 6*4, 3, 1, 1) # Conv7
    loc3 = nn.Conv2d(vg

最低0.47元/天解锁文章

weixin_39989862

关注

2
点赞
踩
13

收藏

觉得还不错? 一键收藏
0
评论
ssd网络结构_SSD 源码实现详细解析 (PyTorch)-模型定义篇

概览SSD 和 YOLO 都是非常主流的 one-stage 目标检测模型, 并且相对于 two-stage 的 RCNN 系列来说, SSD 的实现更加的简明易懂, 接下来我将从以下几个方面展开对 SSD 模型的源码实现讲解: - 模型结构定义 - DefaultBox 生成候选框 - 解析预测结果 - MultiBox 损失函数 - Augmentations Trick - 模型训练 - 模...
复制链接

扫一扫