论文阅读——CMT

目标：优化transformer模型，使其复杂度低于transformer且性能优于高性能卷积模型。

创新点

1、提出了一种新颖的 CMT（CNNs meet transformers）架构用于视觉识别
2、 提出了局部感知单元（LPU）和反向残差前馈网络（IRFFN）

Introduction

Transformer在视觉任务中性能出色，但是性能低于同等大小下的CNN。原因有三：1）在ViT 和其他基于Transformer的模型中，将图像分成patch，它忽略了基于序列的NLP任务和基于图像的视觉任务之间的根本区别，即在每个patch分块中的2维结构和空间局部信息。 2）由于patch大小的固定，Transformer难以明确提取低分辨率和多尺度的特征，这给检测和分割等密集的预测任务带来了很大的挑战。 3）与基于CNN的O（NC^2） 相比Transformer中自注意力模块的计算和存储成本是输入分辨率的平方倍（O(NC^2）,使用Transformer来处理这些图像将不可避免地会导致GPU内存不足和计算效率低的问题。

输入图像首先经过stem进行细粒度特征提取，然后将其送入CMT块堆栈进行表示学习。具体而言，引入的CMT块是变压器块的改进变体，其局部信息通过深度卷积增强。相比于ViT，CMT第一阶段生成的特征可以保持更高的分辨率，即H/4×W/4/ViT中的H/16×W/16，此外，采用了类似于cnn的分阶段架构设计，采用步长为2的4个卷积层，逐步降低分辨率(序列长度)，灵活增加维数。分阶段设计有助于提取多尺度特征，减轻高分辨率带来的计算负担。CMT块中的局部感知单元(LPU)和反向残差前馈网络(IRFFN)可以同时捕获中间特征中的局部和全局结构信息，提高网络的表示能力。最后，利用平均池化代替ViT中的cls token，获得更好的分类效果。
给定一个输入图像，可以得到四个不同分辨率的层次特征图，类似于典型的cnn。利用上述特征图的输入步幅分别为4、8、16和32，我们的CMT可以获得输入图像的多尺度表示，并可以很容易地应用于目标检测和语义分割等下游任务。

相关工作

CNN

网络	特点
LeNet	-
AlexNet & VGGNet & GoogleNet & InceptionNet	-
ResNet	增强泛化能力
SENet	自适应地重新校准通道特征响应

Vision Transformer

网络	特点
ViT	将NLP中的Transformer引入到CV领域
DeiT	蒸馏的方式，使用teacher model来指导Vision Transformer的训练
T2T-ViT	递归的将相邻的tokens聚合成为一个，transformer对visual tokens进行建模
TNT	outer block建模patch embedding 之间的关系，inner block建模pixel embedding之间的关系
PVT	将金字塔结构引入到 ViT 中，可以为各种像素级密集预测任务生成多尺度特征图。
CPVT和CvT	cnn引入到transformer之中

CMT

整体结构
CMT stem(减小图片大小，提取本地信息)
Conv Stride(用来减少feature map,增大channel)
CMT block(捕获全局和局部关系)

CMT模型的核心是CMT Block，其主要由以下几个模块构成：

1）Local Perception Unit（局部感知单元）

• 位置编码会破坏卷积中的平移不变性，忽略了patch之间的局部信息和patch内部的结构信息。
• LPU来缓解这个问题。

2）Lightweight Multi-head Self-attention（轻量级的多头注意力机制）

• 使用深度卷积来减少KV的大小，加入相对位置偏置，构成了轻量级的自注意力计算。

3）Inverted Residual Feed-forward Network（反向残差前馈网络）

• 深度卷积增强局部信息的提取，残差结构来促进梯度的传播能力。
  
  1)Local Perception Unit（局部感知单元）
  本质就是，将输入图片信息，与 3*3 的卷积操作后相加，旨在增加了空间信息，可以和 ViT 的绝对位置编码的对应理解。
  
  
  2）Lightweight Multi-head Self-attention （轻量级的多头注意力机制）
  在原始的self-attention模块中，输入 X 被线性变换为 query、key、value 再进行计算，运算成本高。
  

此模块主要功能就是使用深度卷积计算代替了 key 和 value 的计算，从而减轻了计算开销。
在这里插入图片描述
3）Inverted Residual Feed-forward Network （反向残差前馈网络）
向残差前馈网络（IRFFN）类似于反向残差块，由扩展层、深度卷积和投影层组成。具体来说，我们改变了快捷方式连接的位置，以获得更好的性能。 其中，省略了激活层。深度卷积用于提取局部信息，而额外的计算成本可以忽略不计。插入快捷方式的动机与经典的残差网络相似，可以提高梯度跨层的传播能力。

代码

LPU

class LocalPerceptionUint(nn.Module):
    def __init__(self, dim, act=False):
        super(LocalPerceptionUint, self).__init__()
        self.act = act 
        # 增强本地信息的提取
        self.conv_3x3_dw = ConvDW3x3(dim)
        if self.act:
            self.actation = nn.Sequential(
                nn.GELU(),
                nn.BatchNorm2d(dim)
            )
    def forward(self, x):
        if self.act:
            out = self.actation(self.conv_3x3_dw(x))
            return out 
        else:
            out = self.conv_3x3_dw(x)
            return out 

IRFFN

class InvertedResidualFeedForward(nn.Module):
    def __init__(self, dim, dim_ratio=4.):
        super(InvertedResidualFeedForward, self).__init__()
        output_dim = int(dim_ratio * dim)
        self.conv1x1_gelu_bn = ConvGeluBN(
            in_channel=dim,
            out_channel=output_dim,
            kernel_size=1,
            stride_size=1,
            padding=0
        )
        self.conv3x3_dw = ConvDW3x3(dim=output_dim)  
        self.act = nn.Sequential(
            nn.GELU(),
            nn.BatchNorm2d(output_dim)
        )
        self.conv1x1_pw = nn.Sequential(
            nn.Conv2d(output_dim, dim, 1, 1, 0),
            nn.BatchNorm2d(dim)
        )
        
    def forward(self, x):
        x = self.conv1x1_gelu_bn(x)
        out = x + self.act(self.conv3x3_dw(x))
        out = self.conv1x1_pw(out)
        return out 

LMHSA

class LightMutilHeadSelfAttention(nn.Module):
    """calculate the self attention with down sample the resolution for k, v, add the relative position bias before softmax
    Args:
        dim (int) : features map channels or dims 
        num_heads (int) : attention heads numbers
        relative_pos_embeeding (bool) : relative position embeeding 
        no_distance_pos_embeeding (bool): no_distance_pos_embeeding
        features_size (int) : features shape
        qkv_bias (bool) : if use the embeeding bias
        qk_scale (float) : qk scale if None use the default 
        attn_drop (float) : attention dropout rate
        proj_drop (float) : project linear dropout rate
        sr_ratio (float)  : k, v resolution downsample ratio
    Returns:
        x : LMSA attention result, the shape is (B, H, W, C) that is the same as inputs.
    """
    def __init__(self, dim, num_heads=8, features_size=56, 
                relative_pos_embeeding=False, no_distance_pos_embeeding=False, qkv_bias=False, qk_scale=None, 
                attn_drop=0., proj_drop=0., sr_ratio=1.):
        super(LightMutilHeadSelfAttention, self).__init__()
        assert dim % num_heads == 0, f"dim {
      dim} should be divided by num_heads {
      num_heads}"
        self.dim = dim 
        self.num_heads = num_heads
        head_dim = dim // num_heads   # used for each attention heads
        self.scale = qk_scale or head_dim ** -0.5

        self.relative_pos_embeeding = relative_pos_embeeding
        self.no_distance_pos_embeeding = no_distance_pos_embeeding

        self.features_size = features_size

        self.q = nn.Linear(dim, dim, bias=qkv_bias)
        self.kv = nn.Linear(dim, dim*2, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)
        self.softmax = nn.Softmax(dim=-1)

        self.sr_ratio = sr_ratio
        if sr_ratio > 1:
            self.sr = nn.Conv2d(dim, dim, kernel_size=sr_ratio, stride=sr_ratio)
            self.norm = nn.LayerNorm(dim) 
        
        if self.relative_pos_embeeding:
            self.relative_indices = generate_relative_distance(self.features_size)
            self.position_embeeding = nn.Parameter(torch.randn(2 * self.features_size - 1, 2 * self.features_size - 1))
        elif self.no_distance_pos_embeeding:
            self.position_embeeding = nn.Parameter(torch.randn(self.features_size ** 2, self.features_size ** 2))
        else:
            self.position_embeeding = None

        if self.position_embeeding is not None:
            trunc_normal_(self.position_embeeding, std=0.2)

    def forward(self, x):
        B, C, H, W = x.shape 
        N = H*W
        x_q = rearrange(x, 'B C H W -> B (H W) C')  # translate the B,C,H,W to B (H X W) C
        q = self.q(x_q).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)   # B,N,H,DIM -> B,H,N,DIM
        
        # conv for down sample the x resoution for the k, v
        if self.sr_ratio > 1:
            x_reduce_resolution = self.sr(x)
            x_kv = rearrange(x_reduce_resolution, 'B C H W -> B (H W) C ')
            x_kv = self.norm(x_kv)
        else:
            x_kv = rearrange(x, 'B C H W -> B (H W) C ')
        
        kv_emb = rearrange(self.kv(x_kv), 'B N (dim h l ) -> l B h N dim', h=self.num_heads, l=2)         # 2 B H N DIM
        k, v = kv_emb[0], kv_emb[1]
        
        attn = (q @ k.transpose(-2, -1)) * self.scale    # (B H Nq DIM) @ (B H DIM Nk) -> (B H NQ NK)
        
        # TODO: add the relation position bias, because the k_n != q_n, we need to split the position embeeding matrix
        q_n, k_n = q.shape[1], k.shape[2]
       
        if self.relative_pos_embeeding:
            attn = attn + self.position_embeeding[self.relative_indices[:, :, 0], self.relative_indices[:, :, 1]][:, :k_n]
        elif self.no_distance_pos_embeeding:
            attn = attn + self.position_embeeding[:, :k_n]

        attn = self.softmax(attn)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)  # (B H NQ NK) @ (B H NK dim)  -> (B NQ H*DIM)
        x = self.proj(x)
        x = self.proj_drop(x)
            
        x = rearrange(x, 'B (H W) C -> B C H W ', H=H, W=W)
        return x 

总结

本文提出了一种新的混合结构CMT，用于视觉识别和其他下游计算机视觉任务，如目标检测和实例分割，并解决了在计算机视觉领域以野蛮力方式使用变压器的局限性。所提出的CMT结构利用cnn和变压器来捕获局部和全局信息，提高了网络的表示能力。