全网唯一复现!手机端 1ms 级延迟的主干网模型 MobileOne

技术

picture.image

本文作者:Ezra-Yu(GitHub ID)

在轻量化视觉基础模型系列的前两篇文章中,我们介绍了近期 CNN+Self-Attention 在轻量级基础模型中的工作,但在实际的业务中,工业界还是更喜欢使用 CNN 架构的模型。

主要原因有:Self-Attention 部署有风险 ,目前 Self-Attention 的架构在部署上即使有走通整个流程的案例,但依然存在风险,可能碰到各种各样的问题,耽误进度;Self-Attention 硬件优化不够 ,虽然论文中的 Self-Attention 架构计算量有优势,但实际的推理时延上还是不如传统的卷积神经网络,毕竟目前硬件对 CNN 优化已经很多年了。

本文就给大家介绍基于 CNN 架构中重参数化技术的轻量化主干网络 MobileOne 。虽然 Apple 公司在其项目 ml-mobileone 中已经开源了 MobileOne 的代码以及权重,但并没有公开训练和推理的策略。MMClassification 通过一系列努力,已经完全复现了 MobileOne 的训练精度(全网唯一) ,可以查看相关 issue、复现结果。

issue:

https://github.com/apple/ml-mobileone/issues/7

复现结果:

https://github.com/open-mmlab/mmclassification/tree/1.x/configs/mobileone

2021 年,旷视的 RepVGG 提出重参数化技术,模型在训练和推理时具有不同的结构,推理时结构更简洁,速度更快且易于后期的部署和算子优化。目前重参数化技术已经被广泛使用,如检测任务的 YOLOv6 以及 YOLOv7,且取得了非常好的效果。

MobileOne 模型本身没有特别新颖的创新点,可以看作一个 MobileNet、 重参数技术、训练技巧的组合创新成果。其在苹果 iPhone12 上,以推理时延小于 1ms 的速度在 ImageNet1k 数据集上达到了 75.9 的 TOP1 精度,到达目前轻量级模型最优的效果。

出发点

一直以来移动设备的高效神经网络主干通常是针对计算量或参数量等指标进行优化,然而,这些指标与真实部署在移动设备上的推理延迟存在一个差距。有些模型虽然有较低的计算量和参数量,但它的推理速度仍然非常慢。

为了探索推理时延最优的模型,苹果公司通过在移动设备上部署常见的轻量级模型,分析推理时延瓶颈所在,并提供了对应缓解的方法。MobileOne 与 EfficientFormer 使用的研究方法相似,都是通过实验分析 的方法。

实验分析

作者将常见的移动端轻量级模型在 iPhone12 上使用 coreML 部署,直接测试其推理时延,图 1 是所有的统计结果,可以看出 CNN 架构的模型具有明显的优势。为了定量分析计算量和参数量与推理时间的关系,作者计算了 Spearman 关联系数。图 2 展示了计算量和参数量与推理时延的 Spearman 系数,Spearman 系数值越大,代表它们的关联度越大,可以看出参数量和推理时延的关联程度较弱,计算量与推理时延的关联程度也一般。

picture.image

图 1 常用轻量级模型在 iPhone12 上的表现

picture.image

图 2 推理时延与计算量、参数量的 Spearman 关联系数

作者又分析了基本模块对推理时间的影响,主要是激活函数,残差结果以及 SE 模块的影响。图 3 展示了激活函数对推理时延的影响,使用复杂的激活函数虽然可以带来精度上的非常小的提升,但它会带来推理时延巨大增加。图 4 是 SE 模块与残差模块对推理时延的影响,它们都增加了大量的推理时延,文章中也给出了解释,这类多分支结构,会带来额外的访存开销。

picture.image

图 3 激活函数的影响

picture.image

图 4 SE 模块与残差结构影响

实验结论

  • CNN 结构在推理速度上有很大优势

  • ReLU 的速度最快,其他激活函数速度慢,精度提升并不明显

  • SE 模块与残差结构都显著的影响推理速度,应该尽量避免

MobileOne 结构

将上述实验结论作为设计指导,MobileOne 主要是基于 MobileNet 以及 RepVGG 这样的卷积神经网络,一方面为了减少残差结构的额外开销,使用 RepVGG 中的重参数技术,另一个方面只在最大的模型结构 MobileOne-s4 中才使用少量的 SE 模块。

MobileOne Block 基于深度可分离卷积,由多分支的 DepthWise 卷积模块和 PointWise 卷积模块组成,它有两种状态,训练时状态和推理时状态,图 5 左侧为训练时的状态,右侧为推理时状态。

在训练状态下,DepthWise 卷积模块有三个分支,分别为 1×1 DepthWise 卷积分支,3×3 DepthWise 卷积分支和一个 BN 层分支。PointWise 卷积模块有两个分支,分别为 1×1 卷积分支和一个 BN层 分支。

在推理状态下, DepthWise 卷积模块 和 PointWise 卷积模块只有一个 3×3 DepthWise 卷积和 1×1 PointWise 卷积,没有其他额外的分支,其中卷积和 BN 还可以进一步融合。

picture.image

图 5 MobileOne Block 结构

主要考虑到:

  • 输入分辨率的放大, 会增大计算量以及内存的消耗,这些都是对移动设备不友好的,所以所有的结构都使用 224*224 的输入分辨率。

  • 模型在推理时没有多分支的结构,不会产生额外的访存,所以设计时使用了更大的通道数。

作者提出了 5 种不同大小的模型结构,如图 6 所示,整体延续了 RepVGG 的架构,一共 8 个 stage,其中 6 个 stage 的 Block Type 为 MobileOne Block,1 个 stage 的 Block Type 为 AvgPool,1 个 stage 的 Block Type 为 Linear,只在最大的 MobileOne-S4 的最后的两个 stage 上才添加 SE 模块,所有的激活函数都使用 ReLU。

picture.image

图 6 MobileOne 模型整体结构

训练技术优化

picture.image

图 7 MobileOne 训练技巧消融实验

MobileOne 的训练策略整体上延续了 RepVGG,并在其基础上额外使用了 3 个特有的训练优化技巧,分别是:

  1. 渐进式学习策略

  2. 退火权重衰减系数

  3. 指数平均移动 EMA

渐进式学习策略是由 EfficientNetV2 提出的一种训练策略,简单来说就训练过程渐进地增大图像大小,但在增大图像同时也采用更强的正则化策略,训练的正则化策略包括数据增强和 dropout 等,MobileOne 中更强的正则化指的是采用更大的 AutoAug 的增强强度系数, 在训练的流程中,分别在 37 轮和 112 轮,增加了输入图片的分辨率以及 AutoAug 的强度系数。渐进式学习策略不仅可以加快训练速度,还可以提高模型精度。

picture.image

图 8 MobileOne 训练过程数据增强

与大型模型相比,小型模型需要较少的正则化来防止过度拟合。在训练的早期阶段权重正则化(Weight-Decay)很重要,作者实验发现,相比于消除训练过程中权重正则化,对权重正则化引起的损失进行退火处理更有效。

MobileOne 使用学习率中最常见的 Cosine schedule 来退火权重衰减系数,图 9 是 MobileOne 的权重衰减曲线,从图 9 中可以看到退火权重衰减系数提高了 0.5 %。

picture.image

图 9 MobileOne 权重衰减曲线

EMA 是指数平均移动(Exponential moving average),目前它是深度学习领域非常常见的训练技巧,近年大量新工作使用了该训练技巧。EMA 可以看作一种模型的集成,不过这里集成的是训练迭代过程中权重。其公式非常简单,如下所示:

公式中

是 t 时刻的网络参数,

是 t 时刻滑动平均后的网络参数,那么 t+1 时刻的滑动平均结果

就是这两者的加权融合。

是可以调节的超参数,表示每次更新时原模型的维持比例,

越大,原模型的维持比例就越大,迭代过程中模型权重的变就越迟钝。

目前这几种训练优化都已经在 MMClassification 中实现,欢迎使用,配置如下:


              
# 退火权重衰减系数
              
param_scheduler = [
              
    .......
              
    dict(
              
        type='CosineAnnealingParamScheduler',
              
        param_name='weight_decay',
              
        eta_min=0.00001,
              
        by_epoch=True,
              
        begin=0,
              
        end=300)
              
]
              

              
train_pipeline_37e = [..........] # 37 epoch 后的数据增强
              
train_pipeline_112e = [..........]  # 112 epoch 后的数据增强
              
 
              
custom_hooks = [
              
    # 渐进式学习策略
              
    dict(          
              
        type='SwitchRecipeHook', # 不仅可以切换数据增强,还可以切换batch_aug,以及loss
              
        schedule=[
              
            dict(action_epoch=37, pipeline=train_pipeline_37e),
              
            dict(action_epoch=112, pipeline=train_pipeline_112e),
              
        ]),
              
    # EMA
              
    dict(
              
        type='EMAHook',
              
        momentum=5e-4,
              
        priority='ABOVE_NORMAL',
              
        update_buffers=True)
              
]
          

实验结果

以下展示了桌面端以及移动端上 MobileOne 和其他轻量级模型的一些性能对比,可以看出其具有明显的性能优势。

picture.image

图 10 MobileOne 分类性能比较

picture.image

图 11 MobileOne 检测、分割任务性能比较

在分类任务上,它在速度精度综合评比中不仅超过了 CNN 架构的 MobileNet 、ShuffleNet 系列,还远超前文提到的 MobileViT 系列算法。其中 MobileOne-s1 结构更是在推理时延为 1ms 的前提下,ImageNet 数据集上精度达到 75.6。

MobileOne 不仅在分类任务上取得了速度与精度的最佳权衡,它在检测(图 11.a)和分割(图 11.b)任务中,也取得了非常不错的成果。

总结

本文介绍了最新的基于 CNN 架构的算法 MobileOne,重点主要有以下两点:

  • 将重参数化技术应用到深度可分离卷积中;

  • 使用了 3 个特有的训练优化技巧,分别是:渐进式学习策略;退火权重衰减系数;指数平均移动 EMA。

MobileOne 是目前精度速度综合表现最好的算法,特别适合应用到实际的业务场景中来。并且已经集成在 MMClassification 中,且已经对齐了训练精度,欢迎对相关模型感兴趣的同学尝试试用。

MobileOne:

https://github.com/open-mmlab/mmclassification/tree/dev-1.x/configs/mobileone

(文末点击阅读原文可直达,欢迎体验)

参考文献

  • An Improved One millisecond Mobile Backbone
  • RepVGG: Making VGG-style ConvNets Great Again
  • EfficientNetV2: Smaller Models and Faster Training
  • EfficientFormer: Vision Transformers at MobileNet Speed
0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
字节跳动 NoSQL 的实践与探索
随着 NoSQL 的蓬勃发展越来越多的数据存储在了 NoSQL 系统中,并且 NoSQL 和 RDBMS 的界限越来越模糊,各种不同的专用 NoSQL 系统不停涌现,各具特色,形态不一。本次主要分享字节跳动内部和火山引擎 NoSQL 的实践,希望能够给大家一定的启发。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论