将卷积引入transformer中VcT(Introducing Convolutions to Vision Transformers)的pytorch代码详解

1. Motivation:

​ Transformer已经被证明在很多视觉任务上可以取得不错的性能。最开始的ViT是最先将tansformer应用在分类任务上，并取得了不错的结果，但是它的问题在于需要先在大数据集上（1000万的私有数据）先进行预训练，才能在下游的中小数据集取得不错的结果，如果在同等规模的小数据集上进行训练，实际上ViT的性能是比不上的经典的卷积神经网络（Resnet,VGG等的）。而且ViT的计算随着Token序列的长度的呈指数增长。

​ 问题在于，CNN实际上是非常适合视觉任务，但是CNN不能够图像像素全局之间的长距离联系，而Transformer虽然非常擅长建立长距离联系，ViT首次用transformer建立图像像素的长距离联系，但是却没用CNN擅长建立图像局部关系的能力。

​ 有人认为，实际上CNN就是一种局部的self-attention。而self-attention是复杂化的CNN。

CNN考虑的是卷积核那样大小的感受野，而self-attention考虑的是一张图片那样大的范围。

CNN通过局部的感受野，共享卷积核权重，下采样，上采样等操作，将空间上相邻的那些信息都联系起来，而这些局部的信息是高度相关的，而ViT模型则没有用到这种性质，那么就可能需要更多的数据来训练，这可能就是它的缺陷所在。

思考：那么我们是否能将卷积融入Transformer中，使得Transformer也具有和CNN类似的建模图像空间局部的联系？

其实从理论上来讲，self-attention 强大的建模拟合能力，也是可以学习到和CNN类似的像素局部空间联系，但是ViT需要更多的数据才能将达到和CNN一样的效果，所以将CNN融入Transformer的一个最大好处是可以不需要那么大 的数据集。

那么以上这些就是本文的VcT,将卷积引入Transformer 的动机。

2. Method

本文的方法中，如上图所示主要由三个阶段构成，而每个阶段又是由两个比较重要的小模块而组成。这两个模块如下介绍：

2.1 Convolutional Token Embedding 模块

这个模块的输入是reshap成2D结构的token,主要做的是对这个2D的特征图进行一次卷积操作。卷积的目的在于保证每个阶段都能减小特征图的尺寸，增加特征图通道数，相应的将其reshape成token后，token的数量也会减少，但是token的维度会增加。这使得token能够在越来越大的空间足迹上表示越来越复杂的视觉模式，类似于cnn的特征层。

2.2 Convolutional Projection For Attention 模块

实际上这个模块，做的就是Transformer中的堆叠了多个block模块的事情。而每个block模块则是由self-attention+MLP模块构成。本文的这个模块实际最主要的改变就是对self-attention中的q,k,v矩阵是怎么来的，做了改动。原始的ViT中，q,k,v是通过Linear Projection操作得到，采用的是线性映射。如下图所示：

而本文为了将卷积引入到transformer中，这部分采用的是Convolutional Projection操作，即卷积操作。

细节是：首先将token重新reshape成2D的特征图，再分别通过3个深度可分离卷积（这里如果不理解深度可分离卷积可以去查查资料），得到3个q,k,v特征图。然后再将这三个特征图reshape成token,得到最终的q,k,v。如下公式;

$x_i^{q/k/v}=Flatten(Conv2d(Reshape2D(x_{token}),s))$

其中，Conv2d是深度可分离卷积操作。

然后将得到的token其送入MLP模块，这样一个block模块就构成了。而本模块则是由多个block构成。

一个block的具体结构如下所示：

上述block中的self-attention部分的具体实现过程的图示如下：

2.3 如何实现分类任务

这里采用的思想如ViT中实现分类任务是类似的，但是本文并没有在一开始对图像进行patch_embeding的时候，就加入cls_token，而是在第三个stage中加入cls_token,第二个阶段结束得到的2D特征图，送入Convolutional Token Embedding 模块得到图片数据的token,这是将cls_token和图片数据的token组合起来，送入Convolutional Projection For Attention 模块,输出时会将cls_token和图数据的token分开，用cls_token来进行分类工作。

2.4对于位置编码的思考？

CvT中没有使用显示的位置编码工作，这一点可能是应该CvT将self-attention中的linear Projection换成卷积操作所影响的，因为卷积操作本身就带有对位置编码的能力。

3.实验

3.1实验设置

3.1.1数据集：

1.  ImageNet-1k (1.3M images)，
2.  ImageNet (14M images，22k类)
3.  CIFAR-10/100
4.  Oxford-IIIT-Pet
5.  Oxford-IIIT-Flower

3.1.2 Model Variants（不同尺寸模型）

该表中，Conv.Embed表示Convolutional Token Embedding 模块；Conv.Proj表示Convolutional Projection 模块。

$H_i和D_i$分别表示MHSA(多头注意力机制中的)头的数量和embeding feature维度的大小。$R_i$是MLP中隐藏层特征维度的放大比例。

3.1.3优化器设置

​ AdamW。CvT-13 weight decay = 0.05，CvT-24 weight decay = 0.1。

3.2对比实验(消融实验)

3.2.1位置编码的影响

作者通过消融实验，分别在CvT的三个阶段分别加入位置编码，和不加位置编码对分类性能的影响，通过上表可以发现，CvT不使用位置编码不会影响性能。而DeiT若不使用位置编码则会掉点。根本原因还是如上分析的那样，CvT中因为引入了卷积操作，从而隐含了位置编码信息。

3.2.2Convolutional Token Embedding模块对CvT的影响

作者设计了4组实验，为了找出Convolutional Token Embedding对CvT的影响，作者将其替换为Patch embedding,从上表中可以看出,当使用Convolutional Token Embedding模块并且不使用位置编码，效果最好。当使用Patch embedding并同时使用位置编码，效果次之。

3.2.3Convolutional Projection模块对CvT的影响

1. 首先对比了Convolutional Projection模块中的self-attention中获取q,k,v时卷积操作的stride的大小对实验结果的影响。

​ 上表可以看出，当吧stride=1换为stride=2s时，计算量大幅度减小，但是准确率随之也下降了一点。

2. 将Convolutional Projection模块中的self-attention中q,k,v的映射，替换为传统的linear Project。

由上图可以看出，只有3个阶段都使用 Convolutional Projection 时，性能才是最佳的。

4.小结

CvT是如何将卷积操作映入Transformer的？

• 首先patch_embedding换为了卷积的方式（Convolutional Token Embedding），在每个阶段前都会执行这个操作。
• 然后在self-attention中的线性映射出q,k,v，被替换为卷积方式，具体实现是通过深度可分离卷积。
• 然后不再使用位置编码，主要是因为映入的卷积操作，而卷积操作中隐含了位置信息。

5.CvT的pytorch代码详解

###5.1 Convolutional Token Embedding 模块

#4.Convolutional Token Embedding 模块
'''
输入的是2维的图片数据或者2D的token(即将token  reshape成特征图)，输出的是特征图，目的是让特征图的尺寸变小
图片数据经过一个卷积层，输出一个特征图
'''
class ConvEmbed(nn.Module):
    """ Image to Conv Embedding
    """
    def __init__(self,patch_size=7,in_chans=3,embed_dim=64,stride=4,padding=2,norm_layer=None):
        super().__init__()
        patch_size = to_2tuple(patch_size)
        self.patch_size = patch_size#(s,s),一个patch的大小
        self.proj = nn.Conv2d(in_chans, embed_dim,kernel_size=patch_size,stride=stride,padding=padding)#卷积，得到特征图
        self.norm = norm_layer(embed_dim) if norm_layer else None

    def forward(self, x):#输入是特征图
        x = self.proj(x)#（b,c,h,w）
        B, C, H, W = x.shape
        x = rearrange(x, 'b c h w -> b (h w) c')#（b,h*w,c）
        if self.norm:
            x = self.norm(x)
        x = rearrange(x, 'b (h w) c -> b c h w', h=H, w=W)#（b,c,h,w）
        return x#输出是特征图

5.2 Convolutional Projection For Attention 模块

self-Attention模块和MLP模块构成一个编码器block,而 Convolutional Projection For Attention 模块是多个block堆叠而成。

5.2.1MLP模块

这个模块与Transformer原始的MLP模块是一样的，没有改动

该层主要由两个全连接层构成，输入是token，输出也是token

该层是一个编码器block的前馈神经网络，该层接在self-Attention模块的后面。

#1.MLp层，主要由两个全连接层组成，输入是token,输出是token
class Mlp(nn.Module):
    def __init__(self,in_features,hidden_features=None,out_features=None,act_layer=nn.GELU,drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)
    def forward(self, x):#输入的是token
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x#输出的是token

5.2.2 self-Attention模块

该self-attention模块是本论文中做的最大的改动，将原self-Attention模块中计算q,k,v的方法，由原本的线性映射，改为了卷积映射（深度可分离卷积）.

#2.self-attention层,输入是token,输出是token,
'''
主要内容：
输入是token,将token重新reshape成特征图，然后在特征图的基础进行卷积，
进行三次卷积得到q,k,v三个特征图，然后将这三个特征图再reshape成token,
然后用这个q,k这个两个token,进行attention操作，得到打分值score,
打分值score经过softmax变成概率值，然后与V这个token相乘，得到结果token
'''
class Attention(nn.Module):
    def __init__(self,
                 dim_in,dim_out,num_heads,
                 qkv_bias=False,attn_drop=0.,proj_drop=0.,
                 method='dw_bn',kernel_size=3,stride_kv=1,
                 stride_q=1,padding_kv=1,padding_q=1,
                 with_cls_token=True,**kwargs):
        super().__init__()
        self.stride_kv = stride_kv#1
        self.stride_q = stride_q#1
        self.dim = dim_out
        self.num_heads = num_heads
        # head_dim = self.qkv_dim // num_heads
        self.scale = dim_out ** -0.5
        self.with_cls_token = with_cls_token#True

        #卷积实现得到q,k,v
        self.conv_proj_q = self._build_projection(
            dim_in, dim_out, kernel_size, padding_q,
            stride_q, 'linear' if method == 'avg' else method
        )
        self.conv_proj_k = self._build_projection(
            dim_in, dim_out, kernel_size, padding_kv,
            stride_kv, method
        )
        self.conv_proj_v = self._build_projection(
            dim_in, dim_out, kernel_size, padding_kv,
            stride_kv, method
        )

        #扩充维度
        self.proj_q = nn.Linear(dim_in, dim_out, bias=qkv_bias)
        self.proj_k = nn.Linear(dim_in, dim_out, bias=qkv_bias)
        self.proj_v = nn.Linear(dim_in, dim_out, bias=qkv_bias)

        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim_out, dim_out)#全连接层
        self.proj_drop = nn.Dropout(proj_drop)

    # 卷积操作->BN->然后将卷积图（b,c,h,w）reshape成Token（b,h*w,c)
    def _build_projection(self,dim_in,dim_out,kernel_size,padding,stride,method):
        if method == 'dw_bn':
            proj = nn.Sequential(OrderedDict([
                ('conv', nn.Conv2d(dim_in,dim_in,kernel_size=kernel_size,
                    padding=padding,stride=stride,bias=False,groups=dim_in
                )),
                ('bn', nn.BatchNorm2d(dim_in)),
                ('rearrage', Rearrange('b c h w -> b (h w) c')),
            ]))
        elif method == 'avg':
            proj = nn.Sequential(OrderedDict([
                ('avg', nn.AvgPool2d(kernel_size=kernel_size,padding=padding,stride=stride,ceil_mode=True)
                 ),
                ('rearrage', Rearrange('b c h w -> b (h w) c')),
            ]))
        elif method == 'linear':
            proj = None
        else:
            raise ValueError('Unknown method ({})'.format(method))
        return proj

    def forward_conv(self, x, h, w):#输入是token
        if self.with_cls_token:
            cls_token, x = torch.split(x, [1, h*w], 1)

        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w)#将token转变成特征图

        if self.conv_proj_q is not None:
            q = self.conv_proj_q(x)#特征图->token，'b c h w -> b (h w) c'
        else:
            q = rearrange(x, 'b c h w -> b (h w) c')

        if self.conv_proj_k is not None:
            k = self.conv_proj_k(x)
        else:
            k = rearrange(x, 'b c h w -> b (h w) c')

        if self.conv_proj_v is not None:
            v = self.conv_proj_v(x)
        else:
            v = rearrange(x, 'b c h w -> b (h w) c')

        if self.with_cls_token:
            q = torch.cat((cls_token, q), dim=1)#加上分类的token
            k = torch.cat((cls_token, k), dim=1)
            v = torch.cat((cls_token, v), dim=1)
        return q, k, v#，输出的是token



    def forward(self, x, h, w):#输入x是token
        if (
            self.conv_proj_q is not None
            or self.conv_proj_k is not None
            or self.conv_proj_v is not None
        ):
            q, k, v = self.forward_conv(x, h, w)#'b c h w -> b (h w) c'

        q = rearrange(self.proj_q(q), 'b t (h d) -> b h t d', h=self.num_heads)#先扩宽token的维度，然后再实现mult-head，最后的结构是’b,h,t,d‘
        k = rearrange(self.proj_k(k), 'b t (h d) -> b h t d', h=self.num_heads)
        v = rearrange(self.proj_v(v), 'b t (h d) -> b h t d', h=self.num_heads)

        attn_score = torch.einsum('bhlk,bhtk->bhlt', [q, k]) * self.scale#实现q*kattention,获得分数
        attn = F.softmax(attn_score, dim=-1)#softmax将分数变成概率值
        attn = self.attn_drop(attn)

        x = torch.einsum('bhlt,bhtv->bhlv', [attn, v])#将attention得到概率值与value信息值相乘得到结果，结构是，b,h,t,d
        x = rearrange(x, 'b h t d -> b t (h d)')#reshape成，b,t,(h,d)

        x = self.proj(x)#扩充维度
        x = self.proj_drop(x)

        return x#b,t,(h,d),输出的是token

5.2.3Block模块

以上两个模块结合残差思想组合在一起构成了Convolutional Projection For Attention 模块的一个block而Convolutional Projection For Attention 模块实际就是由多个block堆叠而成

block模块的结构如下所示：

#3.BLOCK层=attention+MLP,加上残差思想
'''
输入是token，输出是token
输入的token先经过reshape成特征图，送入self-attention操作得到结果token
然后将上面的结果送入MLP中进行运算，输出token
'''
class Block(nn.Module):

    def __init__(self,dim_in,dim_out,num_heads,mlp_ratio=4.,qkv_bias=False,drop=0.,attn_drop=0.,drop_path=0.,act_layer=nn.GELU,norm_layer=nn.LayerNorm,
                 **kwargs):
        super().__init__()

        self.with_cls_token = kwargs['with_cls_token']

        self.norm1 = norm_layer(dim_in)#层归一化1
        self.attn = Attention(dim_in, dim_out, num_heads, qkv_bias, attn_drop, drop,**kwargs)#
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim_out)#层归一化2
        dim_mlp_hidden = int(dim_out * mlp_ratio)
        self.mlp = Mlp(in_features=dim_out,hidden_features=dim_mlp_hidden, act_layer=act_layer,drop=drop)
    def forward(self, x, h, w):#输入是token
        res = x
        x = self.norm1(x)#token
        attn = self.attn(x, h, w)#token,
        x = res + self.drop_path(attn)
        x = x + self.drop_path(self.mlp(self.norm2(x)))

        return x#输出的是token

5.3 VisionTransformer模块（一个完整的阶段statge）

该模块由一个Convolutional Token Embedding 模块+Convolutional Projection 模块 两部分组成

#5.VisionTransformer层,由前面的ConvEmbed+ Convolutional Projection For Attention 模块（堆叠多个block模块）组成
'''
输入是图片，输出是特征图和cls_token
图片数据先经过ConvEmbed，得到一个特征图
然后这个特征图会被reshape成token
这个token会组合上cls_token，一起送入堆叠Block中，输出token
最后会将这个token分离出cls_token和图片数据token，然后将图片数据reshape成图片数据的特征图
'''
class VisionTransformer(nn.Module):
    """ Vision Transformer with support for patch or hybrid CNN input stage
    """
    def __init__(self,
                 patch_size=16,patch_stride=16,patch_padding=0,
                 in_chans=3,embed_dim=768,depth=12,
                 num_heads=12,mlp_ratio=4.,qkv_bias=False,
                 drop_rate=0.,attn_drop_rate=0.,drop_path_rate=0.,
                 act_layer=nn.GELU,norm_layer=nn.LayerNorm,init='trunc_norm',**kwargs):
        super().__init__()
        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
        self.rearrage = None

        # 输入的图片大小：img_size=img_size,
        self.patch_embed = ConvEmbed(patch_size=patch_size,in_chans=in_chans,stride=patch_stride,padding=patch_padding,embed_dim=embed_dim,norm_layer=norm_layer )

        #分类的cls_token
        with_cls_token = kwargs['with_cls_token']
        if with_cls_token:
            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        else:
            self.cls_token = None

        self.pos_drop = nn.Dropout(p=drop_rate)
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule,drop随着深度变化


        #堆叠Block
        blocks = []
        for j in range(depth):
            blocks.append(
                Block(dim_in=embed_dim,dim_out=embed_dim,
                    num_heads=num_heads, mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,drop=drop_rate,
                    attn_drop=attn_drop_rate,drop_path=dpr[j],
                    act_layer=act_layer,norm_layer=norm_layer,
                    **kwargs)
            )

        #将堆叠的block组成神经网络
        self.blocks = nn.ModuleList(blocks)

        #给分类用的token初始化
        if self.cls_token is not None:
            trunc_normal_(self.cls_token, std=.02)


        #初始化网络权重参数
        if init == 'xavier':
            self.apply(self._init_weights_xavier)
        else:
            self.apply(self._init_weights_trunc_normal)

    def _init_weights_trunc_normal(self, m):
        if isinstance(m, nn.Linear):
            logging.info('=> init weight of Linear from trunc norm')
            trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                logging.info('=> init bias of Linear to zeros')
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d)):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def _init_weights_xavier(self, m):
        if isinstance(m, nn.Linear):
            logging.info('=> init weight of Linear from xavier uniform')
            nn.init.xavier_uniform_(m.weight)
            if m.bias is not None:
                logging.info('=> init bias of Linear to zeros')
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm2d)):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)


    def forward(self, x):
        x = self.patch_embed(x)#首先输入的是图片（2d），经过卷积，将图片尺寸变小了，输出的是特征图（2d）
        B, C, H, W = x.size()

        x = rearrange(x, 'b c h w -> b (h w) c')#reshape成token

        cls_tokens = None
        if self.cls_token is not None:
            # stole cls_tokens impl from Phil Wang, thanks
            cls_tokens = self.cls_token.expand(B, -1, -1)#cls_token 扩充维度
            x = torch.cat((cls_tokens, x), dim=1)#cls_token与图片数据的token，拼接在一起

        x = self.pos_drop(x)
        for i, blk in enumerate(self.blocks):#执行堆叠的block
            x = blk(x, H, W)

        if self.cls_token is not None:
            cls_tokens, x = torch.split(x, [1, H*W], 1)#将cls_token与图片数据的token分开
        x = rearrange(x, 'b (h w) c -> b c h w', h=H, w=W)#将图片数据token映射回特征图的形式

        return x, cls_tokens#，x是特征图

如果对Transformer原理不了解的，或者还不了解transformer在视觉上的应用（ViT）

可以去参考一下我的另一篇文章:

transformer在图像分类上的应用ViT以及pytorch代码实现

6.参考文献

1. https://mp.weixin.qq.com/s?__biz=MzI5MDUyMDIxNA==&mid=2247550517&idx=1&sn=178d5dbbd7c90917f183361166ec6121&chksm=ec1cebccdb6b62daaeb3f9208f0d9d41233c40841b20ccd88e2079cf8111ffb368c7034a455d&cur_album_id=1685054606675902466&scene=189#rd

2. 论文：CvT: Introducing Convolutions to Vision Transformers

3. 原作者代码