#Vision Transformer的Pytorch实现 ## 1 引言 下面将介绍Vision Transformer的实现过程，下面均简称（Vit）。首先来回顾Vit的整体架构，该模型基本沿用了整个Transformer的Encoder部分，将图片分割成16*16大小的patch，然后将每个patch进行embedding操作，将embedding得到的序列输入到transformer中，从而实现了在不改变transformer结构的基本结构上，实现了CV领域的应用。 ## 2 模型结构  整体的流程可以分为3个部分, 下面依次介绍这三部分： 1. embedding ——将图片进行embedding操作获取sequence 2. transformer encoder —— 利用encoder结构中的attention机制，构建token之间的关联关系 3. mlp head ——将最终encoder的输出映射结果 ### 1 Embedding 第一部分的embedding是Vit的核心，该部分将图片转化为embedding向量，从而使得可以像处理文本序列一样处理图片。这也使得对于图像的特征提取完全依赖于transformer结构，抛弃了以往图像分类、检测、分割等任务对CNN结构的依赖。 该部分embedding由3部分组成分别是： 1. Patch Embedding ——负责将图片转化成embedding序列 2. Position Embedding ——负责提供图片切片的位置信息 3. class token ——负责加入分类头，汇聚所有patch的特征信息  #### 1.1 Patch Embedding **如何将图片转化成embedding序列输入模型？** 在处理输入得到图片时，需要将图片分割成一个个的**patch**，该操作大大降低了transformer的计算量。transformer的计算量与sequence的长度呈2次的平方关系，从而使得将transformer应用于cv领域成为可能。 在标准的vit-base模型中，输入图片尺寸$$224 \times 224$$, 通过将其分成$$16 \times 16$$的patch，得到$$14 \times 14$$大小的feature map，再将其展平flatten操作，得到长度为$$14 \times 14 = 196$$的序列。每个序列可以看成由$$3 \times 16 \times 16$$的patch展平所得到，因此每个序列的维度是$$768$$。 如果加入batch_size的维度，通过Embeddding操作将图片$$(1, 3, 224, 224) -> (1, 196, 768)$$ 通过代码实现以上步骤如下， 在forward中详细标注了每一步骤的维度： ```python class PatchEmbedding(nn.Module): def __init__(self, img_size: int, patch_size: int = 16, in_c: int = 3, embed_dim=None ): super(PatchEmbedding, self).__init__() assert img_size % patch_size == 0, "img_size should be divided by patch_size" self.embed_dim = embed_dim if embed_dim else in_c * patch_size ** 2 # 获得与原有数据一致的emb patch_size = (patch_size, patch_size) self.conv = nn.Conv2d(in_channels=in_c, out_channels=self.embed_dim, kernel_size=patch_size, stride=patch_size) self.norm = nn.LayerNorm(self.embed_dim) def forward(self, x): # [batch_size, 3, h, w] B, C, H, W = x.size() # [batch_size, 3, H, W] -> [batch_size, embed_dim, grid_size, grid_size] x = self.conv(x) # [batch_size, embed_dim, grid_size, grid_size] -> [batch_size, embed_dim, grid_size*grid_size] x = torch.reshape(x, (B, self.embed_dim, -1)) # [batch_size, embed_dim, grid_size*grid_size] -> [batch_size, grid_size * grid_size, embed_dim] x = torch.transpose(x, 1, 2) x = self.norm(x) return x ``` 当设置embed_dim=None的时候，根据图片的patch_size和in_channel自动计算embedding的维度大小，这样可以使得图像的信息与变换前保持一致。 $$embedding = inchannel \times patchsize^2$$ 可以测试运行一下，输入一张$$(3, 224, 224)$$的图片，得到PatchEmbedding输出结果$$(1, 196, 768)$$ ```python if __name__ == "__main__": img = torch.randn(1, 3, 224, 224) patch_emb = PatchEmbedding(img_size=32, patch_size=16, embed_dim=None) img_emb = patch_emb(img) print(img_emb.size()) ``` 输出结果与我们的预期相同： ``` torch.Size([1, 196, 768]) ``` #### 1.2 Position Embedding **为什么要加入位置embedding向量？** 由于transformer结构对于位置没有办法衡量，因此位置编码在transformer中起到重要作用。李宏毅老师在attention机制作用时sequence序列的任意两个元素就是“海内存知己，天涯若比邻”。通过Q,V矩阵计算所得任意2个向量之间的关联性与位置的远近毫无关系。这显然不能很好的建模文本任务这种序列型特征数据，和图像数据这种局部性特征数据。为了解决这个问题，需要加入pos_emb进去。 这样我们就了解了位置编码的作用，文中采用了1d的可训练参数作为位置编码，与《Attention is all you need》这篇transformer开山之作中的位置编码有所不同。 ```python # pos_embedding num_grids = (img_size // patch_size) ** 2 self.pos_emb = nn.Parameter(torch.zeros(1, num_grids+1, dim)) self.pos_drop = nn.Dropout(p=drop_ratio) ``` #### 1.3 class token **如何使用transformer encoder结构进行分类？** 我们的目的在于对图像进行分类，我们知道transformer结构是一个seq2seq的结构。因此用作分类我们在输入的时候加入一个class token———该token与经过patch embedding的向量维度相同。在刚才经过patch embedding之后得到196个维度为768的向量。再加入1个长度为768的class token向量，得到197个向量。再将其输入到encoder结构中。经过多层transformer结构之后，再将class token进行project映射到类别数量上，就可以得到类别。 **为什么这么做可以进行分类？** 在self-attention操作中，我们将任意2个向量的qv矩阵计算得到注意力矩阵。因此class token的向量与图片所有的patch embedding计算其关联程度，然后乘以v矩阵得到结果。因此，class token矩阵集合了整张图片的信息，其输出可以作为整张图片的特征表示。 初始化cls_token如下所示，其中dim=0的维度需要在forward时根据batch维度进行expand ```python # class token self.cls_token = nn.Parameter(torch.zeros(1, 1, dim)) ``` #### 1.4 整合输入部分 **模型定义部分** ```python # patch_embdding self.patch_emb = PatchEmbedding(img_size=img_size, patch_size=patch_size, in_c=in_c) # class token self.cls_token = nn.Parameter(torch.zeros(1, 1, dim)) # pos_embedding num_grids = (img_size // patch_size) ** 2 self.pos_emb = nn.Parameter(torch.zeros(1, num_grids+1, dim)) self.pos_drop = nn.Dropout(p=drop_ratio) ``` **前向传播部分** ```python # [B, C, H, W] -> [B, num_patches, embed_dim] x = self.patch_emb(x) # [1, 1, embed_dim] -> [B, 1, embed_dim] cls_token = self.cls_token.expand(B, -1, -1) # [B, num_patches, embed_dim] + [B, 1, embed_dim] = [B, num_patches + 1, embed_dim] x = torch.cat([cls_token, x], dim=1) x = self.pos_drop(self.pos_emb + x) ``` ### 2 Transformer Encoder 该部分就是标准的transformer中的encoder部分，将我们的输入引入该部分得到输出即可。transformer结构的输入和输出一般不涉及维度变化，因此便于快速构建不同深度的组件。  由上图可以看出，Encoder由N个Encoder Block组成。每个Block先采用Multi-head Attention模块，然后采用MLP模块，前后接Layer Norm和Drop Path。其中MLP的结构如上图右侧所示。 该部分Encoder由2部分组成分别是： 1. Multi-head Attention ——组成Transformer结构的核心机制 2. MLP ——组成Transformer结构的前馈网络 下面将着重介绍Multi-head Attention的结构。 #### 2.1 Multi-head Attention原理及实现 ##### 2.1.1 Attention原理 作为transformer结构中核心的attention机制，是整个模型中的核心。首先还是来回顾一下整个self-attention机制，依此来构建Attention模块，该部分引用李宏毅老师的ppt作为参考。 **Step-1**  首先将输入a乘以Wq，Wk得到q矩阵和k矩阵，将q1和k1-k4进行矩阵乘法，得到a1和其它输入的关联矩阵，通过softmax层得到概率。  其中通过矩阵运算如图所示. **Step-2**  将每个输入a乘以Wv矩阵，得到v矩阵。将关联矩阵a和v矩阵逐个相乘求和，得到a1的输出b1.  其中通过矩阵运算如图所示. ##### 2.1.2 Multi-head Self-Attention 多头注意力的关键点在于分头，和每个头之间的合并操作。除此之外，多头注意力和单头注意力并不存在区别。整体的计算量需要添加一个(dim, dim)维度的矩阵，用于将多个head的输出进行合并。 多头注意力的作用在于采用多个独立的qkv矩阵进行attention操作，在一定程度上丰富所提取到的特征。 **Step-1**  如图所示，假设我们设定num_heas=2,我们将计算得到的qkv三个矩阵，分成q1,q2, k1, k2, v1, v2三个矩阵。角标1的一起计算，角标2一起计算，得到输出b1，b2。这样我们对应一个输入就有2个输出。 **Step-2**  最后将得到的b1，b2乘以一个Wo矩阵，进行映射，得到最终的输出结果。 ##### 2.1.3 Attention代码实现 ```python def __init__(self, dim: int, # 输入token的dim维度 num_heads: int = 8, # 分类头数量 qkv_bias: bool = False, # qkv的bias qk_scale: float = None): super(Attention, self).__init__() assert dim % num_heads == 0, "dim: {} can not be divided by num_heads:{}".format(dim, num_heads) self.dim = dim self.num_heads = num_heads self.dim_per_head = dim // num_heads self.qk_scale = qk_scale or math.sqrt(self.dim_per_head) self.qkv = nn.Linear(in_features=dim, out_features=dim * 3, bias=qkv_bias) self.proj = nn.Linear(in_features=dim, out_features=dim) def forward(self, x): # batch_size, num_patches, dim B, N, D = x.size() # [batch_size, num_patches, dim] -> [batch_size, num_patches, 3*dim] qkv = self.qkv(x) # [batch_size, num_patches, 3*dim] -> [batch_size, num_patches, 3, num_head, dim / num_head] qkv = qkv.reshape(B, N, 3, self.num_heads, self.dim_per_head) # [batch_size, num_patches, 3, num_head, dim / num_head]->[3, batch_size, num_head, num_patches, dim / num_head] qkv = qkv.permute(2, 0, 3, 1, 4) # [batch_size, num_head, num_patches, dim / num_head] q, k, v = qkv[0], qkv[1], qkv[2] # [batch_size, num_head, num_patches, dim / num_head] @ [batch_size, num_head, dim / num_head, num_patches] attn = torch.matmul(q, k.transpose(-1, -2)) * self.qk_scale # [batch_size, num_head, num_patches, num_patches] attn = torch.softmax(attn, dim=-1) # [batch_size, num_head, num_patches, num_patches] @ [batch_size, num_head, num_patches, dim / num_head] val = torch.matmul(attn, v) # [batch_size, num_head, num_patches, dim / num_head] -> [batch_size, num_patches, num_head, dim / num_head] val = torch.transpose(val, 1, 2) # [batch_size, num_patches, num_head, dim / num_head] -> [batch_size, num_patches, dim] val = torch.reshape(val, (B, N, D)) # [batch_size, num_patches, dim] -> [batch_size, num_patches, dim] val = self.proj(val) return val ``` #### 2.2 Encoder结构整合 ##### 2.2.1 MLP的代码实现 ```python class Mlp(nn.Module): def __init__(self, in_channels: int = 768, expand_ratio: int = 4, drop_ratio: float = 0.): super(Mlp, self).__init__() in_channels = in_channels hidden_channels = in_channels * expand_ratio out_channels = in_channels self.fc1 = nn.Linear(in_features=in_channels, out_features=hidden_channels) self.act = nn.GELU() self.fc2 = nn.Linear(in_features=hidden_channels, out_features=out_channels) self.dropout = nn.Dropout(p=drop_ratio) def forward(self, x): x = self.fc1(x) x = self.act(x) x = self.fc2(x) x = self.dropout(x) return x ``` ##### 2.2.2 Encoder代码整合 ```python class Block(nn.Module): def __init__(self, dim: int, # 输入token和qkv矩阵的dim维度 num_heads: int = 8, # 分类头数量 qkv_bias: bool = False, # qkv的bias qk_scale: float = None, in_channels: int = 768, # mlp的输入 expand_ratio: int = 4, # mlp hideen层的拓展维度 mlp_drop_ratio: float = 0., # mlp 最后dropout的概率 drop_path_ratio: float = 0. # residual 结构最后的drop_path概率 ): super(Block, self).__init__() self.norm1 = nn.LayerNorm(dim) self.attn1 = Attention(dim=dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale) self.drop_path1 = DropPath(drop_path_ratio) self.norm2 = nn.LayerNorm(dim) self.mlp1 = Mlp(in_channels=in_channels, expand_ratio=expand_ratio, drop_ratio=mlp_drop_ratio) self.drop_path2 = DropPath(drop_path_ratio) def forward(self, x): x = self.attn1(self.norm1(x)) x = self.drop_path1(x) + x x = self.mlp1(self.norm1(x)) x = self.drop_path2(x) + x return x ``` ## 3 ViT模型整合及代码整合 ### 3.1 MLP分类头的构建 上面通过Transformer Encoder后输出的shape和输入的shape是保持不变的，以ViT-B/16为例，输入的是[197, 768]输出的还是[197, 768]。注意，在Transformer Encoder后其实还有一个Layer Norm没有画出来，后面有我自己画的ViT的模型可以看到详细结构。这里我们只是需要分类的信息，所以我们只需要提取出[class]token生成的对应结果就行，即[197, 768]中抽取出[class]token对应的[1, 768]。接着我们通过MLP Head得到我们最终的分类结果。 ### 3.2 模型的参数配置 在论文的Table1中有给出三个模型（Base/ Large/ Huge）的参数，在源码中除了有Patch Size为16x16的外还有32x32的。其中的Layers就是Transformer Encoder中重复堆叠Encoder Block的次数，Hidden Size就是对应通过Embedding层后每个token的dim（向量的长度），MLP size是Transformer Encoder中MLP Block第一个全连接的节点个数（是Hidden Size的四倍），Heads代表Transformer中Multi-Head Attention的heads数。 | Model | Patch size | Layers | Hidden Size | MLP size | Heads | | ------------ | ------------ | ------------ | ------------ | ------------ | ------------ | | vit-base | 16 | 12 | 768 | 3072 | 12 | | vit-large | 16 | 24 | 1024 | 4096 | 16 | | vit-huge | 14 | 32 | 1280 | 5120 | 16 | ### 3.3 ViT模型整合及代码实现 ```python class Vit(nn.Module): def __init__(self, img_size: int = 224, patch_size: int = 16, in_c: int = 3, num_classes: int = 5, # 输出的类别 depth: int = 12, # block块重复的次数 dim: int = 768, # 输入token和qkv矩阵的dim维度 num_heads: int = 8, # 分类头数量 qkv_bias: bool = False, # qkv的bias qk_scale: float = None, expand_ratio: int = 4, # mlp hideen层的拓展维度 drop_ratio: float = 0., mlp_drop_ratio: float = 0., # mlp 最后dropout的概率 drop_path_ratio: float = 0., # residual 结构最后的drop_path概率 ): super(Vit, self).__init__() # patch_embdding self.patch_emb = PatchEmbedding(img_size=img_size, patch_size=patch_size, in_c=in_c) # class token self.cls_token = nn.Parameter(torch.zeros(1, 1, dim)) # pos_embedding num_grids = (img_size // patch_size) ** 2 self.pos_emb = nn.Parameter(torch.zeros(1, num_grids+1, dim)) self.pos_drop = nn.Dropout(p=drop_ratio) dpr = [x.item() for x in torch.linspace(0, drop_path_ratio, depth)] self.blocks = nn.Sequential(*[ Block(dim=dim, num_heads=num_heads, mlp_drop_ratio=mlp_drop_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale, in_channels=dim, expand_ratio=expand_ratio, drop_path_ratio=dpr[i]) for i in range(len(dpr)) ]) self.norm = nn.LayerNorm(dim) # classifier heads self.head = nn.Linear(dim, num_classes) if num_classes > 0 else nn.Identity() # weights init self._init_weights() def forward(self, x): B, C, H, W = x.size() # [B, C, H, W] -> [B, num_patches, embed_dim] x = self.patch_emb(x) # [1, 1, embed_dim] -> [B, 1, embed_dim] cls_token = self.cls_token.expand(B, -1, -1) # [B, num_patches, embed_dim] + [B, 1, embed_dim] = [B, num_patches + 1, embed_dim] x = torch.cat([cls_token, x], dim=1) x = self.pos_drop(self.pos_emb + x) x = self.blocks(x) x = self.norm(x) x = x[:, 0] # 取得token x = self.head(x) return x ``` ———————————————— 参考链接：CSDN博主「太阳花的小绿豆」 原文链接：https://blog.csdn.net/qq_37541097/article/details/118242600 最后编辑：2024年04月23日 ©著作权归作者所有 赞 0 分享
最新回复