博客

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

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

上一篇文章中,我们讨论了 SSD 算法的工作原理,涵盖了其实现细节并介绍了训练流程。如果你还没有阅读过前一篇博文,建议在继续阅读之前先去看看。

在本系列的第二部分中,我们将重点讨论 SSD 的移动端友好变体——SSDlite。我们的计划是:首先深入分析该算法的主要组件,重点介绍其与原始 SSD 的不同之处;接着探讨已发布模型的训练方式;最后为我们探索的所有新目标检测模型提供详细的基准测试数据。

SSDlite 网络架构

SSDlite 是 SSD 的一种适配版本,最早在 MobileNetV2 论文中被简要介绍,后来在 MobileNetV3 论文中再次使用。由于这两篇论文的主要重点是引入新型 CNN 架构,因此 SSDlite 的大部分实现细节并未明确说明。我们的代码遵循了这两篇论文中提出的所有细节,并在必要时填补了官方实现中的空白。

如前所述,SSD 是一个模型家族,因为可以通过配置不同的主干网络(如 VGG、MobileNetV3 等)和不同的头部(如使用常规卷积、可分离卷积等)来构建它。因此,SSD 的许多组件在 SSDlite 中保持不变。下面我们仅讨论不同的部分。

分类与回归头部 (Classification and Regression Heads)

根据 MobileNetV2 论文第 6.2 节的描述,SSDlite 使用可分离卷积替换了原始头部中使用的常规卷积。因此,我们的实现引入了新的头部,这些头部使用了3×3 深度卷积和 1×1 投影。由于 SSD 方法的所有其他组件保持不变,为了创建 SSDlite 模型,我们的实现会初始化 SSDlite 头部,并将其直接传递给 SSD 构造函数。

主干特征提取器 (Backbone Feature Extractor)

我们的实现引入了一个用于构建 MobileNet 特征提取器的新类。根据 MobileNetV3 论文第 6.3 节,主干网络返回倒置残差模块(Inverted Bottleneck block)中扩展层的输出(输出步长为 16)以及池化层之前那一层的输出(输出步长为 32)。此外,主干网络的所有额外模块都被替换为轻量级等效模块,这些模块使用 1×1 压缩、步长为 2 的 3×3 可分离卷积以及 1×1 扩展。最后,为了确保即使在使用较小的宽度乘数(width multipliers)时,头部也具备足够的预测能力,所有卷积的最小深度(minimum depth)均受 min_depth 超参数控制。

SSDlite320 MobileNetV3-Large 模型

本节讨论所提供的 SSDlite 预训练模型的配置,以及为了尽可能复现论文结果而遵循的训练过程。

训练过程

用于在 COCO 数据集上训练模型的所有超参数和脚本都可以在我们的 references 文件夹中找到。这里我们讨论训练过程中最值得注意的细节。

调优超参数

虽然论文没有提供用于训练模型的超参数(如正则化、学习率和批量大小)的相关信息,但官方仓库配置文件中列出的参数是很好的起点,通过交叉验证,我们将它们调整到了最佳值。所有这些调整使我们在基准 SSD 配置的基础上获得了显著提升。

数据增强

SSDlite 与 SSD 相比的一个关键差异在于,前者的主干网络权重仅为后者的极小一部分。这就是为什么在 SSDlite 中,数据增强更侧重于使模型对不同尺寸的目标具有鲁棒性,而不是为了防止过拟合。因此,SSDlite 仅使用 SSD 变换的一个子集,从而避免了对模型进行过度正则化。

学习率(LR)方案

由于依赖数据增强来使模型对中小型对象具有鲁棒性,我们发现训练方案中使用大量的 epoch 特别有益。具体来说,通过使用大约 3 倍于 SSD 的 epoch,我们的精度提高了 4.2 mAP,使用 6 倍乘数时则提高了 4.9 mAP。进一步增加 epoch 似乎边际效应递减,且使训练变得太慢而不切实际。尽管如此,根据模型配置来看,论文作者似乎使用了相当于 16 倍的乘数。

权重初始化、输入缩放与 ReLU6

最后一组优化使得我们的实现非常接近官方版本,并帮助弥补了精度差距:即从头开始训练主干网络(from scratch)而不是使用 ImageNet 初始化;调整我们的权重初始化方案;更改我们的输入缩放;并将 SSDlite 头部添加的所有标准 ReLU 替换为 ReLU6。请注意,由于我们是从随机权重开始训练模型的,我们额外应用了论文中描述的速度优化,即在主干网络上使用缩减的尾部(reduced tail)

实现差异

将上述实现与官方仓库进行比较,我们发现了一些差异。大多数差异很小,与我们如何初始化权重(例如 正态初始化 vs 截断正态分布)、如何参数化 LR 调度(例如更小 vs 更大的预热率,更短 vs 更长的训练)等有关。已知最大的差异在于我们计算分类损失的方式。更具体地说,官方仓库中带有 MobileNetV3 主干的 SSDlite 实现并未使用 SSD 的 Multibox 损失,而是使用了 RetinaNet 的焦点损失(focal loss)。这与论文的偏差相当大,由于 TorchVision 已经提供了 RetinaNet 的完整实现,我们决定使用常规的 Multi-box SSD 损失来实现 SSDlite。

关键精度提升拆解

如前几篇文章所讨论的,复现研究论文并将其转化为代码并非是一个精度单调递增的过程,特别是在不知道全部训练和实现细节的情况下。通常,这一过程涉及大量的回溯,因为人们需要识别出那些对精度有重大影响的实现细节和参数,并将它们与那些无关紧要的区分开来。下面我们尝试可视化那些将我们的精度从基准水平提升上来的最重要的迭代步骤:

迭代mAP
带有“SSD 风格”超参数的基准模型10.6
+ 调优后的超参数14.2
+ SSDlite 数据增强15.2
+ 3x LR 方案19.4
+ 6x LR 方案20.1
+ 权重初始化 & 输入缩放 & ReLU621.3

上述呈现的优化顺序是准确的,尽管在某些情况下显得有些理想化。例如,虽然在超参数调优阶段测试了不同的调度器,但它们都没有提供显著的改进,因此我们保留了基准测试中使用的 MultiStepLR。不过,在后期尝试不同的 LR 方案时,我们发现切换到 CosineAnnealingLR 很有益,因为它需要的配置更少。因此,我们认为从上述总结中应该得出的主要结论是:即使从正确的实现和同一模型系列的优化超参数集开始,通过优化训练配方和微调实现,总能找到提升精度的空间。不可否认,上述是一个相当极端的例子,精度翻了一番,但在许多情况下,确实有大量的优化可以帮助我们显著提高精度。

基准测试

以下是如何初始化这两个预训练模型的方法

ssdlite = torchvision.models.detection.ssdlite320_mobilenet_v3_large(pretrained=True)
ssd = torchvision.models.detection.ssd300_vgg16(pretrained=True)

以下是新检测模型与选定的先前检测模型之间的基准测试对比

模型mAPCPU 推理 (秒)参数数量 (M)
SSDLite320 MobileNetV3-Large21.30.09113.44
SSD300 VGG1625.10.830335.64
SSD512 VGG16 (未发布)28.82.249437.08
SSD512 ResNet50 (未发布)30.21.113742.70
Faster R-CNN MobileNetV3-Large 320 FPN (低分辨率)22.80.167919.39
Faster R-CNN MobileNetV3-Large FPN (高分辨率)32.80.840919.39

正如我们所见,SSDlite320 MobileNetV3-Large 模型是迄今为止最快、最小的模型,因此它是现实世界移动应用的出色候选者。虽然其精度低于预训练的低分辨率 Faster R-CNN 等效模型,但 SSDlite 框架具有适应性,可以通过引入更多卷积的更重型头部来提升其精度。

另一方面,SSD300 VGG16 模型速度较慢且精度较低。这主要是由于其 VGG16 主干网络。尽管 VGG 架构非常重要且具有影响力,但在今天看来已经相当过时。因此,尽管该特定模型具有历史和研究价值,因而被包含在 TorchVision 中,但我们建议那些希望在现实应用中使用高分辨率检测器的用户,要么将 SSD 与其他主干网络结合使用(参阅此示例了解如何创建),要么使用预训练的 Faster R-CNN 模型之一。

希望你喜欢 SSD 系列的第二部分,这也是该系列的最后一部分。我们期待收到你的反馈。