EdgeFormer: 向视觉 Transformer 学习,构建一个比 MobileViT 更好更快的卷积网络

 前言 本文主要探究了轻量模型的设计。通过使用 Vision Transformer 的优势来改进卷积网络,从而获得更好的性能。

欢迎关注公众号CV技术指南,专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

论文://arxiv.org/abs/2203.03952

代码://github.com/hkzhang91/EdgeFormer

 

核心内容


本文主要探究了轻量模型的设计。通过使用 Vision Transformer 的优势来改进卷积网络,从而获得更好的性能。

本文提出的核心算子,即 global circular convolution (GCC),是一个卷积操作,但是会引入位置嵌入,同时还具有全局的感受野。另一个改进点是使用提出的 GCC 和 SE 操作构建了类似于 Vision Transformer 的基础操作单元。借助于 SE 引入了样本相关的注意力机制。

提出的结构组合起来可以作为一个基础的单元用于替换卷积网络或者 Transformer 架构中的相关模块。

基于这些提出的结构构建的轻量卷积结构在更少的参数量和推理速度的同时,获得了比轻量卷积网络和 Vision Transformer 更好的性能。

对于 ImageNet-1K 的分类,相较于基线模型 MobileViT,EdgeFormer 以约 5M 参数实现 78.6%的 TOP-1 准确性,节省 11%的参数和 13%的计算成本,但准确性提高了 0.2%,并且推理速度更快(在基于 ARM 的 Rockchip RK3288 上)。仅使用了 DeIT 0.5 倍参数量,但与 DeiT 相比,准确率提升了 2.7%。在 MS-CCOCO 目标检测和 Pascal VOC 分割任务上,EdgeFormer 还显示出更好的性能。

 

ViT 和 ConvNet


二者都是必不可少的。

  • 从实用角度而言,二者各有优缺点:

    • ViT 模型逐渐体现出更好的性能,但是通常面领着高昂的计算成本,并且难以训练。

    • 相较于 ViT 类结构,卷积网络显示出了稍差一点的性能,但是仍然有着独特的优势。例如卷积网络有着更好的硬件支持,容易训练。另外,卷积网络在针对移动或者小模型上仍然大有可为。

  • 从信息处理的视角而言,二者有着独特的特性:

    • ViT 架构擅长于提取全局信息,并且在数据驱动下,使用注意力机制从不同位置上提取信息。

    • 卷积网络专注于建模局部关系,并且有着更强的归纳偏置带来的先验信息。

经过这样的对比,作者们思考,是否可以从 ViT 的设计上学习一些东西来提升针对移动或者边缘设备的卷积网络的性能。这也正是本文的目的所在。

文中,作者们想要设计一个新的轻量的卷积网络,来进一步增强其在移动设备和边缘计算设备友好模型领域中的力量。纯卷积是更加移动设备友好的,因为卷积操作被现有的工具链已经高度优化,并且被广泛使用于部署模型到资源受限设备上。

甚至,由于这些年来卷积网络的流行,一些现有的神经网络加速器围绕着卷积风格的操作被设计出来,而并没有有效支持例如 softmax 这样复杂的非线性操作以及大矩阵乘法所需要的数据总线带宽数据。这些硬件和软件上的限制使得纯卷积的轻量模型相对于那些即使在其他方面具有同样竞争力的 ViT 类的模型仍然是更好的选择。

为了设计这样一个卷积网络,作者们比较并总结了 ViT 相较于卷积网络的三个主要不同:

  • ViT 更擅长于提取全局特征

  • ViT 应用了 meta-former block

  • 在 ViT 中,信息集成是数据驱动的(但是这种数据驱动的形式过于耗费内存,于是大量基于先验知识对此集成区域进行约束,压缩空间,从而提升 Attention 操作的效率)

为此,作者们针对这些点设计了 EdgeFormer 模块,来改进卷积网络:

  • 提出一个全局循环卷积(Global Circular Convolution,GCC)操作来提取全局特征从而改善传统卷积操作的有限的感受野。其中,基于实例的卷积核(这里想说的应该是卷积核的参数与输入相关)和位置嵌入策略分别被用来应对输入尺寸的变化,并且向输出特征图中注入位置信息。

  • 基于提出的 GCC,构建了一个纯卷积风格的 meta-former 模块作为基础结构。

  • 通过对 FFN 添加通道注意力(SE Block)来引入输入相关的核权重。

对于最终的模型 EdgeFormer,则是借鉴 CoatNet 和 MobileViT 的结构风格,采用了双叉形式的结构(bifurcate structure)。根据这些设计构建的最终模型在三个主流的视觉任务上的轻量模型的对比中获得了不错的性能表现。

 

从相关工作看论文


NLP 中的 Transformer 被 ViT 引入视觉任务后,出现了大量工作对其结构进行优化,以提升其计算效率和训练效率。

  • 从提升训练策略的角度来看,DeiT(Training data-efficient image transformers & distillation through attention)使用知识蒸馏策略来训练 ViT,从而使用更少的预训练数据实现了极具竞争力的性能。

     

  • 从提升模型架构的角度来看:

    • LeViT(Levit: a vision transformer in convnet’s clothing for faster inference)混合了卷积网络和 Transformer,实现了很好的速度和准确率的权衡。

    • BoTNet(Bottleneck transformers for visual recognition)则替换 ResNet 深层的卷积层为多头自注意力结构。

    • ViT-C(Early convolutions help transformers see better)则将原始的 ViT 的 patchify stem 结构(核大小等于跨步的卷积)替换为 convolutional stem 结构从而有效提升了 ViT 优化的稳定性和峰值性能。

    • ConViT(Convit: Improving vision transformers with soft convolutional inductive biases)则将通过门控位置自注意力合并了 soft convolutional inductive biases。

    • CMT(Cmt: Convolutional neural networks meet vision transformers)则应用了基于深度卷积的局部感知单元和一个轻量的 Transformer 模块(使用跨步等于卷积核的深度卷积缩减 K 和 V)。

    • CoatNet(Coatnet: Marrying convolution and attention for all data sizes)合并了卷积和自注意力来设计新的 transformer 模块(浅层使用卷积,深层使用额外引入全局静态权重的 Attention),这可以同时关注局部和全局信息。

    • PVT(Pyramid vision transformer: A versatile backbone for dense prediction without convolutions)和 CVT(Cvt: Introducing convolutions to vision transformers)在不同阶段中插入卷积操作来减少 token 的数量,构建了分层的多阶段结构。

    • Swin(Swin transformer: Hierarchical vision transformer using shifted windows)则是将 Attention 的计算约束在了固定尺寸的局部窗口之中。

    • PiT(Rethinking spatial dimensions of vision transformers)则是通过联合使用池化层和深度卷积层来实现特征图通道的扩张和空间的缩减。

    • 一些研究者试图通过借鉴卷积网络的思想来优化 ViT,他们将卷积网络的一些操作整合到 ViT 中。

    • 而另一些工作则是通过直接组合两种风格来设计新的混合架构。他们都相似地在模型起始阶段使用卷积 stem 来提取局部特征,而在后续的结构中使用 transformer 风格的结构提取“全局”或者“局部-全局”特征。本文提出的纯卷积架构的工作同样使用类似的思想。

 

针对本文着重探究的轻量模型的设计的领域,作者们也详细的回顾了卷积神经网络和视觉 Transformer 的一些工作。

  • 卷积神经网络:ShuffleNet 系列,MobileNet 系列,MicroNet,GhostNet,EfficientNet,TinyNet 和 MnasNet。相较于标准卷积网络,他们具有更少的参数,更低的计算成本,更快的推理速度。另外,轻量的卷积神经网络可以被应用到各种设备上。但是尽管有这样的优势,但是这些轻量模型性能稍差于更重的模型。

  • 卷积神经网络与视觉 Transformer 的组合:

    • Mobile-Former(Mobile-former: Bridging mobilenet and transformer)并行组合了 MobileNet 和 Transformer,利用前者提取局部特征,利用后者捕获全局信息。

    • MobileViT(Mobilevit: light-weight, general-purpose, and mobilefriendly vision transformer)中使用提出的 MobileViT 模块来替换 MobileNetV2 的深层模块。提出的模块中,通过卷积提取的局部表征,与全局表征被拼接来获得“局部-全局”表征。

 

就本文而言,提出的结构与 Mobile-Former 和 MobileViT 比较相关。不同点在于这两个方法仍然去保留 Transformer 块,而本文结构却是一个纯的卷积网络,这使得提出的结构对移动端更加友好。论文针对低算力平台的的实验也证实了这一点。在通过向视觉 Transformer 学习来设计一个纯卷积网络这一点上,本文和 ConvNeXt 有相近之处。主要有两点不同:

  • 想法和架构并不相同。ConvNeXt 通过将一系列增量但是有效的设计引入,从而将一个标准的 ResNet 朝着视觉 Transformer 的方向来改动。而本文则直接从三个卷积网络和视觉 Transformer 的差异的角度入手来设计网络,这从一个宏观的层面上填补了缺口。由于想法不同,对应的结构也就有所差异。

  • 提出的目的不同。本文主要针对移动设备。相较于 ConvNeXt,在轻量模型的设定下,本文的方法更有优势。

     

提出的方法


EdgeFormer Block

图中显示了卷积网络和 ViT 类结构主要差异。而图中下半部分则是提出的模块的结构。

 

使用全局循环卷积提取全局特征

卷积操作中,每个位置上的输出特征是通过对以其为中心的局部邻域上的输入特征进行加权聚合得到的。而 ViT 类结构中,自注意力模块提取特征则是通过对全局空间范围上的特征进行聚合后得到的,而且聚合权重是基于特征相似性生成的动态参数。相较于卷积的局部感受野,后者可以从整个空间维度上学习全局特征。

针对这点不足,本文提出了全局循环卷积 GCC。提出的 GCC 通过使用水平和垂直方向进行分解从而得到了一个单一来看足够轻量的操作。二者各自可以同时覆盖到向同行和相同列的所有输入位置。

具体而言,以垂直方向的 GCC 为例,即 GCC-V,其步骤如下:

  1. 输入特征为 CxHxW

  2. 引入一个基础位置嵌入信息为 CxBx1,通过扩展它,使其完全等于输入特征的形状,从而可以灵活地适应不同尺寸的输入特征。

    1. 通过双线性插值,变为 CxHx1

    2. 轴向复制扩展,变为 CxHxW

    3. 加到输入上

  3. 将得到的特征言 H 方向堆叠,得到 Cx(2H-1)xW,这步操作非常重要,由此便可以实现具有全局感受野且仍然参数共享的标准卷积操作了。

  4. 接下来开始构造卷积核参数,其同样由一个大小为 CxBx1 的基础量插值到 CxHx1 而得到。此时卷积核的空间感受野为 Hx1。搭配着 GCC-H 中的 Cx1xW,便近似得到了全局的覆盖。当然,这里同样可以引入偏置,具体的实现也同样是基于一个基础量扩展而成。

  5. 使用插值后的卷积核和嵌入位置信息的输入张量计算标准卷积操作,得到大小为 CxHxW 的输出。

这个流程中,GCC-V 中大小为 CxHx1 的卷积核的输出特征坐标 i,实际上对应着输入特征上的坐标范围 (i+t)%h,这里的 t 是一个 0 到 H-1 之间的值,不同的 t 对应着卷积核覆盖的局部邻域中的相对坐标。

也就是说,通过这样一个取模的计算,作者们构建出了一个“循环”的概念。可以说是非常会变通了。但是作者们的实现,却是更考验着思维上的转化了。因为针对这样卷积操作计算方式的改变,一般有两种实现策略,一种是从卷积运算的本身着手,另一种则是从被卷积的特征的角度去变换。这一点在我自己的一份关于动态滤波器的工作中也有所体现。

一般而言,前者虽然直接但是实现却非常复杂,可能需要重新写算子。而后者则更考验思维变通的能力。本文的实现思路就是后者。所以更具参考和启发意义。

这部分结构主要实现可见:

//github.com/hkzhang91/EdgeFormer/blob/df5ab09c7f60187d022a891ad083f4596838ff54/cvnets/modules/edgeformer_block.py#L534-L677

具体而言,作者们直接将两个特征沿着指定轴向维度进行拼接(注意这里后面的特征只取了 H-1 行的数据,不然输出的特征形状就比输入大 1 了):torch.cat((x_1_r, x_1_r[:, :, :-1, :]), dim=2)

相较于传统卷积操作,提出的 GCC 引入了三处修改:

  1. 感受野扩大到了全局空间。要注意和直接扩大传统局部卷积核到全局范围的设定区分开。因为传统卷积需要引入 padding 操作来维持特征尺寸,如果直接增大卷积核,绝大多数位置上的处理过程中,仍然只有部分像素被覆盖到。尤其是在边缘区域。大部分区域都是 0 值。

  2. 位置嵌入被用来保证输出特征对于空间位置的敏感性。GCC 可以提取全局特征,但是也打乱了原始输入的空间结构(因为循环的存在,所以也就无法严格按照局部邻域的本身的相对顺序来处理了)。对于分类任务而言,空间结构的保持可能不是一个大的问题,但是文中的实验中也展现了分割检测这样的密集预测任务对于保持空间结构的需求还是有的。也就是说位置信息的嵌入对于这样的任务而言有一定的益处。

  3. 卷积核和位置嵌入是根据输入尺寸动态插值生成的。GCC 中,卷积核的尺寸和位置嵌入和输入是一致的,这增强了对于不同尺寸输入的适应能力。

     

使用 GCC 构建 EdgeFormer 块

从卷积网络到 ViT 架构,一个相当重要的修改就是 meta-former block 替换了残差块。meta-former block 一般包含着两个核心组件:token mixer 和 channel mixer。前者用于在不同空间位置上的 token 之间交换信息,后者则是重点关注通道的交互。而且二者都会应用残差结构。

受此启发,GCC 被插入 meta-former 风格的结构中,从而构建了 EdgeFormer block。其中替换自注意力模块为 GCC,从而作为一个新的空间处理单元。这主要有两点原因:

  1. GCC 可以提取全局特征,并在不同像素之间进行信息交互,这满足 token mixer 的要求。

  2. 自注意力模块的计算复杂度是平方级别的,通过改用 GCC,可以显著降低计算成本,有助于实现设计轻量的卷积网络。

     

向 Channel Mixer 中添加通道注意力

ViT 架构中,自注意力模块可以根据输入调整权重,这造就了 ViT 这来数据驱动的模型。通过引用注意力机制,数据驱动模型可以专注于重要的特征,并且抑制不必要的信息,这带来了更好的性能。曾经的一些工作(例如 SENet 等)也强调了保持数据驱动的重要性。

通过替换自注意力操作为 GCC 结构,我们得到了一个可以提取全局特征的纯卷积网络。但是替换后的模型不再是数据驱动的结构了,为此,文中向 Channel Mixer 中引入了通道注意力。这里参考了 SENet 的结构。

 

EdgeFormer

通过为提出的 EdgeFormer block 选择一个外框架(outer frame),构建了本文的模型 EdgeFormer。

现存的混合架构基本可以分为三类,包括序列结构、并行结构和双叉式(bifurcate)结构。其中双叉式结构目前获得了最好的效果,典型如 CoatNet,而对于轻量结构,MobileViT 也采用了这一形式。因此本文也延续了这一设计,基于 MobileViT 构建了最终的外框架。具体而言,针对 baseline MobileViT,本文做出了如下改进:

  1. MobileViT 包含两个主要类型的模块,ViT 和 MobileNetV2。本文保持了所有 MobileNetV2 模块,替换了 ViT 结构。这一替换使得最终模型从混合架构变成了纯卷积架构。

  2. 适当地增加了 EdgeFormer 块的宽度。即便如此,替换的模型仍然具有较少的参数和较少的计算成本。

  3. 现有的双叉结构包含着一些交互模块,这负责本地和全局特征模块之间的信息交互。原始 MobileViT 中,ViT 块是最重的结构,经过替换后,这些交互结构反而成了大头。因此,本文将分组卷积和点卷积引入这些模块来减少了参数数量还不会伤害性能。

     

实验结果


性能对比

作者们在分类检测分割上都验证了提出模型和组件的性能。

在轻量模型的对比中,可以看到混合架构超越了以往的纯卷积架构和 ViT 类架构,而本文的升级版的纯卷积架构获得了整体更好的效果。

 

消融实验

这里的消融实验在三个任务上验证了提出组件的效果,这里的 baseline 为 MobileViT。主要做了以下几个方面的对比:

  • 用于替换自注意力的(对比消融 GCC)卷积核:普通卷积、两种尺度的大核卷积、GCC

    • 这里一开始就引入了通道注意力,似乎不太满足控制变量的要求

    • 2、3、6 的对比可见 GCC 的优势。

    • 这里在大核卷积的实验中(2、3),随着卷积核的增大性能区域饱和(倒是再增大是否会下降就很难说,只能说可能有这个趋势。这一点现象据说在 ConvNeXt 中也有反映)。而直接覆盖了全局的 GCC 却可以获得更好的表现。

  • meta-former 架构的应用

    • 这里为了替换 meta-former 结构,作者们将 GCC 集成到 ResNeXt block 中作为替代品。

    • 4、7 的对比可以看到使用 meta-former 构建纯卷积网络的潜力。

  • 通道注意力的应用

    • 5、7 的对比可以看到通道注意力的性能增益。

  • 位置嵌入的引入

    • 位置嵌入在检测和分割这样的任务中有着积极的意义。这可能也是因为二者相较于分类任务,对于空间位置更加敏感。

       

低功率平台的推理速度

作者们在 Rockchip RK3288 上对 EdgeFormer-S 和基线模型 MobileViT-S 模型进行了对比。通过使用 ONNX 和 MNN 将这些模型迁移到 RK3288 上,对他们迭代一百次从而测试平均推理速度。

对比中可以看到,EdgeFormer 快了 23%。除了更少的 Madd 操作,作者们认为速度提升可能还因为:

  1. 卷积已经被现有的广泛用于部署模型到这些资源有限的设备上的工具链高度优化。

  2. 相较于卷积,Transformer 需要更多的数据带宽来计算 attention map,这涉及到两个大的矩阵 K 和 Q,然而在卷积中,核相对于输入特征而言,是一个很小的矩阵。如果带宽需求超过了芯片设计,那么 CPU 将闲置来等待数据,从而导致 CPU 利用率较低和总体推理速度较慢。

欢迎关注公众号CV技术指南,专注于计算机视觉的技术总结、最新技术跟踪、经典论文解读、CV招聘信息。

CV技术指南创建了一个免费的知识星球。关注公众号添加编辑的微信号可邀请加入。

征稿通知

欢迎可以写以下内容的朋友联系我(微信号:“FewDesire”)。

  1. TVM入门到实践的教程
  2. TensorRT入门到实践的教程
  3. MNN入门到实践的教程
  4. 数字图像处理与Opencv入门到实践的教程
  5. OpenVINO入门到实践的教程
  6. libtorch入门到实践的教程
  7. Oneflow入门到实践的教程
  8. Detectron入门到实践的教程
  9. CUDA入门到实践的教程
  10. caffe源码阅读
  11. pytorch源码阅读
  12. 深度学习从入门到精通(从卷积神经网络开始讲起)
  13. 最新顶会的解读。例如最近的CVPR2022论文。
  14. 各个方向的系统性综述、主要模型发展演变、各个模型的创新思路和优缺点、代码解析等。
  15. 若自己有想写的且这上面没提到的,可以跟我联系。

声明:有一定报酬,具体请联系详谈。若有想法写但觉得自己能力不够,也可以先联系本人(微信号:FewDesire)了解。添加前请先备注“投稿”。

其它文章

招聘 | ​迁移科技招聘深度学习、视觉、3D视觉、机器人算法工程师等多个职位

计算机视觉入门路线

YOLO系列梳理(一)YOLOv1-YOLOv3

YOLO系列梳理(二)YOLOv4

YOLO系列梳理(三)YOLOv5

Attention Mechanism in Computer Vision

从零搭建Pytorch模型教程(三)搭建Transformer网络

从零搭建Pytorch模型教程(二)搭建网络

从零搭建Pytorch模型教程(一)数据读取

StyleGAN大汇总 | 全面了解SOTA方法、架构新进展

一份热力图可视化代码使用教程

一份可视化特征图的代码

工业图像异常检测研究总结(2019-2020)

关于快速学习一项新技术或新领域的一些个人思维习惯与思想总结