跳转到主要内容
博客

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

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

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

在本系列的第二部分中,我们将重点介绍SSD的移动友好型变体——SSDLite。我们的计划是首先介绍算法的主要组件,突出与原始SSD不同的部分,然后讨论已发布模型的训练方式,最后提供我们探索的所有新目标检测模型的详细基准。

SSDLite网络架构

SSDLite是SSD的一种改编,最初在MobileNetV2论文中简要介绍,后来在MobileNetV3论文中重复使用。由于这两篇论文的主要重点是引入新颖的CNN架构,因此SSDLite的大部分实现细节并未阐明。我们的代码遵循两篇论文中提出的所有细节,并在必要时从官方实现中填补空白。

如前所述,SSD是一个模型家族,因为可以使用不同的骨干网络(如VGG、MobileNetV3等)和不同的头部(如使用常规卷积、可分离卷积等)进行配置。因此,SSDLite中许多SSD组件保持不变。下面我们只讨论那些不同的部分。

分类和回归头部

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

骨干特征提取器

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

SSDLite320 MobileNetV3-Large模型

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

训练过程

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

调整的超参数

尽管论文没有提供用于训练模型的超参数信息(如正则化、学习率和批量大小),但官方仓库中配置文件中列出的参数是很好的起点,我们通过交叉验证将其调整到最优值。所有这些都使我们比基线SSD配置有了显著提升。

数据增强

SSDLite与SSD的一个重要区别在于,前者骨干网络的权重仅为后者的一部分。这就是为什么在SSDLite中,数据增强更侧重于使模型对可变大小的对象具有鲁棒性,而不是试图避免过拟合。因此,SSDLite仅使用SSD转换的子集,从而避免了模型的过度正则化。

学习率方案

由于依赖数据增强使模型对中小尺寸物体具有鲁棒性,我们发现使用大量 epoch 对训练方案特别有利。更具体地说,通过使用大约 3 倍于 SSD 的 epoch 数,我们能够将精度提高 4.2mAP 点;通过使用 6 倍乘数,我们能够提高 4.9mAP。进一步增加 epoch 数似乎收益递减,并使训练过于缓慢和不切实际,然而根据模型配置,论文作者似乎使用了等效的 *16 倍乘数*。

权重初始化 & 输入缩放 & ReLU6

最终优化方案集使我们的实现与官方实现非常接近,并帮助我们弥补了精度差距:从头开始训练骨干网络,而不是从ImageNet初始化;调整我们的权重初始化方案;改变我们的输入缩放;并将SSDLite头部中所有标准ReLUs替换为ReLU6。请注意,由于我们从随机权重训练模型,我们还应用了论文中描述的减尾速度优化。

实现差异

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

关键精度改进分解

正如前面文章讨论的,复现研究论文并将其转化为代码并非一个精度单调递增的过程,尤其是在训练和实现细节不完全已知的情况下。通常,这个过程会涉及大量的回溯,因为需要识别那些对精度有显著影响的实现细节和参数,并将其与那些没有影响的区分开来。下面我们尝试可视化从基线提高我们精度的最重要迭代。

迭代mAP
基线搭配“SSD风格”超参数10.6
+ 调优的超参数14.2
+ SSDlite数据增强15.2
+ 3倍学习率方案19.4
+ 6倍学习率方案20.1
+ 权重初始化 & 输入缩放 & ReLU621.3

上面列出的优化顺序是准确的,尽管在某些情况下略有理想化。例如,虽然在超参数调优阶段测试了不同的调度器,但它们都没有提供显著的改进,因此我们保持了基线中使用的MultiStepLR。然而,在后来尝试不同的学习率方案时,我们发现切换到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系列的第二部分也是最后一部分。期待您的反馈。