在 TorchVision v0.10 中,我们发布了两个基于 SSD 架构的全新目标检测模型。我们计划通过两篇系列文章,详细介绍这些算法的关键实现细节以及它们的训练过程。
在本系列的第一部分,我们将重点介绍 SSD 论文 Single Shot MultiBox Detector 中描述的原始算法实现。我们将简要概述该算法的工作原理,然后深入探讨其主要组件,突出代码中的关键部分,最后讨论我们如何训练发布的模型。我们的目标是涵盖复现该模型所需的所有细节,包括那些论文中未提及但属于原始实现一部分的优化。
SSD 是如何工作的?
强烈建议阅读上述论文,但这里有一个简单且过度简化的复习。我们的目标是检测图像中物体的位置及其所属类别。以下是 SSD 论文 中展示模型预测示例的图 5:

SSD 算法使用一个 CNN 主干网络(Backbone),将输入图像传入其中,并提取网络不同层级的卷积输出。这些输出列表被称为特征图(feature maps)。随后,这些特征图会被送入分类(Classification)和回归(Regression)头,由它们负责预测边界框的类别和位置。
由于每个图像的特征图包含来自网络不同层级的输出,其尺寸各异,因此能够捕捉到不同尺寸的物体。在每个特征图之上,我们平铺了多个默认框(default boxes),可以将其视为我们初步的先验猜测。对于每个默认框,我们预测是否存在物体(及其类别)以及它的偏移量(对原始位置的修正)。在训练阶段,我们需要先将真实标注(ground truth)与默认框进行匹配,然后利用这些匹配来计算损失。在推理阶段,相似的预测框会被合并以估计最终结果。
SSD 网络架构
在本节中,我们将讨论 SSD 的关键组件。我们的代码紧密遵循该论文,并使用了官方实现中包含的许多未记录的优化。
DefaultBoxGenerator
DefaultBoxGenerator 类 负责生成 SSD 的默认框,其运作方式与 FasterRCNN 的 AnchorGenerator 相似(关于它们的区别,请参阅论文第 4-6 页)。它产生一组具有特定宽度和高度的预定义框,这些框平铺在图像上,作为物体可能位置的第一批粗略先验猜测。以下是 SSD 论文中关于真实标注和默认框可视化的图 1:

该类由一组控制形状和平铺方式的超参数进行参数化。对于那些想要尝试新主干网络/数据集的用户,该实现提供了自动生成的良好猜测,用户也可以传入优化的自定义值。
SSDMatcher
SSDMatcher 类 扩展了 FasterRCNN 使用的标准 Matcher,负责将默认框与真实标注匹配。在计算所有组合的 IoU 后,我们使用该匹配器为每个默认框找到重叠度高于 IoU 阈值 的最佳候选真实标注。SSD 版本的匹配器有一个额外的步骤,以确保每个真实标注都与具有最高重叠度的默认框匹配。匹配器的结果将用于模型训练过程中的损失估计。
分类和回归头
SSDHead 类 负责初始化网络的分类和回归部分。以下是其代码的一些显著细节:
- 分类头和回归头都继承自同一个基类,该基类负责对每个特征图进行预测。
- 特征图的每一层都使用单独的 3×3 卷积来估计类别逻辑值 (logits) 和边界框位置。
- 每个头在每一层所做的预测数量取决于默认框的数量和特征图的尺寸。
主干特征提取器
特征提取器通过额外层对标准 VGG 主干进行了重构和增强,如 SSD 论文图 2 所示。

该类支持 TorchVision 中的所有 VGG 模型,用户也可以为其他类型的 CNN 创建类似的提取器类(参见此 ResNet 示例)。以下是该类的一些实现细节:
- 对第 3 层 Maxpool 的
ceil_mode参数进行修改是获得与论文相同特征图尺寸的必要操作。这是由于 PyTorch 与该模型原始 Caffe 实现之间的细微差异所致。 - 它在 VGG 之上添加了一系列额外特征层。如果在构造过程中
highres参数为True,它将附加一个额外的卷积。这对于 SSD512 版本模型非常有用。 - 如论文第 3 节所述,原始 VGG 的全连接层被转换为卷积,其中第一个使用了空洞卷积 (Atrous)。此外,maxpool5 的步长和卷积核大小也进行了修改。
- 如第 3.1 节所述,L2 正则化应用于 conv4_3 的输出,并且引入了一组可学习的权重来控制其缩放比例。
SSD 算法
实现中的最后一个关键部分是 SSD 类。以下是一些显著细节:
- 该算法通过一组与其它检测模型相似的参数进行参数化。必须的参数包括:负责估计特征图的主干网络,应为
DefaultBoxGenerator类配置好的实例的anchor_generator,输入图像将被调整到的尺寸,以及分类时的num_classes(不包含背景类)。 - 如果未提供头,构造函数将初始化默认的
SSDHead。为此,我们需要知道主干网络产生的每个特征图的输出通道数。我们首先尝试从主干网络检索此信息,如果不可用,我们将动态估计它。 - 该算法重用了其它检测模型使用的标准 BoxCoder 类。该类负责编码和解码边界框,并配置为使用与原始实现相同的先验方差。
- 虽然我们重用了标准 GeneralizedRCNNTransform 类 来调整和归一化输入图像,但 SSD 算法对其进行了配置,以确保图像尺寸保持固定。
以下是该实现的两个核心方法:
compute_loss方法估计了 SSD 论文第 5 页描述的标准 Multi-box 损失。它使用 平滑 L1 损失 (smooth L1 loss) 进行回归,使用带有 难例挖掘 (hard-negative sampling) 的标准 交叉熵损失 进行分类。- 与所有检测模型一样,
forward方法在训练模式和评估模式下的行为有所不同。它首先调整并归一化输入图像,然后将其通过主干网络获取特征图。接着,特征图被传入头网络得到预测结果,随后方法会生成默认框。- 如果模型处于训练模式,
forward将估计默认框与真实标注的 IoU,使用SSDmatcher生成匹配,最后通过调用compute_loss方法估计损失。 - 如果模型处于评估模式,我们首先通过保留通过分数阈值的检测结果来选择最佳预测,挑选出最有希望的框,并运行 NMS(非极大值抑制)来清理和筛选最佳预测。最后,我们对预测结果进行后处理,将其调整回原始图像尺寸。
- 如果模型处于训练模式,
SSD300 VGG16 模型
SSD 是一个模型家族,因为它可以通过不同的主干和不同的头配置进行设置。在本节中,我们将重点介绍提供的 SSD 预训练模型。我们将讨论其配置细节以及用于复现报告结果的训练过程。
训练过程
该模型使用 COCO 数据集进行训练,其所有超参数和脚本均可在我们的 references 文件夹中找到。下面我们提供训练过程中最引人注目的细节。
论文超参数
为了在 COCO 上取得尽可能好的结果,我们采用了论文第 3 节中描述的有关优化器配置、权重正则化等的超参数。此外,我们发现采用官方实现中有关 DefaultBox 生成器平铺配置的优化也非常有帮助。这种优化在论文中未提及,但对于提高小目标的检测精度至关重要。
数据增强
实现论文第 6 页和第 12 页描述的 SSD 数据增强策略对于复现结果至关重要。具体而言,使用随机“放大 (Zoom In)”和“缩小 (Zoom Out)”变换使得模型对各种输入尺寸具有鲁棒性,并提高了对中小型物体的检测精度。最后,由于 VGG16 拥有相当多的参数,增强中包含的光度失真具有正则化作用,有助于避免过拟合。
权重初始化与输入缩放
我们发现遵循论文提出的权重初始化方案也是有益的。为此,我们不得不调整我们的输入缩放方法:撤销由 ToTensor() 执行的 0-1 缩放,并使用带有此缩放比例的预训练 ImageNet 权重(感谢 Max deGroot 在他的仓库中提供这些权重)。所有新卷积层的权重都使用 Xavier 初始化,偏置设置为零。初始化后,网络进行了端到端训练。
学习率 (LR) 方案
正如论文中所报告的,在应用激进的数据增强后,有必要更长时间地训练模型。我们的实验证实了这一点,我们不得不调整学习率、批次大小和总体步骤以获得最佳结果。我们提出的学习方案配置得较为保守,显示出在步骤间趋于平稳的迹象,因此只需执行我们 epoch 总数 66% 的训练量,很可能也能训练出类似的模型。
关键精度提升细分
需要注意的是,直接从论文实现模型是一个循环往复的过程,涉及编码、训练、调试和调整配置,直到达到论文中报告的精度。这通常还涉及简化训练配方或使用更现代的方法进行增强。这绝不是一个线性过程,即通过一次提升一个方向就能获得递增的精度,而是涉及探索不同的假设、在不同方面进行改进以及大量的回溯。
考虑到这一点,下面我们尝试总结对我们精度影响最大的优化。我们将各种实验分为 4 个主要组,并将改进归因于最匹配的项。注意,图表的 Y 轴从 18 而不是 0 开始,以便更清晰地展示优化之间的差异:

| 模型配置 | mAP 增量 | mAP |
|---|---|---|
| 具有“FasterRCNN 风格”超参数的基准线 | – | 19.5 |
| + 论文超参数 | 1.6 | 21.1 |
| + 数据增强 | 1.8 | 22.9 |
| + 权重初始化与输入缩放 | 1 | 23.9 |
| + LR 方案 | 1.2 | 25.1 |
我们的最终模型达到了 25.1 的 mAP,完全复现了论文中报告的 COCO 结果。这里有精度指标的详细分析。
希望您觉得本系列的第一部分有趣。在第二部分中,我们将重点讨论 SSDlite 的实现并探讨其与 SSD 的区别。在此之前,我们期待您的反馈。