作者 Vasilis Vryniotis

在 TorchVision v0.10 中,我们发布了两个基于 SSD 架构的新目标检测模型。我们的计划是用两篇文章来涵盖算法的关键实现细节以及关于它们如何训练的信息。

在该系列的第一部分,我们将重点关注 Single Shot MultiBox Detector 论文中描述的 SSD 算法的原始实现。我们将简要地对算法的工作原理进行高级描述,然后介绍其主要组成部分,突出其代码的关键部分,最后讨论我们如何训练已发布的模型。我们的目标是涵盖重现模型所需的所有细节,包括论文中未涉及但在原始实现中的那些优化。

SSD 如何工作?

强烈建议阅读上述论文,但这里有一个快速过分简化的回顾。我们的目标是检测图像中物体的位置以及它们的类别。这里是 SSD 论文中的图 5,展示了模型的预测示例

SSD 算法使用 CNN 主干网络,将输入图像通过它传递,并获取来自网络不同层次的卷积输出。这些输出列表称为特征图。这些特征图然后通过分类和回归头(Classification and Regression heads),它们负责预测类别和框的位置。

由于每个图像的特征图包含来自网络不同层次的输出,它们的大小不同,因此可以捕获不同尺寸的物体。在每个特征图顶部,我们平铺了几个默认框(default boxes),可以将它们视为我们粗略的先验猜测。对于每个默认框,我们预测是否存在物体(及其类别)及其偏移量(相对于原始位置的修正)。在训练时,我们首先需要将真实标注(ground truth)与默认框匹配,然后使用这些匹配来估计损失。在推理时,将相似的预测框组合起来以估计最终预测。

SSD 网络架构

在本节中,我们将讨论 SSD 的关键组成部分。我们的代码紧密遵循该论文,并使用了官方实现中包含的许多未文档化的优化。

DefaultBoxGenerator

DefaultBoxGenerator 类负责生成 SSD 的默认框,其工作方式与 FasterRCNN 的 AnchorGenerator 类似(有关它们差异的更多信息,请参见论文第 4-6 页)。它生成一组特定宽度和高度的预定义框,这些框在图像上平铺,作为物体可能位置的第一个粗略先验猜测。这里是 SSD 论文中的图 1,展示了真实标注和默认框的可视化

该类由一组控制其形状平铺的超参数进行参数化。对于想要实验新主干网络/数据集的人来说,该实现将通过默认参数自动提供好的猜测,但也可以传递优化的自定义值

SSDMatcher

SSDMatcher 类扩展了 FasterRCNN 使用的标准 Matcher,负责将默认框与真实标注进行匹配。在估计所有组合的 IoU 后,我们使用 Matcher 为每个默认框找到最佳的候选真实标注,其重叠高于IoU 阈值。SSD 版本的 Matcher 有一个额外的步骤,以确保每个真实标注都与具有最高重叠的默认框匹配。Matcher 的结果用于模型训练过程中的损失估计。

分类和回归头

SSDHead 类负责初始化网络的分类和回归部分。以下是关于其代码的一些值得注意的细节

主干网络特征提取器

特征提取器重新配置并增强了一个标准的 VGG 主干网络,增加了额外的层,如 SSD 论文图 2 所示

该类支持 TorchVision 的所有 VGG 模型,并且可以为其他类型的 CNN 创建类似的提取器类(参见 ResNet 的这个例子)。以下是该类的一些实现细节

  • 修补(Patching)第三个 Maxpool 层的 ceil_mode parameter 是必要的,以便获得与论文相同的特征图大小。这是由于 PyTorch 与模型的原始 Caffe 实现之间存在微小差异。
  • 它在 VGG 之上添加了一系列额外的特征层。如果在构建期间 highres 参数为 True,它将附加一个额外的卷积层。这对于 SSD512 版本的模型很有用。
  • 如论文第 3 节所讨论的,原始 VGG 的全连接层被转换为卷积层,其中第一个使用空洞卷积(Atrous)。此外,maxpool5 的步长(stride)和核大小(kernel size)被修改
  • 如第 3.1 节所述,L2 归一化用于conv4_3 的输出,并引入一组可学习的权重来控制其缩放。

SSD 算法

实现的最后一个关键部分是SSD 类。这里有一些值得注意的细节

以下是实现的两个核心方法

SSD300 VGG16 模型

SSD 是一个模型家族,因为它可以配置不同的主干网络和不同的头部配置。在本节中,我们将重点介绍提供的SSD 预训练模型。我们将讨论其配置的细节以及用于重现报告结果的训练过程。

训练过程

该模型使用 COCO 数据集进行训练,其所有超参数和脚本可在我们的参考资料文件夹中找到。下面我们提供关于训练过程中最值得注意方面的详细信息。

论文超参数

为了在 COCO 上获得最佳结果,我们采用了论文第 3 节中描述的超参数,涉及优化器配置、权重正则化等。此外,我们发现采用官方实现中关于 DefaultBox 生成器的平铺配置的优化很有用。这项优化在论文中没有描述,但对于提高小物体的检测精度至关重要。

数据增强

实现论文第 6 页和第 12 页描述的SSD 数据增强策略对于重现结果至关重要。更具体地说,使用随机“Zoom In”和“Zoom Out”变换使模型对各种输入尺寸具有鲁棒性,并提高了其在小物体和中等物体上的精度。最后,由于 VGG16 参数较多,增强中包含的光度失真具有正则化效果,有助于避免过拟合。

权重初始化 & 输入缩放

我们发现另一个有益的方面是遵循论文提出的权重初始化方案。为此,我们不得不通过撤销 ToTensor() 执行的 0-1 缩放来调整我们的输入缩放方法,并使用与此缩放匹配的预训练 ImageNet 权重(感谢 Max deGroot 在他的仓库中提供它们)。新卷积层的所有权重都使用 Xavier 初始化,并且它们的偏置项设置为零。初始化后,网络进行了端到端训练

学习率方案

正如论文报道的,应用激进的数据增强后,需要训练更长时间的模型。我们的实验证实了这一点,我们不得不调整学习率、批量大小和总步骤以获得最佳结果。我们提出的学习方案配置相对保守,在步骤之间显示出平台迹象,因此您很可能只需我们周期数的 66% 就能训练出类似的模型。

关键精度提升分解

值得注意的是,直接从论文实现一个模型是一个迭代过程,它在编码、训练、错误修复和调整配置之间循环,直到我们匹配论文报道的精度。通常,它还涉及简化训练配方或用更近期的方法增强它。这绝不是一个线性过程,其中增量精度提升是通过一次改进一个方向来实现的,而是涉及探索不同的假设,在不同方面进行增量改进,并进行大量的回溯。

考虑到这一点,下面我们尝试总结对我们精度影响最大的优化。我们通过将各种实验分为 4 个主要组,并将实验改进归因于最接近的匹配来实现这一点。请注意,图的 Y 轴从 18 开始而不是从 0 开始,以使优化之间的差异更明显

模型配置 mAP 增量 mAP
基线模型,带有“FasterRCNN 风格”超参数 - 19.5
+ 论文超参数 1.6 21.1
+ 数据增强 1.8 22.9
+ 权重初始化 & 输入缩放 1 23.9
+ 学习率方案 1.2 25.1

我们的最终模型达到了 25.1 的 mAP,并完全重现了论文报告的 COCO 结果。这是精度指标的详细分解

希望您觉得本系列的第一部分很有趣。在第二部分,我们将重点介绍 SSDlite 的实现,并讨论其与 SSD 的差异。在此之前,我们期待您的反馈。