作者:Vasilis Vryniotis 和 Francisco Massa

在 TorchVision v0.9 中,我们发布了一系列可用于分类、目标检测和语义分割的新的移动端友好模型。在本文中,我们将深入探讨这些模型的代码,分享值得注意的实现细节,解释我们如何配置和训练它们,并强调在调优过程中做出的重要权衡。我们的目标是公开原始论文和模型仓库中通常未记录的技术细节。

网络架构

MobileNetV3 架构的实现紧密遵循原始论文。它是可定制的,并提供不同的配置用于构建分类、目标检测和语义分割骨干网络。它的设计遵循与 MobileNetV2 相似的结构,两者共享共同的构建块

我们直接提供了论文中描述的两种变体:LargeSmall。两者使用相同的代码构建,唯一区别在于它们的配置,配置描述了块的数量、大小、激活函数等。

配置参数

尽管可以直接编写自定义的 InvertedResidual 设置并将其传递给 MobileNetV3 类,但对于大多数应用,我们可以通过将参数传递给模型构建方法来调整现有配置。以下是一些关键的配置参数:

  • width_mult 参数是一个乘数,影响模型的通道数。默认值为 1,通过增加或减少它,可以改变所有卷积层的滤波器数量,包括第一层和最后一层的滤波器。实现确保滤波器数量始终是 8 的倍数。这是一种硬件优化技巧,可以加快操作的向量化。

  • reduced_tail 参数将网络最后几个块的通道数减半。某些目标检测和语义分割模型使用此版本。这是MobileNetV3 论文中描述的一种速度优化,据报道可以在对精度没有显著负面影响的情况下将延迟降低 15%。

  • dilated 参数影响模型最后 3 个 InvertedResidual 块,将其常规的深度可分离卷积转换为空洞卷积。这用于控制这些块的输出步幅,并对语义分割模型的精度产生显著的积极影响

实现细节

下面我们提供有关该架构一些值得注意的实现细节的额外信息。MobileNetV3 类负责根据提供的配置构建网络。以下是该类的一些实现细节:

  • 最后一个卷积块将最后一个 InvertedResidual 块的输出扩展了6 倍。该实现与论文中描述的 Large 和 Small 配置一致,并且可以适应乘数参数的不同值。

  • 与 MobileNetV2 等其他模型类似,在分类器的最终全连接层之前放置了一个 dropout 层。

InvertedResidual 类是网络的主要构建块。以下是该块的一些值得注意的实现细节,以及来自论文图 4 的可视化:

  • 如果输入通道数和扩展后的通道数相同,则没有扩展步骤。这发生在网络的第一个卷积块上。

  • 即使扩展后的通道数与输出通道数相同,也始终存在投影步骤

  • 深度可分离块的激活函数放置在 Squeeze-and-Excite 层之前,因为这可以略微提高精度。

分类

在本节中,我们将提供预训练模型的基准测试,并详细介绍它们的配置、训练和量化过程。

基准测试

以下是如何初始化预训练模型

large = torchvision.models.mobilenet_v3_large(pretrained=True, width_mult=1.0,  reduced_tail=False, dilated=False)
small = torchvision.models.mobilenet_v3_small(pretrained=True)
quantized = torchvision.models.quantization.mobilenet_v3_large(pretrained=True)

下面是新模型与选定现有模型之间的详细基准测试。正如我们所见,对于那些愿意牺牲一点精度以换取大约 6 倍速度提升的用户来说,MobileNetV3-Large 是 ResNet50 的一个可行替代方案。

模型 Acc@1 Acc@5 CPU 推理时间(秒) 参数数量(百万)
MobileNetV3-Large 74.042 91.340 0.0411 5.48
MobileNetV3-Small 67.668 87.402 0.0165 2.54
量化后的 MobileNetV3-Large 73.004 90.858 0.0162 2.96
MobileNetV2 71.880 90.290 0.0608 3.50
ResNet50 76.150 92.870 0.2545 25.56
ResNet18 69.760 89.080 0.1032 11.69

请注意,推理时间是在 CPU 上测量的。它们不是绝对的基准测试,但可用于模型之间的相对比较。

训练过程

所有预训练模型均配置宽度乘数为 1,具有完整尾部,非空洞卷积,并在 ImageNet 上拟合。Large 和 Small 变体均使用相同的超参数和脚本进行训练,这些内容可在我们的参考文件夹中找到。下面我们将详细介绍训练过程中最值得注意的方面。

实现快速稳定的训练

正确地配置 RMSProp 对于实现具有数值稳定性的快速训练至关重要。论文作者在他们的实验中使用了 TensorFlow,并在他们的运行中报告使用了比默认值高很多rmsprop_epsilon。通常,此超参数取较小的值,因为它用于避免分母为零,但在该特定模型中,选择正确的值似乎对于避免损失函数中的数值不稳定至关重要。

另一个重要细节是,尽管 PyTorch 和 TensorFlow 的 RMSProp 实现通常表现相似,但仍存在一些差异,其中最值得注意的是在我们的设置中 epsilon 超参数的处理方式。更具体地说,PyTorch 在平方根计算的外部添加 epsilon,而 TensorFlow 则在内部添加。此实现细节的结果是,在移植论文中的超参数时,需要调整 epsilon 值。一个合理的近似可以使用公式 PyTorch_eps = sqrt(TF_eps) 得出。

通过调整超参数和改进训练方案来提高精度

配置好优化器以实现快速稳定的训练后,我们转而优化模型的精度。有一些技术帮助我们实现了这一点。首先,为了避免过拟合,我们使用 AutoAugment 算法对数据进行增强,接着是 RandomErasing。此外,我们还使用交叉验证调优了权重衰减等参数。我们还发现,在训练结束后,对不同 epoch 检查点进行权重平均是有益的。最后,虽然未在发布的训练方案中使用,但我们发现使用 Label Smoothing、Stochastic Depth 和 LR 噪声注入可以将总体精度提高 1.5 个点以上。

图表展示了提高 MobileNetV3 Large 变体精度的最重要迭代的简化摘要。请注意,模型训练时的实际迭代次数要多得多,并且精度提升并非总是单调递增。另请注意,图表的 Y 轴从 70% 开始而不是 0%,以使迭代之间的差异更明显。

迭代 Acc@1 Acc@5
基线(采用“MobileNetV2 式”超参数) 71.542 90.068
+ 使用默认 eps 的 RMSProp 70.684 89.38
+ 使用调整后的 eps 和 LR 方案的 RMSProp 71.764 90.178
+ 数据增强和调优后的超参数 73.86 91.292
+ 检查点平均 74.028 91.382
+ Label Smoothing、Stochastic Depth 和 LR 噪声 75.536 92.368

请注意,一旦我们达到可接受的精度,我们就会在之前未用于训练或超参数调优的保留测试数据集上验证模型性能。此过程有助于我们检测过拟合,并且在发布所有预训练模型之前都会执行此操作。

量化

我们目前为 MobileNetV3-Large 变体的 QNNPACK 后端提供了量化权重,可提供 2.5 倍的加速。为了量化模型,使用了量化感知训练 (QAT)。用于训练模型的超参数和脚本可在我们的参考文件夹中找到。

请注意,QAT 使我们能够模拟量化的影响并调整权重,从而提高模型精度。与简单的训练后量化相比,这带来了 1.8 个点的精度提升。

量化状态 Acc@1 Acc@5
未量化 74.042 91.340
量化感知训练 73.004 90.858
训练后量化 71.160 89.834

目标检测

在本节中,我们将首先提供已发布模型的基准测试,然后讨论如何在特征金字塔网络中使用 MobileNetV3-Large 骨干网络以及 FasterRCNN 检测器来执行目标检测。我们还将解释如何训练和调优网络以及我们必须做出的任何权衡。我们不会涵盖与SSDlite一起使用的细节,这将在未来的文章中讨论。

基准测试

以下是如何初始化模型

high_res = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(pretrained=True) 
low_res = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_320_fpn(pretrained=True)

下面是新模型与选定现有模型之间的一些基准测试。正如我们所见,对于那些愿意牺牲几个精度点以换取 5 倍速度提升的用户来说,使用 MobileNetV3-Large FPN 骨干网络的高分辨率 Faster R-CNN 似乎是等效 ResNet50 模型的一个可行替代方案。

模型 mAP CPU 推理时间(秒) 参数数量(百万)
Faster R-CNN MobileNetV3-Large FPN(高分辨率) 32.8 0.8409 19.39
Faster R-CNN MobileNetV3-Large 320 FPN(低分辨率) 22.8 0.1679 19.39
Faster R-CNN ResNet-50 FPN 37.0 4.1514 41.76
RetinaNet ResNet-50 FPN 36.4 4.8825 34.01

实现细节

该检测器使用 FPN 式骨干网络,从 MobileNetV3 模型的不同卷积层提取特征。默认情况下,预训练模型使用第 13 个 InvertedResidual 块的输出以及池化层之前的卷积输出,但该实现支持使用更多阶段的输出。

从网络中提取的所有特征图都由 FPN 块将其输出投影到256 个通道,因为这极大地提高了网络的速度。FasterRCNN 检测器使用 FPN 骨干网络提供的这些特征图来提供不同尺度的边界框和类别预测。

训练与调优过程

我们目前提供两个能够在不同分辨率下执行目标检测的预训练模型。这两个模型都在 COCO 数据集上使用相同的超参数和脚本进行训练,这些内容可在我们的参考文件夹中找到。

高分辨率检测器使用 800-1333 像素的图像进行训练,而移动端友好的低分辨率检测器则使用 320-640 像素的图像进行训练。我们提供两组独立的预训练权重的原因是,直接在较小图像上训练检测器,与将小图像输入预训练高分辨率模型相比,mAP 精度提高了 5 个点。两个骨干网络均使用 ImageNet 上拟合的权重进行初始化,并在训练过程中对权重的最后 3 个阶段进行了微调。

通过调整 RPN NMS 阈值,可以对移动端友好模型进行额外的速度优化。通过仅牺牲 0.2 mAP 的精度,我们将模型的 CPU 速度提高了大约 45%。优化的详细信息如下所示:

调优状态 mAP CPU 推理时间(秒)
之前 23.0 0.2904
之后 22.8 0.1679

下面提供了一些 Faster R-CNN MobileNetV3-Large FPN 模型的预测可视化示例

语义分割

在本节中,我们将首先提供已发布预训练模型的一些基准测试。然后我们将讨论如何将 MobileNetV3-Large 骨干网络与 LR-ASPPDeepLabV3FCN 等分割头结合进行语义分割。我们还将解释网络的训练过程,并为对速度要求严格的应用提出一些可选的优化技术。

基准测试

以下是如何初始化预训练模型

lraspp = torchvision.models.segmentation.lraspp_mobilenet_v3_large(pretrained=True) 
deeplabv3 = torchvision.models.segmentation.deeplabv3_mobilenet_v3_large(pretrained=True)

下面是新模型与选定现有模型之间的详细基准测试。正如我们所见,对于大多数应用而言,使用 MobileNetV3-Large 骨干网络的 DeepLabV3 是 ResNet50 FCN 的一个可行替代方案,因为它在实现类似精度的同时获得了 8.5 倍的速度提升。我们还注意到 LR-ASPP 网络在所有指标上均优于等效的 FCN。

模型 mIoU 全局像素精度 CPU 推理时间(秒) 参数数量(百万)
LR-ASPP MobileNetV3-Large 57.9 91.2 0.3278 3.22
DeepLabV3 MobileNetV3-Large 60.3 91.2 0.5869 11.03
FCN MobileNetV3-Large(未发布) 57.8 90.9 0.3702 5.05
DeepLabV3 ResNet50 66.4 92.4 6.3531 39.64
FCN ResNet50 60.5 91.4 5.0146 32.96

实现细节

在本节中,我们将讨论已测试分割头的重要实现细节。请注意,本节描述的所有模型都使用了空洞卷积的 MobileNetV3-Large 骨干网络。

LR-ASPP

LR-ASPP 是 MobileNetV3 论文作者提出的 Reduced Atrous Spatial Pyramid Pooling 模型的精简版。与 TorchVision 中的其他分割模型不同,它不使用辅助损失。它而是使用输出步幅分别为 8 和 16 的低层和高层特征

与论文中使用具有可变步幅的 49x49 平均池化层不同,我们的实现使用 AdaptiveAvgPool2d 层来处理全局特征。这是因为论文作者将该头模型专门针对 Cityscapes 数据集进行了调整,而我们的重点是提供一个适用于多个数据集的通用实现。最后,我们的实现总是在返回输出之前进行双线性插值,以确保输入和输出图像的尺寸完全匹配。

DeepLabV3 和 FCN

MobileNetV3 与 DeepLabV3 和 FCN 的结合方式与其他模型类似,并且这些方法的阶段估计与 LR-ASPP 相同。唯一值得注意的区别是,我们不使用高层和低层特征,而是将常规损失附加到输出步幅为 16 的特征图上,并将辅助损失附加到输出步幅为 8 的特征图上(链接)

最后需要注意的是,该模型的 FCN 版本并未发布,因为它在速度和精度方面完全被 LR-ASPP 取代。但是,预训练权重仍然可用,并且只需对代码进行微小更改即可使用。

训练与调优过程

我们目前提供两个能够执行语义分割的 MobileNetV3 预训练模型:LR-ASPP 和 DeepLabV3。模型的骨干网络使用 ImageNet 权重进行初始化并进行端到端训练。这两种架构均在 COCO 数据集上使用相同的脚本和相似的超参数进行训练。详细信息可在我们的参考文件夹中找到。

通常,推理时图像被调整为 520 像素。一种可选的速度优化方法是构建模型的低分辨率配置,方法是使用高分辨率预训练权重,并将推理时的图像大小调整减少到 320 像素。这将使 CPU 执行时间提高大约 60%,同时牺牲几个 mIoU 点。此优化的详细数据可在下表中找到:

低分辨率配置 mIoU 差异 速度提升 mIoU 全局像素精度 CPU 推理时间(秒)
LR-ASPP MobileNetV3-Large -2.1 65.26% 55.8 90.3 0.1139
DeepLabV3 MobileNetV3-Large -3.8 63.86% 56.5 90.3 0.2121
FCN MobileNetV3-Large(未发布) -3.0 57.57% 54.8 90.1 0.1571

以下是 LR-ASPP MobileNetV3-Large 模型预测的可视化示例

希望您觉得这篇文章有价值。我们期待您的反馈,以了解您是否希望我们更频繁地发布此类内容。如果社区认为这类文章有用,我们将乐于发布更多涵盖新引入的机器学习模型实现细节的文章。