EfficientFormer: Vision Transformers at MobileNet Speed

最新推荐文章于 2024-04-10 09:53:55 发布

kebijuelun

最新推荐文章于 2024-04-10 09:53:55 发布

阅读量635

点赞数

分类专栏： paper_reading 文章标签：深度学习人工智能

本文链接：https://blog.csdn.net/kebijuelun/article/details/125801890

版权

paper_reading 专栏收录该内容

92 篇文章 11 订阅

订阅专栏

Paper name

EfficientFormer: Vision Transformers at MobileNet Speed

Paper Reading Note

URL: https://arxiv.org/pdf/2206.01191.pdf

TL;DR

本文目标是回答一个问题：transformer 是否可以在比 cnn 运行更快的时候同时精度更高？本文提出了 EfficientFormer 在修改一些耗时多的模块与进行网络结构搜索的情况下对前面的问题给了正向回复。

Introduction

背景

Vision Transformer (ViT) 在计算机视觉的很多任务中取得了优异的表现，不过因为模型的大参数量和模型结构(attention mechanism) 导致其在运行时一般比轻量级卷积网络慢
transformer 运行速度慢的原因分析：
- 参数量大
- 随着 token 长度增加计算复杂度呈二次增加
- norm layer 不可融合
- 缺乏编译器级别优化 (比如 CNN 中的 Winograd)
之前的一些 ViT 的轻量化方案主要是通过网络结构搜索或结合 MobileNet block 的混合结构设计，但是这样的测试时间依然难以达到实时应用需求
这就引起了一个问题：transformer 是否可以在比 cnn 运行更快的时候同时精度更高

本文方案

本位提出了一种 dimension-consistent 的纯 transformer 网络，另外通过 latency-driven slimming 来得到一系列最终模型 (面向运行时间优化，而不是 MAC 或参数量)，起名为 EfficientFormer
EfficentFormer-L1 在 Imagenet-1K 上 top1 acc 为 79.2%，iphone12 上运行耗时为 1.6 ms (基于 CoreML 编译测试)
- 对比数据： MobileNetV2×1.4 (1.6 ms, 74.7% top1)
EfficientFormer-L7 (7ms, 83.3% top1)

Dataset/Algorithm/Model/Experiment Detail

实现方式

模型运行延迟分析

在这里插入图片描述
从以上测试结果得到的发现：

大 kernel 和 stride 的 patch embedding 是在移动设备上的速度瓶颈。由 DeiT-S、PoolFormer 与 LeViTate-256 的对比可以看出。慢的原因主要是 large-kernel 卷积在编译器级别没有类似 Winograd 之类的优化，这里替换成几个 3x3 卷积代替直连能加速
一致的特征维数对于 token mixer 的选择很重要。token mixer 的可选方案有传统的拥有去哪聚感受野的 MHSA mixer、更复杂的 shifted window attention、类似 pooling 的非参数化算子。其中 shifted window 算子目前大部分移动设备编译器都不支持，主要关心运算高效的 pooling token mixer 和精度更优的 MHSA。从上图对比 PoolFormer-s24 与 LeViT-256 可以看出 reshape 是瓶颈所在，保持特征维数一致非常重要。对比 DeiT-Small 和 LeViT-256，可以发现 MHSA 在特征维度保持一致的时候带来的延迟并不多。本文设计在 4D 特征和 3D MHSA 中尽量不用 reshape。
conv-bn 的延迟比 LN(GN)-Linear 更优，掉点可接受。因为 BN 在测试阶段吸到 conv 中能降低测试延迟，从 DeiT-Small and PoolFormer-S24 等模型可以看出无法被吸的 norm 层在测试阶段占了总时长 10%-20%。本文中在 4D 特征中尽量使用 conv-bn 结构，3D 特征中使用 LN 为了获取更高精度。
非线性层的延迟取决于硬件和编译器。GeLU 在 iphone12 上几乎不比 relu 慢，但是 HardSwish 很慢（LeViT-256 在使用 HardSwish 时延迟为 44.5 ms，使用 GeLU 是 11.9 ms）。本文使用 GeLU 作为非线性层。

EfficientFormer 设计

在这里插入图片描述

整个网络由一个 patch embeding 和一系列 meta transformer blocks (MB) 构成

其中 MB 是 tokenmixer 后接 MLP 构成

详细结构

4D partition 具体结构
- 2 个 stride=2 的 3x3 conv 进行 patch embedding
- pool 与 conv bn gelu 等算子用于提取 low level 特征
3D partition 部分结构依然是和传统 ViT 结构一致

Latency Driven Slimming

Design of Supernet，定义了每个 metapath 中可以选择的模块，其中 I 是 identity path，j 代表 j-th stage，i 代表 i-th block

其中只在最后两个 stage 把 MB3D 放在搜索空间中，因为 MHSA 在 early stage 耗时明显，另外放置在后两个 stage 也符合 early stage 获取 low-level 特征，last stage 获取 long-term dependencies 的设计直觉。
Searching Space：主要是每个 stage 的宽度（Cj）和每个 stage 的 block 数量（Nj），以及从最后的第 N 个 block 开始用 MB3D
Searching Algorithm：提出了一种简单快速的基于梯度的搜索算法
- step1：使用 Gumble Softmax sampling 来训练 supernet，获取每个 metapath 中的每个 block 的重要性得分
  
  其中 $\alpha$ 是代表选择一个 block 的概率，也即代表了该 block 的重要性； $\epsilon \in U(0, 1)$ 确保探索； $\tau$ 代表温度； $n$ 代表 metapath 中的 block 类型；
- step2：获取不同宽度的 MB4D 与 MB3D 在端上设备运行延迟的 LUT 表
- step3：基于 LUT 表对 step1 中训练的 supernet 进行 gradual slimming；基于以下定义计算每个 block 的重要性得分，然后对每个 stage 中的所有 block 重要求和计算得到 stage 的重要性得分
  
  action space 如下：
  - action1：对于最不重要的 MP 使用 Identity path
  - action2：移除第一个 MB3D
  - action3：降低最不重要的 MP 的宽度
    结合延迟 LUT 表计算以上 action 的延迟，然后评估每个 action 的精度降低程度，选择降低延迟大且掉点低的 action