将卷积引入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 Nonedef forward(self, x):#输入是特征图x = self.proj(x)#（b,c,h,w）B, C, H, W = x.shapex = 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_featureshidden_features = hidden_features or in_featuresself.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):#输入的是tokenx = 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#1self.stride_q = stride_q#1self.dim = dim_outself.num_heads = num_heads# head_dim = self.qkv_dim // num_headsself.scale = dim_out ** -0.5self.with_cls_token = with_cls_token#True#卷积实现得到q,k,vself.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 = Noneelse:raise ValueError('Unknown method ({})'.format(method))return projdef forward_conv(self, x, h, w):#输入是tokenif 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)#加上分类的tokenk = torch.cat((cls_token, k), dim=1)v = torch.cat((cls_token, v), dim=1)return q, k, v#，输出的是tokendef forward(self, x, h, w):#输入x是tokenif (self.conv_proj_q is not Noneor self.conv_proj_k is not Noneor 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,dx = 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)#层归一化1self.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)#层归一化2dim_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):#输入是tokenres = xx = self.norm1(x)#tokenattn = 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 modelsself.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_tokenwith_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 = Noneself.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随着深度变化#堆叠Blockblocks = []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成tokencls_tokens = Noneif self.cls_token is not None:# stole cls_tokens impl from Phil Wang, thankscls_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):#执行堆叠的blockx = 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. 原作者代码