YOLO算法改进Backbone系列之:EdgeViT

摘要：在计算机视觉领域，基于Self-attention的模型(如ViTs)已经成为CNN之外的一种极具竞争力的架构。尽管越来越强的变种具有越来越高的识别精度，但由于Self-attention的二次复杂度，现有的ViT在计算和模型大小方面都有较高的要求。虽然之前的CNN的一些成功的设计选择(例如，卷积和分层结构)已经被引入到最近的ViT中，但它们仍然不足以满足移动设备有限的计算资源需求。这促使人们最近尝试开发基于最先进的MobileNet-v2的轻型MobileViT，但MobileViT与MobileNet-v2仍然存在性能差距。

在这项工作中，作者进一步推进这一研究方向，引入了EdgeViTs，一个新的轻量级ViTs家族，也是首次使基于Self-attention的视觉模型在准确性和设备效率之间的权衡中达到最佳轻量级CNN的性能。这是通过引入一个基于Self-attention和卷积的最优集成的高成本的local-global-local (LGL) 信息交换瓶颈来实现的。对于移动设备专用的评估，不依赖于不准确的proxies，如FLOPs的数量或参数，而是采用了一种直接关注设备延迟和能源效率的实用方法。

基于自注意力机制的视觉Transformer(ViT)在视觉任务上，已经形成和CNN一样强有力的架构，但其计算量和模型大小很大。虽然一些工作通过引入先验信息或级联多阶段结构到ViT中，但在移动设备上仍不够高效。本文研究基于MobileNetV2的轻量化ViT，其通过引入局部-全局-局部(LGL)的bottleneck实现，其结合了注意力机制和CNN的优势。

作者提出了将VIT模型应用于移动端需要考虑的三个问题：

l 推理速度要快。当前的一些指标如 FLOPs 难以反映模型在移动端的速度，因为内存访问速度、并行性等因素还要综合考虑

l 模型可以大；当前手机可以拥有32GB内存，存储模型并应该做为限制因素

l 实现的友好性。SWIN里的 cyclic shift 不便于在移动端实现，因此模型要考虑是否便于在移动端实现

以上面三个原则为指导，作者提出了 EdgeViTs，设计了一个高效的局部-全局-局部（LGL）模块，能够实现更好的准确性和计算效率。为了设计适用于移动/边缘设备的轻量级ViT，作者采用了最近ViT变体中使用的分层金字塔结构。Pyramid Transformer模型通常在不同阶段降低了空间分辨率同时也扩展了通道维度。每个阶段由多个基于Transformer Block处理相同形状的张量，类似ResNet的层次设计结构。基于Transformer Block严重依赖于具有二次复杂度的Self-attention操作，其复杂度与视觉特征的空间分辨率呈2次关系。通过逐步聚集空间Token，Pyramid Transformer可能比各向同性模型（ViT）更有效。

在这项工作中，作者深入到Transformer Block，并引入了一个比较划算的Bottlneck，Local-Global-Local(LGL)。LGL通过一个稀疏注意力模块进一步减少了Self-attention的开销，实现了更好的准确性-延迟平衡。

Self-attention已被证明是非常有效的学习全局信息或长距离空间依赖性的方法，这是视觉识别的关键。另一方面，由于图像具有高度的空间冗余(例如，附近的Patch在语义上是相似的)，将注意力集中到所有的空间Patch上，即使是在一个下采样的特征映射中，也是低效的。因此，与以前在每个空间位置执行Self-attention的Transformer Block相比，LGL Bottleneck只对输入Token的子集计算Self-attention，但支持完整的空间交互，如在标准的Multi-Head Self-attention (MHSA)中。这样既会减少Token的作用域，同时也保留建模全局和局部上下文的底层信息流。

为了实现这一点，作者将Self-attention分解为连续的模块，处理不同范围内的空间Token。这里引入了3种有效的操作：

（1）Local aggregation：仅集成来自局部近似Token信号的局部聚合；对于每个Token，利用Depth-wise和Point-wise卷积在大小为k×k的局部窗口中聚合信息

（2）Global sparse attention：建模一组代表性Token之间的长期关系，其中每个Token都被视为一个局部窗口的代表；对均匀分布在空间中的稀疏代表性Token集进行采样，每个r×r窗口有一个代表性Token。这里r表示子样本率，然后只对这些被选择的Token应用Self-attention。这与所有现有的ViTs不同，在那里所有的空间Token都作为Self-attention计算中的query被涉及到。

（3）Local propagation：将委托学习到的全局上下文信息扩散到具有相同窗口的非代表Token。通过转置卷积将代表性Token中编码的全局上下文信息传播到它们的相邻的Token中。

将这些结合起来，LGL Bottleneck就能够以低计算成本在同一特征映射中的任何一对Token之间进行信息交换。

EdgeViT模型变体配置信息：

在YOLOv5项目中添加模型作为Backbone使用的教程：

（1）将YOLOv5项目的models/yolo.py修改parse_model函数以及BaseModel的_forward_once函数

（2）在models/backbone（新建）文件下新建Edgevit.py，添加如下的代码：

（3）在models/yolo.py导入模型并在parse_model函数中修改如下（先导入文件）：

（4）在model下面新建配置文件：yolov5_edgevit.yaml

（5）运行验证：在models/yolo.py文件指定--cfg参数为新建的yolov5_edgevit.yaml