目标检测 SSD 算法结构解析以及基于 pytorch 代码复现（基于 VOC 数据集训练测试）

  SSD 作为目标检测单阶段方法中最经典的算法之一，是大多数人入门的第一选择，网上参考资料也很多，本文旨在记录自己的学习过程，希望对其他的初学者能够有所帮助。

一、SSD 基础知识

  目标检测单阶段方法一般先对图像进行特征提取，然后再根据提取的特征进行分类与边框回归。SSD 也是如此，比较独特的是，它在 VGG16 基础上通过添加额外的不同尺度特征提取层实现检测性能的提升。具体结构如下图：

共提取出 6 个不同尺度的特征图，分别为（38，38），（19，19），（10，10），（5，5），（3，3），（1，1），从而实现检测。（这样设计的目的是不同尺度的特征图包含的信息不同，一般浅层感受野较小，包含的位置信息较好，并且有利于小目标的检测，高层一般感受野较大，包含更多的全局信息）不过 SSD 只是简单的利用了不同尺度信息，后面也有人融合了特征金字塔思想，提出了 FSSD ，提升了检测性能。    第二个比较核心的点是 SSD 采用卷积操作实现最后的检测而不是使用全连接层，采用 3*3 卷积核，产生类别的分数或产生相对于默认框坐标的形状偏移。通过卷积操作，减少了参数，从而达到了比 YOLO 更快的速度。   最后一点是 SSD 预先设置了先验框，每一个目标设置四个先验框，如下图所示。其中包含中心点坐标以及长、宽（w,h）四个参数。可以看到大尺度特征图更适合检测小目标，小尺度特征图适合检测大一些的目标。   以上是 SSD 的核心设计思路以及整体结构。采用的损失函数思想是局部化损失（例如平滑 L1）和置信损失（例如 Softmax）之间的加权和。

二、pytorch 代码复现

  在有了以上知识的基础上，下面介绍基于 pytorch 的代码复现。原论文基于 VOC 以及 COCO 数据集均做了实验，由于 COCO 数据集太大，本文仅基于 VOC 数据集进行训练及验证。

pytorch代码链接：https://github.com/amdegroot/ssd.pytorch一个较好的教程：https://www.jianshu.com/p/04ea90ebfd95，主要根据这个教程对代码做一些更改，以兼容高版本pytorch

（1）代码复现过程   下面是一些补充知识以及遇到的问题如何解决：   基本配置为 python+pytorch1.2.0，首先下载 VOC 数据集,可通过网盘链接自行下载：

链接：https://pan.baidu.com/s/1ZOcEwE5cbDgb_1R5FocAWg 提取码：mwh4

  下载好以后可以在代码工程 data\scripts 文件夹下创建名为 VOC2007 的文件夹，并把训练集 VOC2007 下的文件拷贝过来，并在代码中更改数据集路径 （参考上面的教程，Linux系统采用源代码指令即可）   根据教程改一些代码，然后根据自己电脑修改初始参数（RTX1660Ti 最高批次尺寸可设为25，推荐使用16，若使用源代码中32，损失函数会变为无穷大；RTX2080Ti 可设为64）。然后运行 train.py，代码每迭代5000次会保存一次权重，可根据实际情况停止，官方训练迭代了40K次。

注意：1. 如果上面这句报错，可去掉.double()    2. 损失函数无穷大，可考虑调低学习率

  训练完成后，可更改 eval.py 中权重名称，并运行，验证结果。官方 demo 给了视频检测，这里贴一个图像检测的 demo 程序

from __future__ import print_function import torch from torch.autograd import Variable from matplotlib import pyplot as plt import cv2 import time import argparse import sys from os import path sys.path.append(path.dirname(path.dirname(path.abspath(__file__)))) from data import BaseTransform, VOC_CLASSES as labelmap from ssd import build_ssd weight_path = 'D:\workspace\pytorch\project\ssd.pytorch\ssd.pytorch\weights\ssd300_VOC07_50000.pth' parser = argparse.ArgumentParser(description='Single Shot MultiBox Detection') parser.add_argument('--weights', default=weight_path, type=str, help='Trained state_dict file path') parser.add_argument('--cuda', default=False, type=bool, help='Use cuda in live demo') args = parser.parse_args() COLORS = [(255, 0, 0), (0, 255, 0), (0, 0, 255)] FONT = cv2.FONT_HERSHEY_SIMPLEX def cv2_demo(net, transform): def predict(frame): height, width = frame.shape[:2] x = torch.from_numpy(transform(frame)[0]).permute(2, 0, 1) x = Variable(x.unsqueeze(0)) y = net(x) # forward pass detections = y.data # scale each detection back up to the image scale = torch.Tensor([width, height, width, height]) for i in range(detections.size(1)): j = 0 while detections[0, i, j, 0] >= 0.6: score = float(detections[0, i, j, 0]) pt = (detections[0, i, j, 1:] * scale).cpu().numpy() cv2.rectangle(frame, (int(pt[0]), int(pt[1])), (int(pt[2]), int(pt[3])), COLORS[i % 3], 2) cv2.putText(frame, labelmap[i - 1] + '_' + str(score)[:4], (int(pt[0]), int(pt[1])), FONT, 1, (255, 0, 255), 2, cv2.LINE_AA) j += 1 return frame #for i in range() frame = cv2.imread("car.jpg") frame = predict(frame) IMAGE_SIZE = (12, 8) plt.figure(figsize=IMAGE_SIZE) plt.imshow(frame) plt.show() if __name__ == '__main__': net = build_ssd('test', 300, 21) # initialize SSD net.load_state_dict(torch.load(args.weights)) transform = BaseTransform(net.size, (104 / 256.0, 117 / 256.0, 123 / 256.0)) cv2_demo(net.eval(), transform)

frame = cv2.imread(“car.jpg”)    #这句读自己要测试的图片 while detections[0, i, j, 0] >= 0.6  #这句可更改阈值，若无检测结果，可降低阈值重新检测。

（2）代码复现结果

有其他问题欢迎留言交流。