【ARXIV2205】EdgeViTs: Competing Light-weight CNNs on Mobile Devices with Vision Transformers

【ARXIV2205】EdgeViTs: Competing Light-weight CNNs on Mobile Devices with Vision Transformers

The paper ：https://arxiv.org/abs/2205.03436

Code ：https://github.com/1hunters/EdgeViT

Vision based on self attention mechanism Transformer(ViT) On visual tasks , Has formed and CNN The same powerful architecture , But its calculation and model size are very large . Although some work is done by introducing prior information or cascading multistage structures to ViT in , But it's still not efficient enough on mobile devices . This paper is based on MobileNetV2 Lightweight ViT, By introducing local - overall situation - Local (LGL) Of bottleneck Realization , It combines attention mechanisms with CNN The advantages of .

The author proposes that VIT Three issues need to be considered when applying the model to mobile terminals ：

• 1） Reasoning speed should be fast . Some current indicators such as FLOPs It is difficult to reflect the speed of the model at the mobile end , Because memory access speed 、 Parallelism and other factors need to be comprehensively considered .
• 2） The model can be large . Current mobile phones can have 32GB Of memory , The storage model should not be a limiting factor .
• 3） Achieved friendliness . SWIN Inside cyclic shift It is not convenient to realize on the mobile end , Therefore, the model should consider whether it is convenient to implement in the mobile terminal .

Guided by the above three principles , The author puts forward EdgeViTs, Designed an efficient Local - overall situation - Local （LGL） modular , It can achieve better accuracy and calculation efficiency .

The model is shown above , The point is LGL modular , It includes three key parts ：

• local aggregation： By convolution sum depth conv form
• global sparse attention: Calculate attention after average pooling
• local propagation: Use deconvolution to restore the reduced feature map to its original size .

Let's see the specific code , It's not difficult to understand .

class LocalAgg():
    def __init__(self, dim):
        self.conv1 = Conv2d(dim, dim, 1)
        self.conv2 = Conv2d(im, dim, 3, padding=1, groups=dim)
        self.conv3 = Conv2d(dim, dim, 1)
        self.norm1 = BatchNorm2d(dim)
        self.norm2 = BatchNorm2d(dim)
     forward(self, x):
        x = self.conv1(self.norm1(x))
        x = self.conv2(x)
        x = self.conv3(self.norm2(x))
        return x

class GlobalSparseAttn():
    def __init__(self, dim, sample_rate, scale):
        self.scale = scale
        self.qkv = Linear(dim, dim * 3)
        self.sampler = AvgPool2d(1, stride=sample_rate)
        kernel_size=sr_ratio
        self.LocalProp = ConvTranspose2d(dim, dim, kernel_size, stride=sample_rate, groups=dim
        )
        self.norm = LayerNorm(dim)
        self.proj = Linear(dim, dim)
    def forward(self, x):
        x = self.sampler(x)
        q, k, v = self.qkv(x)
        attn = q @ k * self.scale
        attn = attn.softmax(dim=-1)
        x = attn @ v
        x = self.LocalProp(x)
        x = self.proj(self.norm(x))
        return x

Actually , The whole network is based on CNN Of , It's just used SWIN The typical architecture of . The experimental results are shown in the following table . Although the author said , And MobileViTs comparison ,EdgeViTs Under three kinds of complexity settings 5.4%、2.8% and 2.7% Raise , But I feel from FLOPs And other indicators , There is no advantage of the dating . Here is my personal understanding , Where there are different opinions, you can communicate at any time .