由瓦西利斯·弗里尼奥蒂斯

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

在本系列的第 1 部分,我们将重点关注 SSD 算法的原版实现,正如在《单次多框检测器》论文中所描述的那样。我们将简要介绍算法的工作原理,然后介绍其主要组件,突出代码中的关键部分,并最终讨论我们如何训练发布的模型。我们的目标是涵盖所有必要的细节,包括那些在论文中没有涉及但在原版实现中存在的优化,以便重现该模型。

SSD 是如何工作的?

阅读上述论文是非常推荐的,但这里有一个简化的快速复习。我们的目标是检测图像中物体的位置及其类别。以下是 SSD 论文中的图 5,展示了模型的预测示例:

SSD 算法使用 CNN 骨干网络,将输入图像通过它传递,并从网络的各个层次提取卷积输出。这些输出的列表被称为特征图。然后,这些特征图通过分类和回归头部传递,这些头部负责预测类别和框的位置。

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

SSD 网络架构

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

DefaultBoxGenerator

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

该类由一组超参数参数化,这些超参数控制其形状和平铺。实现将自动提供默认参数的良好猜测,供那些想要尝试新的骨干/数据集的人使用,但也可以传递优化的自定义值。

SSDMatcher

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

分类和回归头

SSDHead 类负责初始化网络的分类和回归部分。以下是它们代码的一些显著细节:

  • 分类和回归头部都继承自同一个类,该类负责为每个特征图生成预测。
  • 每个特征图的每一层都使用一个独立的 3x3 卷积来估计类别对数和框的位置。
  • 每个头部在每个级别生成的预测数量取决于默认框的数量和特征图的大小。

主干特征提取器

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

该类支持所有 TorchVision 的 VGG 模型,可以为其他类型的 CNN 创建类似的提取器类(请参阅 ResNet 的示例)。以下是该类的几个实现细节:

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

SSD 算法

实现中的最后一项关键部分是在 SSD 类上。以下是一些值得注意的细节:

  • 算法通过一组类似于其他检测模型的参数进行参数化。必填参数包括:负责估计特征图的骨干网络,应配置为 DefaultBoxGenerator 类的实例的 anchor_generator ,输入图像将被调整的大小以及用于分类(不包括背景)的 num_classes
  • 如果没有提供头部,构造函数将初始化默认的 SSDHead 。为此,我们需要知道骨干网络产生的每个特征图的输出通道数。最初,我们尝试从骨干网络中检索此信息,如果不可用,我们将动态估计它。
  • 该算法重用了其他检测模型使用的标准 BoxCoder 类。该类负责编码和解码边界框,并配置为使用与原始实现相同的先验方差。
  • 虽然我们重用了标准 GeneralizedRCNNTransform 类来调整大小和归一化输入图像,但 SSD 算法配置它以确保图像大小保持固定。

这里是实现的两个核心方法:

  • compute_loss 方法估计了 SSD 论文第 5 页所述的标准多框损失。它使用平滑 L1 损失进行回归,并使用带硬负样本采样的标准交叉熵损失进行分类。
  • 在所有检测模型中,前向方法当前的行为取决于模型是在训练模式还是评估模式。它首先对输入图像进行调整大小和归一化,然后将它们通过骨干网络以获取特征图。然后,将特征图通过头部以获取预测,并生成默认框。
    • 如果模型处于训练模式,前向将估计默认框与真实值之间的 IoU,使用 SSDmatcher 生成匹配,并最终通过调用 compute_loss method 来估计损失。
    • 如果模型处于评估模式,我们首先通过仅保留通过分数阈值的检测来选择最佳检测,选择最有希望的框并运行 NMS 进行清理和选择最佳预测。最后,我们对预测进行后处理,将它们调整到原始图像大小。

SSD300 VGG16 模型

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

训练过程

该模型使用 COCO 数据集进行训练,所有超参数和脚本都可以在我们的参考文献文件夹中找到。以下我们提供训练过程中最显著方面的详细信息。

论文超参数

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

数据增强

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

权重初始化与输入缩放

我们还发现遵循论文中提出的权重初始化方案是有益的。为此,我们必须调整我们的输入缩放方法,通过撤销 ToTensor() 执行的 0-1 缩放,并使用与此缩放相匹配的预训练 ImageNet 权重(向 Max deGroot 致敬,他在他的仓库中提供了这些权重)。所有新卷积层的权重都使用 Xavier 初始化,其偏差设置为零。初始化后,网络从头到尾进行训练。

学习率方案

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

关键精度提升分解

需要注意的是,直接从论文中实现模型是一个迭代过程,它围绕着编码、训练、错误修复和配置调整循环,直到我们达到论文上报告的精度。通常,这也涉及到简化训练方案或用更现代的方法增强它。这绝对不是一个线性过程,其中通过每次只改进一个方向来逐步提高精度,而是涉及到探索不同的假设,在不同方面进行增量改进,并进行大量的回溯。

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

模型配置 mAP 增量 平均精度(mAP)
基于“FasterRCNN-style”超参数的基线 - 19.5
+ 论文超参数 1.6 21.1
+ 数据增强 1.8 22.9
+ 权重初始化与输入缩放 1 23.9
+ 学习率方案 1.2 25.1

我们的最终模型实现了 25.1 的 mAP,并精确地重现了论文中报告的 COCO 结果。以下是准确度指标的详细分解。

希望您觉得本系列的第 1 部分有趣。在第 2 部分,我们将重点关注 SSDlite 的实现,并讨论其与 SSD 的不同之处。在此之前,我们期待您的反馈。