论文链接：https://arxiv.org/abs/2010.11929
项目链接：https://github.com/google-research/vision_transformer
本文代码链接：https://gitcode.com/gh_mirrors/de/deep-learning-for-image-processing/tree/master/pytorch_classification/vision_transformer

ViT架构简述

Vision Transformer，简称ViT，是Transformer在视觉任务中的第一个落地模型。自此，NLP和CV开始大一统，并进一步推动了多模态模型的发展。ViT的很多设计思路参考了NLP中的BERT模型，其本质是希望将Transformer架构不经任何修改的前提下，通过巧妙的图像预处理方法，将Transformer无缝衔接到任何一个CV任务中。
ViT的整体架构图如下所示，其基本的思想是将一幅$224\times 224$的图像切分为多个$16\times 16$的Patches，并将这些Patch作为Transformer中的每个Token直接应用到Transformer中。由于在之前的文章中已经详细解读过Transformer的架构，ViT中的Transformer没有任何变动，因此重点讲述ViT对图像专门设计的方法。

由上图可以看出，整个ViT架构包括三个主要的组成部分：Linear Projection（也就是Image Embedding）、Transformer Encoder和MLP Head，其中Transformer Encoder的部分和一般的Transformer并无差异，主要的改进工作集中于图像Embedding层。

Image Embedding

对于一个标准的Transformer，它要求的输出一个大小为$[batch,length,dim]$的向量，但是对于一个图像而言，其在深度学习中的基本单位是$[batch,width,height,channel=3]$，这明显不是Transformer期望的输入格式。因此，需要一个Embedding层将图像数据嵌入到一个理想的特征空间内。一种直观的思路是，以像素点为单位，将每个2D图像flatten成一个1D的像素序列，这样就可以输入Transformer中。但是，哪怕是对于一幅$224\times 224$大小的图像，Flatten后就会有$224\times 224$个像素点参与注意力计算，随着输入图片分辨率的增加，模型计算压力会呈指数级增长，显然这不是最优的解法。
正如作者的论文题目”AN IMAGE IS WORTH 16X16 WORDS“，每个图像相当于一系列$16\times 16$大小Patch的组合，此时对于一幅$224\times 224$的图像而言，就可以拆分成$num=\frac{224^2}{16^2}=196$个Patch。换言之，一张$224\times 224$的图，就是由196个Words组成的！
那如何规定输入的Word/Patch的维度呢？这时候就可以采用之前提到的最直接的思路：每个Word/Patch大小是$16\times 16\times 3$，3为通道数，因此可以把每个Word/Patch的2D数据进行Flatten，成为一个$16\times 16\times 3=768$维的1D向量。综上，我们就把一个$[batch,width,height,3]$的输入，转换为了$[batch,196,768]$格式的输入，这样就可以送入Transformer中进行编码了。

Posional Embedding

对图像数据进行Embedding后，我们可以得到一个大小为$[batch,196,768]$的Image Embedding，根据之前对Transformer的介绍，为了标注每个词在句子中的先后顺序，Transformer采用正余弦位置编码给每个词向量加入了一个Positional Embedding。对于图像而言，我们将图像切分为Patch后，同样也需要一个位置信息，来标注Patch的顺序，因为注意力计算本身是忽略Token之间位置信息的。原文中，作者采用了一个可学习的1D位置编码，并且通过实验验证，2D位置编码并没有展现出更出色的性能。

关于CLS

这里作者参考了BERT的Class Token的设计，在原有的Patch序列中加入了一个可学习的Token，即CLS。认为这个Token在不断的执行注意力计算的过程中，是可以学习到整个图像的全局信息的，并且只需要将这个Token送入后续的下游任务，即可取得十分出色的效果。这个操作类似对196个Patch的向量进行平均，从实现效果的角度看没有太大差别，但是对学习率有不同的要求：

代码调试记录

本文采用ViT的torch复现代码，下面将结合代码对照论文进行讲解。
代码中提供了多种ViT版本，我们这里只选择调试最基础的Vit_base_patch_16_224。

def vit_base_patch16_224(num_classes: int = 1000):  """  ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).    ImageNet-1k weights @ 224x224, source https://github.com/google-research/vision_transformer.    weights ported from official Google JAX impl:    链接: https://pan.baidu.com/s/1zqb08naP0RPqqfSXfkB2EA  密码: eu9f  """    model = VisionTransformer(img_size=224,  patch_size=16,  embed_dim=768,  depth=12,  num_heads=12,  representation_size=None,  num_classes=num_classes)  return model  def vit_base_patch16_224_in21k(num_classes: int = 21843, has_logits: bool = True):  """  ViT-Base model (ViT-B/16) from original paper (https://arxiv.org/abs/2010.11929).    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.    weights ported from official Google JAX impl:    https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch16_224_in21k-e5005f0a.pth    """    model = VisionTransformer(img_size=224,  patch_size=16,  embed_dim=768,  depth=12,  num_heads=12,  representation_size=768 if has_logits else None,  num_classes=num_classes)  return model  def vit_base_patch32_224(num_classes: int = 1000):  """  ViT-Base model (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).    ImageNet-1k weights @ 224x224, source https://github.com/google-research/vision_transformer.    weights ported from official Google JAX impl:    链接: https://pan.baidu.com/s/1hCv0U8pQomwAtHBYc4hmZg  密码: s5hl  """    model = VisionTransformer(img_size=224,  patch_size=32,  embed_dim=768,  depth=12,  num_heads=12,  representation_size=None,  num_classes=num_classes)  return model  def vit_base_patch32_224_in21k(num_classes: int = 21843, has_logits: bool = True):  """  ViT-Base model (ViT-B/32) from original paper (https://arxiv.org/abs/2010.11929).    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.    weights ported from official Google JAX impl:    https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_base_patch32_224_in21k-8db57226.pth    """    model = VisionTransformer(img_size=224,  patch_size=32,  embed_dim=768,  depth=12,  num_heads=12,  representation_size=768 if has_logits else None,  num_classes=num_classes)  return model  def vit_large_patch16_224(num_classes: int = 1000):  """  ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).    ImageNet-1k weights @ 224x224, source https://github.com/google-research/vision_transformer.    weights ported from official Google JAX impl:    链接: https://pan.baidu.com/s/1cxBgZJJ6qUWPSBNcE4TdRQ  密码: qqt8  """    model = VisionTransformer(img_size=224,  patch_size=16,  embed_dim=1024,  depth=24,  num_heads=16,  representation_size=None,  num_classes=num_classes)  return model  def vit_large_patch16_224_in21k(num_classes: int = 21843, has_logits: bool = True):  """  ViT-Large model (ViT-L/16) from original paper (https://arxiv.org/abs/2010.11929).    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.    weights ported from official Google JAX impl:    https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch16_224_in21k-606da67d.pth    """    model = VisionTransformer(img_size=224,  patch_size=16,  embed_dim=1024,  depth=24,  num_heads=16,  representation_size=1024 if has_logits else None,  num_classes=num_classes)  return model  def vit_large_patch32_224_in21k(num_classes: int = 21843, has_logits: bool = True):  """  ViT-Large model (ViT-L/32) from original paper (https://arxiv.org/abs/2010.11929).    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.    weights ported from official Google JAX impl:    https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-vitjx/jx_vit_large_patch32_224_in21k-9046d2e7.pth    """    model = VisionTransformer(img_size=224,  patch_size=32,  embed_dim=1024,  depth=24,  num_heads=16,  representation_size=1024 if has_logits else None,  num_classes=num_classes)  return model  def vit_huge_patch14_224_in21k(num_classes: int = 21843, has_logits: bool = True):  """  ViT-Huge model (ViT-H/14) from original paper (https://arxiv.org/abs/2010.11929).    ImageNet-21k weights @ 224x224, source https://github.com/google-research/vision_transformer.    NOTE: converted weights not currently available, too large for github release hosting.    """model = VisionTransformer(img_size=224,  patch_size=14,  embed_dim=1280,  depth=32,  num_heads=16,  representation_size=1280 if has_logits else None,  num_classes=num_classes)  return model

由于只是调试代码，我们并不调用正经的train.py，而是用一个简单的调用脚本来实现：

if __name__ == "__main__":  # 构建ViT模型  model = vit_base_patch16_224()  # 构建输入数据[b, w, h, c]  image = torch.randn((2, 3, 224, 224))  # 输入模型  outputs = model(image)  print(outputs.shape)

首先需要加载一个vit_base模型，进入ViT的初始化部分：

class VisionTransformer(nn.Module):  def __init__(self, img_size=224, patch_size=16, in_c=3, num_classes=1000,  embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, qkv_bias=True,  qk_scale=None, representation_size=None, distilled=False, drop_ratio=0.,  attn_drop_ratio=0., drop_path_ratio=0., embed_layer=PatchEmbed, norm_layer=None,  act_layer=None):  """  Args:  img_size (int, tuple): input image size    patch_size (int, tuple): patch size    in_c (int): number of input channels   num_classes (int): number of classes for classification head    embed_dim (int): embedding dimension    depth (int): depth of transformer    num_heads (int): number of attention heads    mlp_ratio (int): ratio of mlp hidden dim to embedding dim    qkv_bias (bool): enable bias for qkv if True    qk_scale (float): override default qk scale of head_dim ** -0.5 if set    representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set    distilled (bool): model includes a distillation token and head as in DeiT models    drop_ratio (float): dropout rate    attn_drop_ratio (float): attention dropout rate    drop_path_ratio (float): stochastic depth rate    embed_layer (nn.Module): patch embedding layer    norm_layer: (nn.Module): normalization layer"""super(VisionTransformer, self).__init__()  self.num_classes = num_classes  self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models  self.num_tokens = 2 if distilled else 1  norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) # nn.LayerNormact_layer = act_layer or nn.GELU  # nn.GELUself.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_c=in_c, embed_dim=embed_dim)  num_patches = self.patch_embed.num_patches  self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))  self.dist_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) if distilled else None  self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))  self.pos_drop = nn.Dropout(p=drop_ratio)  dpr = [x.item() for x in torch.linspace(0, drop_path_ratio, depth)]  # stochastic depth decay rule  self.blocks = nn.Sequential(*[  Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,  drop_ratio=drop_ratio, attn_drop_ratio=attn_drop_ratio, drop_path_ratio=dpr[i],  norm_layer=norm_layer, act_layer=act_layer)  for i in range(depth)  ])  self.norm = norm_layer(embed_dim)  # Representation layer  if representation_size and not distilled:  self.has_logits = True  self.num_features = representation_size  self.pre_logits = nn.Sequential(OrderedDict([  ("fc", nn.Linear(embed_dim, representation_size)),  ("act", nn.Tanh())  ]))  else:  self.has_logits = False  self.pre_logits = nn.Identity()  # Classifier head(s)  self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()  self.head_dist = None  if distilled:  self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if num_classes > 0 else nn.Identity()  # Weight init  nn.init.trunc_normal_(self.pos_embed, std=0.02)  if self.dist_token is not None:  nn.init.trunc_normal_(self.dist_token, std=0.02)  nn.init.trunc_normal_(self.cls_token, std=0.02)  self.apply(_init_vit_weights)

可以看到这是一个12头的ViT，一共有12层，num_classes代表最后的分类任务有几个类别，同时也是最后一层MLP的输出维度。
对于Image Embedding部分，代码中用一个独立的class实现：

class PatchEmbed(nn.Module):  """  2D Image to Patch Embedding    """    def __init__(self, img_size=224, patch_size=16, in_c=3, embed_dim=768, norm_layer=None):  super().__init__()  img_size = (img_size, img_size)  patch_size = (patch_size, patch_size)  self.img_size = img_size  self.patch_size = patch_size  self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])  self.num_patches = self.grid_size[0] * self.grid_size[1]  self.proj = nn.Conv2d(in_c, embed_dim, kernel_size=patch_size, stride=patch_size)  self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()  def forward(self, x):  B, C, H, W = x.shape  assert H == self.img_size[0] and W == self.img_size[1], \  f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."  # flatten: [B, C, H, W] -> [B, C, HW]  # transpose: [B, C, HW] -> [B, HW, C]        x = self.proj(x).flatten(2).transpose(1, 2)  x = self.norm(x)  return x

并且在Transformer类中用以下代码实现调用：

self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_c=in_c, embed_dim=embed_dim)

可以看到，Patch Embedding中主要包括一个Conv2D和Norm层，在第24行代码中，输入的x[b, c, h, w]会通过一个Conv2D转换为一个768维度的向量，得到一个$[b,dim,h,w]$的特征，随后会将最后两个维度进行Flatten，得到$[b,dim,h\ast w]$的特征嵌入，并交换最后两个维度即可得到前文所述的$[b,h\ast w,768]$的特征。PatchEmbedding中的Norm Layer是一个Identity层，即直接输出不做任何操作。
第37行代码是对CLS Token的设置：

self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))

通过nn.Parameter来定义一个可学习的Token，这里是首先定义一个单位token，在后续会拓展到和输入特征一致的大小。
第39行代码是Positional Embedding的实现：

self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim)) 

同样也是定义了一个可学习的位置编码，而不是transformer中的正余弦编码，大小为$[1,196+1,786]$，初始化为0。
随后ViT定义了12层的Transformer Encoder：

具体的代码如下：

class Block(nn.Module):  def __init__(self,  dim,  num_heads,  mlp_ratio=4.,  qkv_bias=False,  qk_scale=None,  drop_ratio=0.,  attn_drop_ratio=0.,  drop_path_ratio=0.,  act_layer=nn.GELU,  norm_layer=nn.LayerNorm):  super(Block, self).__init__()  self.norm1 = norm_layer(dim)  self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,  attn_drop_ratio=attn_drop_ratio, proj_drop_ratio=drop_ratio)  # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here  self.drop_path = DropPath(drop_path_ratio) if drop_path_ratio > 0. else nn.Identity()  self.norm2 = norm_layer(dim)  mlp_hidden_dim = int(dim * mlp_ratio)  self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop_ratio)  def forward(self, x):  x = x + self.drop_path(self.attn(self.norm1(x)))  x = x + self.drop_path(self.mlp(self.norm2(x)))  return x

可以看到，整个Block中包括两层Layer_Norm、一个多头自注意力和一个MLP，和论文中的示例图可以对应。
我们再回到Transformer类中，在完成了Transformer Encoder的定义后，最后需要定义一个分类头，即一个简单的MLP即完成了对Transformer的定义过程。整个Transformer的结构图下：

VisionTransformer((patch_embed): PatchEmbed((proj): Conv2d(3, 768, kernel_size=(16, 16), stride=(16, 16))(norm): Identity())(pos_drop): Dropout(p=0.0, inplace=False)(blocks): Sequential((0-11): Block((norm1): LayerNorm((768,), eps=1e-06, elementwise_affine=True)(attn): Attention((qkv): Linear(in_features=768, out_features=2304, bias=True)(attn_drop): Dropout(p=0.0, inplace=False)(proj): Linear(in_features=768, out_features=768, bias=True)(proj_drop): Dropout(p=0.0, inplace=False))(drop_path): Identity()(norm2): LayerNorm((768,), eps=1e-06, elementwise_affine=True)(mlp): Mlp((fc1): Linear(in_features=768, out_features=3072, bias=True)(act): GELU(approximate='none')(fc2): Linear(in_features=3072, out_features=768, bias=True)(drop): Dropout(p=0.0, inplace=False))))(norm): LayerNorm((768,), eps=1e-06, elementwise_affine=True)(pre_logits): Identity()(head): Linear(in_features=768, out_features=1000, bias=True)
)

下面我们通过输入一个Image信息来调试ViT的运行流程，首先看一下Transformer的forward函数：

def forward_features(self, x):  # [B, C, H, W] -> [B, num_patches, embed_dim]  x = self.patch_embed(x)  # [B, 196, 768]  # [1, 1, 768] -> [B, 1, 768]    cls_token = self.cls_token.expand(x.shape[0], -1, -1)  if self.dist_token is None:  x = torch.cat((cls_token, x), dim=1)  # [B, 197, 768]  else:  x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)  x = self.pos_drop(x + self.pos_embed)  x = self.blocks(x)  x = self.norm(x)  if self.dist_token is None:  return self.pre_logits(x[:, 0])  else:  return x[:, 0], x[:, 1]  def forward(self, x):  x = self.forward_features(x)  if self.head_dist is not None:  x, x_dist = self.head(x[0]), self.head_dist(x[1])  if self.training and not torch.jit.is_scripting():  # during inference, return the average of both classifier predictions  return x, x_dist  else:  return (x + x_dist) / 2  else:  x = self.head(x)  return x

forward接受一个输入图像$[2,3,224,224]$，随后首先通过一个self.forward_features(x)函数对原始图像信息进行编码，即使调用PatchEmbed进行特征嵌入，会返回一个大小为$[2,196,768]$的Patch特征。随后对cls_token在Batch维度上进行扩展，即对每个Batch上的数据都添加一个CLS Token。

cls_token = self.cls_token.expand(x.shape[0], -1, -1)

随后通过以下代码，将CLS Token和Patch Token拼接起来：

if self.dist_token is None:  x = torch.cat((cls_token, x), dim=1)  # [B, 197, 768]  
else:  x = torch.cat((cls_token, self.dist_token.expand(x.shape[0], -1, -1), x), dim=1)

将Positional Embedding和token Embedding进行加和：

x = self.pos_drop(x + self.pos_embed)

这样得到的x特征即可送入Attention模块中进行编码，Attention Block的计算流程如下：

def forward(self, x):  x = x + self.drop_path(self.attn(self.norm1(x)))  x = x + self.drop_path(self.mlp(self.norm2(x)))  return x

和下图中的计算流程是一致的。

我们重点看一下Attention的计算流程：

class Attention(nn.Module):  def __init__(self,  dim,   # 输入token的dim  num_heads=8,  qkv_bias=False,  qk_scale=None,  attn_drop_ratio=0.,  proj_drop_ratio=0.):  super(Attention, self).__init__()  self.num_heads = num_heads  head_dim = dim // num_heads  self.scale = qk_scale or head_dim ** -0.5  self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)  self.attn_drop = nn.Dropout(attn_drop_ratio)  self.proj = nn.Linear(dim, dim)  self.proj_drop = nn.Dropout(proj_drop_ratio)  def forward(self, x):  # [batch_size, num_patches + 1, total_embed_dim]  B, N, C = x.shape  # qkv(): -> [batch_size, num_patches + 1, 3 * total_embed_dim]  # reshape: -> [batch_size, num_patches + 1, 3, num_heads, embed_dim_per_head]        # permute: -> [3, batch_size, num_heads, num_patches + 1, embed_dim_per_head]        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)  # [batch_size, num_heads, num_patches + 1, embed_dim_per_head]  q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)  # transpose: -> [batch_size, num_heads, embed_dim_per_head, num_patches + 1]        # @: multiply -> [batch_size, num_heads, num_patches + 1, num_patches + 1]        attn = (q @ k.transpose(-2, -1)) * self.scale  attn = attn.softmax(dim=-1)  attn = self.attn_drop(attn)  # @: multiply -> [batch_size, num_heads, num_patches + 1, embed_dim_per_head]  # transpose: -> [batch_size, num_patches + 1, num_heads, embed_dim_per_head]        # reshape: -> [batch_size, num_patches + 1, total_embed_dim]        x = (attn @ v).transpose(1, 2).reshape(B, N, C)  x = self.proj(x)  x = self.proj_drop(x)  return x

输入的特征通过一个线性层将输入的$[2,197,768]$转换成$[2,12,197,64]$大小的特征，即把768维度的特征分配到12个头上，每个头的维度就是64维。这里和Transformer中的区别在于，QKV直接计算注意力，没有经过pad_attn_mask。
从多头注意力层返回的特征并不是完整的输出，而是只返回CLS Token：

if self.dist_token is None:  return self.pre_logits(x[:, 0])  
else:  return x[:, 0], x[:, 1]

返回的这个CLS Token可以作为后续分类任务的特征。至此，整个ViT的计算流程就完成了。