博客

关于 TorchVision 的 MobileNetV3 实现,你需要知道的一切

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

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

网络架构

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

开箱即用,我们提供了论文中描述的两个变体:Large(大型)Small(小型)。两者均使用相同的代码构建,唯一的区别在于配置参数,这些参数定义了块的数量、块的大小、激活函数等。

配置参数

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

  • width_mult 参数是一个乘数,影响模型的通道数。默认值为 1,增大或减小该值可以改变所有卷积的滤波器数量,包括第一层和最后一层。该实现确保滤波器数量始终是8 的倍数。这是一种硬件优化技巧,可以实现更快的运算向量化。
  • reduced_tail 参数会将网络最后几个块的通道数减半。此版本被某些目标检测和语义分割模型使用。这是一种在MobileNetV3 论文中描述的速度优化方案,据称在不显著影响准确率的情况下可降低 15% 的延迟。
  • dilated 参数会影响模型中最后 3 个 InvertedResidual 块,将其常规的深度卷积(Depthwise Convolutions)转变为带孔卷积(Atrous Convolutions)。这用于控制这些块的输出步长(output stride),并对语义分割模型的准确率有显著的正面影响

实现细节

下面我们提供有关该架构一些显著实现细节的补充信息。MobileNetV3 类负责根据提供的配置构建网络。以下是该类的一些实现细节:

  • 最后一个卷积块将最后一个 InvertedResidual 块的输出扩展了 6 倍。该实现与论文中描述的 Large 和 Small 配置保持一致,并可适应不同的乘数参数值。
  • 与 MobileNetV2 等其他模型类似,Dropout 层被放置在分类器的最终线性层(Linear layer)之前。

InvertedResidual 类是网络的主要构建模块。以下是该模块的一些显著实现细节,以及参考论文图 4 的可视化内容:

  • 如果输入通道和扩展后的通道相同,则没有扩展步骤(expansion step)。这种情况发生在网络的第一个卷积块中。
  • 即使扩展后的通道与输出通道相同,也始终存在一个投影步骤(projection step)
  • 深度卷积块的激活方法放置在 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@1Acc@5CPU 推理 (秒)参数量 (M)
MobileNetV3-Large74.04291.3400.04115.48
MobileNetV3-Small67.66887.4020.01652.54
量化版 MobileNetV3-Large73.00490.8580.01622.96
MobileNetV271.88090.2900.06083.50
ResNet5076.15092.8700.254525.56
ResNet1869.76089.0800.103211.69

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

训练过程

所有预训练模型均配置为宽度乘数为 1,具有完整尾部,无膨胀(non-dilated),并在 ImageNet 上进行训练。Large 和 Small 变体都使用相同的超参数和脚本进行训练,这些可以在我们的参考文档文件夹中找到。下面我们提供训练过程中最引人注目的细节。

实现快速且稳定的训练

正确配置 RMSProp 对于实现快速且数值稳定的训练至关重要。论文作者在他们的实验中使用了 TensorFlow,并报告说他们使用了相比默认值相当高rmsprop_epsilon。通常,由于该超参数用于避免零除错误,它会取很小的值,但在该特定模型中,选择正确的值对于避免损失函数中的数值不稳定似乎很重要。

另一个重要的细节是,尽管 PyTorch 和 TensorFlow 的 RMSProp 实现通常表现类似,但存在一些差异,在我们的设置中,最明显的是如何处理 epsilon 超参数。更具体地说,PyTorch 在平方根计算之外添加 epsilon,而 TensorFlow 则在内部添加。这一实现细节的结果是,在移植论文的超参数时,需要调整 epsilon 值。一个合理的近似公式为 PyTorch_eps = sqrt(TF_eps)

通过调整超参数和改进训练配方来提高准确率

在配置优化器以实现快速稳定的训练后,我们转向优化模型的准确率。有几种技术帮助我们实现了这一点。首先,为了避免过拟合,我们使用了 AutoAugment 算法进行数据增强,随后进行了 RandomErasing。此外,我们通过交叉验证调整了权重衰减(weight decay)等参数。我们还发现,在训练结束后对不同的 epoch 检查点执行权重平均(weight averaging)是有益的。最后,虽然没有用于我们发布的训练配方中,但我们发现使用标签平滑(Label Smoothing)、随机深度(Stochastic Depth)和学习率噪声注入可以将总体准确率提高超过 1.5 个点。

该图表描绘了改进 MobileNetV3 Large 变体准确率的最重要迭代的简化摘要。请注意,训练模型时实际进行的迭代次数要多得多,而且准确率的提升并不总是单调递增的。还要注意,图表的 Y 轴从 70% 开始而不是 0%,以使迭代之间的差异更明显。

迭代Acc@1Acc@5
带有“MobileNetV2 风格”超参数的基准71.54290.068
+ 使用默认 eps 的 RMSProp70.68489.38
+ 使用调整后的 eps 和 LR 方案的 RMSProp71.76490.178
+ 数据增强与调整后的超参数73.8691.292
+ 检查点平均化74.02891.382
+ 标签平滑 & 随机深度 & LR 噪声75.53692.368

请注意,一旦达到可接受的准确率,我们会在之前未用于训练或超参数调整的留出测试数据集(hold-out test dataset)上验证模型性能。此过程有助于我们检测过拟合,并且在所有预训练模型发布前都会执行。

量化

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

请注意,QAT 允许我们模拟量化的影响并调整权重,从而提高模型准确率。与简单的训练后量化(post-training quantization)相比,准确率提升了 1.8 个点。

量化状态Acc@1Acc@5
非量化74.04291.340
量化感知训练73.00490.858
训练后量化71.16089.834

目标检测

在本节中,我们将首先提供已发布模型的基准测试,然后讨论如何将 MobileNetV3-Large 主干网络与特征金字塔网络(FPN)以及 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 模型的可行替代方案。

模型mAPCPU 推理 (秒)参数量 (M)
Faster R-CNN MobileNetV3-Large FPN (高分辨率)32.80.840919.39
Faster R-CNN MobileNetV3-Large 320 FPN (低分辨率)22.80.167919.39
Faster R-CNN ResNet-50 FPN37.04.151441.76
RetinaNet ResNet-50 FPN36.44.882534.01

实现细节

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

从网络提取的所有特征图其输出都会通过 FPN 块投影至 256 个通道,因为这极大地提高了网络速度。FPN 主干网提供的这些特征图被 FasterRCNN 检测器用于在不同尺度上提供框和类预测。

训练与调优过程

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

高分辨率检测器使用 800-1333px 的图像进行训练,而移动端友好的低分辨率检测器使用 320-640px 的图像进行训练。我们提供两套预训练权重的理由是:相比于将小图像传递给预训练好的高分辨率模型,直接在小图像上训练检测器可使精度提升 5 mAP。两个主干网络均使用在 ImageNet 上拟合的权重进行初始化,并在训练过程中微调其最后 3 个阶段的权重。

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

调优状态mAPCPU 推理 (秒)
之前23.00.2904
之后22.80.1679

以下是一些可视化 Faster R-CNN MobileNetV3-Large FPN 模型预测结果的示例:

语义分割

在本节中,我们将首先提供已发布的预训练模型的基准测试。然后,我们将讨论如何将 MobileNetV3-Large 主干网与诸如 LR-ASPPDeepLabV3FCN 等分割头部(segmentation heads)相结合来进行语义分割。我们还将解释网络是如何训练的,并为速度关键型应用提出几种可选的优化技术。

基准测试

这是初始化预训练模型的方法:

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

以下是新模型与选定现有模型之间的详细基准对比。正如我们所见,配备 MobileNetV3-Large 主干网的 DeepLabV3 对于大多数应用而言是 FCN (ResNet50) 的可行替代方案,因为它以 8.5 倍的速度提升实现了类似的准确率。我们还观察到,LR-ASPP 网络在所有指标上都超越了同类 FCN。

模型平均交并比 (mIoU)全局像素准确率CPU 推理 (秒)参数量 (M)
LR-ASPP MobileNetV3-Large57.991.20.32783.22
DeepLabV3 MobileNetV3-Large60.391.20.586911.03
FCN MobileNetV3-Large (未发布)57.890.90.37025.05
DeepLabV3 ResNet5066.492.46.353139.64
FCN ResNet5060.591.45.014632.96

实现细节

在本节中,我们将讨论已测试分割头部的关键实现细节。请注意,本节描述的所有模型都使用膨胀后的(dilated)MobileNetV3-Large 主干网络。

LR-ASPP

LR-ASPP 是 MobileNetV3 论文作者提出的轻量级(Lite)简化版空洞空间金字塔池化模型。与 TorchVision 中的其他分割模型不同,它不使用辅助损失函数。相反,它利用输出步长分别为 8 和 16 的低级和高级特征

与论文中使用带有可变步长的 49×49 平均池化层不同,我们的实现使用 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.165.26%55.890.30.1139
DeepLabV3 MobileNetV3-Large-3.863.86%56.590.30.2121
FCN MobileNetV3-Large (未发布)-3.057.57%54.890.10.1571

以下是一些可视化 LR-ASPP MobileNetV3-Large 模型预测结果的示例:

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