【精选】狗都能看懂的Vision Transformer的讲解和代码实现 |
您所在的位置：网站首页 › transformer详解代码 › 【精选】狗都能看懂的Vision Transformer的讲解和代码实现 |
【精选】狗都能看懂的Vision Transformer的讲解和代码实现
|
文章目录
1、ViT介绍2、模型详解Embedding层详解Transformer Encoder详解MLP Head详解完整模型（ViT-B/16为例）
3、代码复现4、总结
1、ViT介绍
从深度学习暴发以来，CNN一直是CV领域的主流模型，而且取得了很好的效果，相比之下，基于self-attention结构的Transformer在NLP领域大放异彩。虽然Transformer结构已经成为NLP领域的标准，但在计算机视觉领域的应用还非常有限。 ViT（vision transformer）是Google在2020年提出的直接将Transformer应用在图像分类的模型，通过这篇文章的实验，给出的最佳模型在ImageNet1K上能够达到88.55%的准确率（先在Google自家的JFT数据集上进行了预训练），说明Transformer在CV领域确实是有效的，而且效果还挺惊人。 在讲解ViT原理之前，读者需要有self-attention的相关知识，之前博文有讲过，这里就不复述了。 下图是Vision Transformer的结构，乍一看和self-attention的结构非常像。主要由三个模块组成： Linear Projection (Patch + Position 的Embedding层)Transformer Encoder（详细结构见图右边）MLP Head（分类层）
Embedding层详解
从图中看，Transformer Encoder读入的是一小块一小块这样的图片。这样做的原因是：将一个个小块图像视为token（token可以理解为NLP中的一个字或词），在Transformer中计算每个token之间的相关性。这一点就和卷积神经网络有很大区别了。以往的CNN，以卷积 + 池化的方式不断下采样，这样理论上模型可以通过加深模型深度，达到增大感受野的目的。不过这样会有两个缺点： 实际结果中显示，CNN对边缘的响应很弱。这也非常好理解，越靠边缘的像素，因为被卷积次数少，自然在梯度更新时，贡献更少。CNN只能和临近像素计算相关性。由于其滑窗卷积的特性，无法对非领域的像素共同计算，例如左上角的像素无法和右下角的像素联合卷积。这就导致了某些空间信息是无法利用的。同时根据MAE论文中所说的，自然图像具有冗余性，即相邻像素点代表的信息是差不多的，所以只计算领域像素无法最大化利用图像特征。回到ViT中，仅仅把图像拆分成小块（patch）是不够的，Transformer Encoder需要的是一个向量，shape为[num_token, token_dim]。对于图片数据来说，shape为[H,W,C]是不符合要求的，所以就需要转换，要将图片数据通过这个Embedding层转换成token。以ViT-B/16为例： 假设输入图像为224x224x3，一个token原始图像shape为16x16x3，那这样就可以将图像拆分成 ( 224 / 16 ) 2 = 196 (224/16)^2 = 196 (224/16)2=196个patch，然后将每个patch线性映射至一维向量中，那么这个一维向量的长度即为 16 ∗ 16 ∗ 3 = 768 16*16*3=768 16∗16∗3=768维。将196个token叠加在一起最后维度就是[196, 768]。 在代码实现中，patch的裁剪是用一个patch_size大小的卷积同时以patch_size的步长进行卷积实现的： class PatchEmbed(nn.Module): """ 2D Image to Patch Embedding """ def __init__(self, image_size=224, patch_size=16, in_c=3, embed_dim=768, norm_layer=None): """ Map input tensor to patch. Args: image_size: input image size patch_size: patch size in_c: number of input channels embed_dim: embedding dimension. dimension = patch_size * patch_size * in_c norm_layer: The function of normalization """ super().__init__() image_size = (image_size, image_size) patch_size = (patch_size, patch_size) self.image_size = image_size self.patch_size = patch_size self.grid_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1]) self.num_patches = self.grid_size[0] * self.grid_size[1] # The input tensor is divided into patches using 16x16 convolution 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.image_size[0] and W == self.image_size[1], \ f"Input image size ({H}*{W}) doesn't match model ({self.image_size[0]}*{self.image_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同时需要注意的是，token可以添加位置编码。由于self-attention的机制，在没有字符位置的情况下，“我爱你”和“你爱我”的计算结果是一样的。所以当添加了位置编码，分类准确率将会更高。所以迁移到ViT的Embeding层中，位置信息是以Add的方式融合在输入tensor中的。 另外，输入tensor还需要有一个class token（分类层），数据格式和其他token一样，长度为768的向量，与位置编码的融合方式不一样，这里做的是Concat，这样做是因为分类信息是在后面需要取出来单独做预测的，所以不能以Add方式融合，shape也就从[196, 768]变为[197, 768]。 Transformer Encoder也就是堆叠encoder block几次，和常规的CNN类似，主要由几个部分组成： Layer Normalization 相比于Batch Normalization来说，BN是针对所有的样本，对某一个特征图计算均值和方差，然后然后对这个特征图神经元做归一化。LN是对某一个样本，计算该样本所有特征图的均值和方差，然后对这个样本做归一化。BN适用于不同mini batch数据分布差异不大的情况，而且BN需要开辟变量存每个节点的均值和方差，空间消耗略大；而且 BN适用于有mini_batch的场景。LN只需要一个样本就可以做归一化，可以避免 BN 中受 mini-batch 数据分布影响的问题，也不需要开辟空间存每个节点的均值和方差。Multi-Head Attention 这个结构是self-attention中的一种，在之前的博客中有详细的讲解。DropOut/DropPath 论文中描述是使用了DropOut，但实际别人复现的代码中使用的是DropPath。两者对最后的结果影响不大，这里也不展开详细说了。MLP Block 如下图右侧所示，即一个全连接层 + GELU激活函数 + DropOut组成，采用的是倒瓶颈结构，输入特征层经过一次全连接之后，通道膨胀为原来的4倍，后一个全连接层再恢复成原来的数目。所以经过MLP Block之后tensor的shape是不变的。
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 MLP Head详解经过Transformer Encoder后输出tensor的shape和输出tensor的shape是一样的，以ViT-B/16为例，输入是[197, 768]，输出还是[197, 768]。下游任务如果是分类模型，需要提取出对应的class token获取分类结果。Vision Transformer中说MLP是由一个全连层 + Tanh激活 + 全连接层组成。但实际使用起来一层全连接层直接做分类即可。
具体代码如下： class VisionTransformer(nn.Module): def __init__(self, image_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.5, embed_layer=PatchEmbed, norm_layer=None, act_layer=None): """ Args: image_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, dim = patch_size * patch_size * in_c 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) # partial类似闭包函数，第一个参数是函数对象，后面的参数是第一个函数对象的实参 act_layer = act_layer or nn.GELU self.patch_embed = embed_layer(image_size=image_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)) # nn.Parameter转换函数，把参数转换为可训练变量 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() # placeholder # 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) 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 3、代码复现code(觉得有用请点star~，这对我很重要):https://github.com/Runist/torch_Vision_Transformer paper:https://arxiv.org/pdf/2010.11929.pdf 4、总结至此，Vision Transformer的原理及模型结构以及全部介绍完毕。从实验给出的结果，在当时也达到了SOTA，但是相比于CNN，它需要更多的数据集。在小数据集上训练出来的精度是不如CNN的，但在大数据集上ViT精度更高。一个直观的解释是：ViT因为self-attention独特的机制，更多的利用token与token跨像素之间的信息，而CNN只是对领域的像素进行计算，所以相同参数的情况下，ViT获得的信息更多，在某种程度上可以看成是模型深度更深。所以小数据集上ViT是欠拟合的。实际开发中的做法是：基于大数据集上训练，得到一个预训练权重，然后再在小数据集上Fine-Tune。 后续将会补充Vision Transformer实现的完整代码，先挖个坑。 已补充 |
【本文地址】
今日新闻 |
推荐新闻 |