笔记02----重新思考轻量化视觉Transformer中的局部感知CloFormer（即插即用）

1. 基本信息

论文标题: 《Rethinking Local Perception in Lightweight Vision Transformer》
中文标题: 《重新思考轻量化视觉Transformer中的局部感知》
作者单位: 清华大学
发表时间: 2023
论文地址: https://arxiv.org/abs/2303.17803
代码地址: https://github.com/qhfan/CloFormer

2. 应用场景

图像分类、目标检测、语义分割等领域。

3. 研究背景

现阶段，Transformer在图像分类、目标检测、语义分割等领域表现出优异的性能。然而Transformer参数量和计算量太大，不适合部署到移动设备。
在现有的轻量级Transformer模型中，大多数方法只注重设计稀疏注意力以有效处理低频全局信息，而处理高频局部信息的方法相对简单。

4. 方法概述

为了同时利用共享权重和上下文感知权重的优势，提出了CloFormer，这是一种具有上下文感知局部增强功能的轻量级视觉转换器，具体贡献如下：

在 CloFormer 中，引入了一种名为 AttnConv 的卷积算子，它采用注意力机制，充分利用共享权重和上下文感知权重的优势来实现局部感知。此外，它使用了一种新方法，该方法结合了比普通局部自注意力更强的非线性来生成上下文感知权重。
在CloFormer 中，采用双分支架构，其中一个分支使用AttnConv 捕获高频信息，而另一个分支使用带有下采样的普通注意力捕获低频信息。双分支结构使 CloFormer 能够同时捕获高频和低频信息。
该方法在图像分类、目标检测和语义分割方面的广泛实验证明了 CloFormer 的优越性。 CloFormer 在 ImageNet1k 上仅用 4.2M 参数和 0.6G FLOP 就实现了 77.0% 的准确率，明显优于其他模型。

4.1 整体网络结构

在这里插入图片描述

如上图所示，CloFormer包含一个卷积主干和四个阶段。每个阶段由Clo block和ConvFFN组成，先通过卷积主干传递输入图像以获得tokens。该系统由四个卷积组成，每个卷积的步幅分别为2、2、1和1。随后，标记经过四个阶段的Clo block和ConvFFN来提取层次特征。最后，利用全局平均池化和全连接层来生成预测。

ConvFFN

为了将局部信息整合到FFN过程中，用ConvFFN取代了传统的FFN。 ConvFFN和常用的FFN之间的主要区别是，ConvFFN在GELU激活后使用深度卷积（DWconv），这使得ConvFFN能够聚合局部信息。由于DWconv，下行采样可以直接在ConvFFN中执行，而无需引入PatchMerge模块。 CloFormer使用了两种类型的ConvFFN。第一种是级内ConvFFN，它直接利用跳过连接。另一个是连接两个阶段的ConvFFN。在这种类型的ConvFFN的跳过连接中，使用DWconv和全连接层分别对输入进行下采样和上维。

在这里插入图片描述

Clo block

每个块由一个本地分支和一个全局分支组成。在全局分支中，首先对K和V进行下采样，然后对Q、K和V进行标准attention处理，提取低频全局信息。

在这里插入图片描述

4.2 AttnConv模块

全局分支的模式有效地减少了需要注意的flop的数量，也产生了一个全局接受野。然而，它在有效捕获低频全局信息的同时，对高频局部信息的处理能力不足。
在AttnConv中，首先应用线性变换得到Q，K， V，这与标准注意力相同，在进行线性变换后，首先对V进行共享权值的局部特征聚合处理，然后基于处理后的V和Q， K进行上下文感知的局部增强。具体分为为如下三个步骤：

Local Feature Aggregation

使用一个简单的深度卷积（DWconv）来对 V 进行局部信息聚合。

Context-aware Local Enhancement

使用两个DWconv分别聚合Q和K的本地信息。然后，计算Q和K的Hadamard积，并对结果进行一系列变换，以获得−1到1之间的上下文感知权重。最后，利用生成的权值对局部特征进行增强。

Fusion with Global Branch

使用简单的方法将局部分支的输出与全局分支的输出融合。

4.3 代码

可以将Clo block当作注意力机制使用，具体代码如下：

import torch  
import torch.nn as nn  
from efficientnet_pytorch.model import MemoryEfficientSwish  # 从 EfficientNet 的库中引入高效激活函数 Swish  
class AttnMap(nn.Module):  
    def __init__(self, dim):  
        super().__init__()  
        # 定义一个包含两层卷积和激活函数的块，用于生成注意力映射  
        self.act_block = nn.Sequential(  
            nn.Conv2d(dim, dim, 1, 1, 0),  # 1x1 卷积，保持通道数不变  
            MemoryEfficientSwish(),       # Swish 激活函数  
            nn.Conv2d(dim, dim, 1, 1, 0)  # 再次使用 1x1 卷积  
        )  
  
    def forward(self, x):  
        return self.act_block(x)  # 前向传播，返回处理后的张量  
  
class CloAttention(nn.Module):  
    def __init__(self, dim, num_heads=8, group_split=[4, 4], kernel_sizes=[5], window_size=4,  
                 attn_drop=0., proj_drop=0., qkv_bias=True):  
        super().__init__()  
        # 参数初始化和断言检查  
        assert sum(group_split) == num_heads  # 确保分组的头总数等于注意力头总数  
        assert len(kernel_sizes) + 1 == len(group_split)  # 核大小和分组数一致  
  
        self.dim = dim  # 输入通道数  
        self.num_heads = num_heads  # 总的多头注意力头数  
        self.dim_head = dim // num_heads  # 每个头的通道数  
        self.scalor = self.dim_head ** -0.5  # 注意力缩放因子  
        self.kernel_sizes = kernel_sizes  # 高频分支的卷积核大小  
        self.window_size = window_size  # 低频分支窗口大小  
        self.group_split = group_split  # 每个分支分配的头数  
  
        # 创建高频和低频分支的模块  
        convs = []  # 高频卷积  
        act_blocks = []  # 高频注意力模块  
        qkvs = []  # 高频分支的 QKV 卷积  
  
        for i in range(len(kernel_sizes)):  
            kernel_size = kernel_sizes[i]  
            group_head = group_split[i]  
            if group_head == 0:  
                continue  # 如果分组头数为 0，跳过此分支  
            convs.append(nn.Conv2d(3 * self.dim_head * group_head, 3 * self.dim_head * group_head, kernel_size,  
                                   1, kernel_size // 2, groups=3 * self.dim_head * group_head))  # 高频卷积  
            act_blocks.append(AttnMap(self.dim_head * group_head))  # 注意力映射模块  
            qkvs.append(nn.Conv2d(dim, 3 * group_head * self.dim_head, 1, 1, 0, bias=qkv_bias))  # QKV 卷积  
  
        # 定义低频全局注意力分支  
        if group_split[-1] != 0:  
            self.global_q = nn.Conv2d(dim, group_split[-1] * self.dim_head, 1, 1, 0, bias=qkv_bias)  # Q 卷积  
            self.global_kv = nn.Conv2d(dim, group_split[-1] * self.dim_head * 2, 1, 1, 0, bias=qkv_bias)  # KV 卷积  
            self.avgpool = nn.AvgPool2d(window_size, window_size) if window_size != 1 else nn.Identity()  # 平均池化  
  
        # 将模块添加到 ModuleList 中  
        self.convs = nn.ModuleList(convs)  
        self.act_blocks = nn.ModuleList(act_blocks)  
        self.qkvs = nn.ModuleList(qkvs)  
        self.proj = nn.Conv2d(dim, dim, 1, 1, 0, bias=qkv_bias)  # 投影层  
        self.attn_drop = nn.Dropout(attn_drop)  # 注意力权重的 dropout        self.proj_drop = nn.Dropout(proj_drop)  # 输出的 dropout  
    def high_fre_attntion(self, x: torch.Tensor, to_qkv: nn.Module, mixer: nn.Module, attn_block: nn.Module):  
        '''  
        高频分支的注意力计算  
        x: (b c h w) 输入特征  
        '''        b, c, h, w = x.size()  
        qkv = to_qkv(x)  # 计算 QKV，得到 (b, 3*m*d, h, w)        qkv = mixer(qkv).reshape(b, 3, -1, h, w).transpose(0, 1).contiguous()  # 混合后得到 (3, b, m*d, h, w)        q, k, v = qkv  # 分解为 Q、K、V  
        attn = attn_block(q.mul(k)).mul(self.scalor)  # 计算缩放后的注意力  
        attn = self.attn_drop(torch.tanh(attn))  # 使用 tanh 激活并应用 dropout        res = attn.mul(v)  # 应用注意力权重到 V        return res  
  
  
    def low_fre_attention(self, x: torch.Tensor, to_q: nn.Module, to_kv: nn.Module, avgpool: nn.Module):  
        '''  
        低频分支的注意力计算  
        x: (b c h w) 输入特征  
        '''        b, c, h, w = x.size()  
        q = to_q(x).reshape(b, -1, self.dim_head, h * w).transpose(-1, -2).contiguous()  # 计算 Q 并调整形状为 (b, m, h*w, d)        kv = avgpool(x)  # 对输入特征进行平均池化  
        kv = to_kv(kv).view(b, 2, -1, self.dim_head, (h * w) // (self.window_size ** 2)).permute(1, 0, 2, 4, 3).contiguous()  # 计算 KV        k, v = kv  # 分解为 K、V  
        attn = self.scalor * q @ k.transpose(-1, -2)  # 计算缩放后的注意力  
        attn = self.attn_drop(attn.softmax(dim=-1))  # 对注意力进行 softmax 和 dropout        res = attn @ v  # 应用注意力权重到 V        res = res.transpose(2, 3).reshape(b, -1, h, w).contiguous()  # 调整形状为原始形状  
        return res  
  
  
    def forward(self, x: torch.Tensor):  
        '''  
        x: (b c h w) 输入特征  
        '''        res = []  # 保存各分支的输出  
        for i in range(len(self.kernel_sizes)):  
            if self.group_split[i] == 0:  
                continue  
            res.append(self.high_fre_attntion(x, self.qkvs[i], self.convs[i], self.act_blocks[i]))  # 高频分支输出  
        if self.group_split[-1] != 0:  
            res.append(self.low_fre_attention(x, self.global_q, self.global_kv, self.avgpool))  # 低频分支输出  
        return self.proj_drop(self.proj(torch.cat(res, dim=1)))  # 合并分支输出并应用投影  
  
  
# 输入 N C HW,  输出 N C H W
if __name__ == '__main__':  
    block = CloAttention(64).cuda()  # 初始化 CloAttention 模块  
    input = torch.rand(1, 64, 64, 64).cuda()  # 创建一个随机输入  
    output = block(input)  # 前向传播  
    print(f"Input_Size：{input.size()}\nOutput_Size：{output.size()}")  # 打印输入和输出的张量形状

5. 结果

表中报告了ImageNet1K分类结果。结果表明，当模型大小和FLOPs相似时，模型比以前的模型性能更好。其中，CloFormer-XXS仅使用4.2万个参数和0.6G FLOPs， Top-1准确率达到77.0%，分别超过ShuffleNetV22x、MobileViT-XS和EdgeViT-XXS 1.6%、2.2%和2.6%

在这里插入图片描述