Vision Transformer学习笔记

最新推荐文章于 2024-09-16 13:36:19 发布

CUCKyrie

最新推荐文章于 2024-09-16 13:36:19 发布

阅读量711

点赞数 29

分类专栏： IMC学习笔记文章标签：学习笔记 transformer

本文链接：https://blog.csdn.net/weixin_45671402/article/details/141194040

版权

IMC学习笔记专栏收录该内容

2 篇文章 0 订阅

订阅专栏

论文链接：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，它要求的输出一个大小为 $[ba t c h, l e n g t h, d im]$ 的向量，但是对于一个图像而言，其在深度学习中的基本单位是 $[ba t c h, w i d t h, h e i g h t, c hann e l = 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$ 的图像而言，就可以拆分成 $\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向量。综上，我们就把一个 $[ba t c h, w i d t h, h e i g h t, 3]$ 的输入，转换为了 $[ba t c h, 196, 768]$ 格式的输入，这样就可以送入Transformer中进行编码了。

Posional Embedding

对图像数据进行Embedding后，我们可以得到一个大小为 $[ba t c h, 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.LayerNorm
        act_layer = act_layer or nn.GELU  # nn.GELU
  
        self.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, d im, h, w]$ 的特征，随后会将最后两个维度进行Flatten，得到 $[b, d im, h * w]$ 的特征嵌入，并交换最后两个维度即可得到前文所述的 $[b, h * 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的计算流程就完成了。

CUCKyrie

关注

29
点赞
踩
27

收藏

觉得还不错? 一键收藏
打赏
0
评论
复制链接

分享到 QQ

分享到新浪微博

扫一扫

专栏目录