博客

关于 Torchvision 的 SSD 实现,你需要知道的一切

作者: 2021年6月16日2024年11月16日暂无评论

在 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 类 负责初始化网络的分类和回归部分。以下是其代码的一些显著细节:

主干特征提取器

特征提取器通过额外层对标准 VGG 主干进行了重构和增强,如 SSD 论文图 2 所示。

该类支持 TorchVision 中的所有 VGG 模型,用户也可以为其他类型的 CNN 创建类似的提取器类(参见此 ResNet 示例)。以下是该类的一些实现细节:

SSD 算法

实现中的最后一个关键部分是 SSD 类。以下是一些显著细节:

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

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.621.1
+ 数据增强1.822.9
+ 权重初始化与输入缩放123.9
+ LR 方案1.225.1

我们的最终模型达到了 25.1 的 mAP,完全复现了论文中报告的 COCO 结果。这里有精度指标的详细分析

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