• VIT 源码详解

    1.项目配置说明

            参数说明：

    数据集：  

           --name cifar10-100_500

           --dataset cifar10

    哪个版本的模型： 

           --model_type ViT-B_16

    预训练权重： 

           --pretrained_dir checkpoint/ViT-B_16.npz

    2.patch embeding与position_embedding

            对于图像编码，以VIT - B/16为例，首先用卷积核大小为16*16、步长为16的卷积，对图像进行变换，此时图像维度变成16 * 768 * 14 * 14，再变换维度为[16, 196, 768]，然后将维度为16*1*768的0patch相连。

            对于位置编码，构建一个1 * 197 * 768的向量

            最后，将图像编码与位置编码相加就完成了本次编码。

    代码如下：

    1. class Embeddings(nn.Module):
    2.     """Construct the embeddings from patch, position embeddings.
    3.     """
    4.     def __init__(self, config, img_size, in_channels=3):
    5.         super(Embeddings, self).__init__()
    6.         self.hybrid = None
    7.         img_size = _pair(img_size)
    8.  
    9.         # patch_size 大小 与 patch数量  n_patches
    10.         if config.patches.get("grid") is not None:
    11.             grid_size = config.patches["grid"]
    12.             patch_size = (img_size[0] // 16 // grid_size[0], img_size[1] // 16 // grid_size[1])
    13.             n_patches = (img_size[0] // 16) * (img_size[1] // 16)
    14.             self.hybrid = True
    15.         else:
    16.             patch_size = _pair(config.patches["size"])
    17.             n_patches = (img_size[0] // patch_size[0]) * (img_size[1] // patch_size[1])
    18.             self.hybrid = False
    19.  
    20.         # 使用混合模型
    21.         if self.hybrid:
    22.             self.hybrid_model = ResNetV2(block_units=config.resnet.num_layers,
    23.                                          width_factor=config.resnet.width_factor)
    24.             in_channels = self.hybrid_model.width * 16
    25.         # patch_embeding 16 * 768 * 14 * 14
    26.         self.patch_embeddings = Conv2d(in_channels=in_channels,
    27.                                        out_channels=config.hidden_size,
    28.                                        kernel_size=patch_size,
    29.                                        stride=patch_size)
    30.         # 初始化 position_embeddings: 1 * 197 * 768
    31.         self.position_embeddings = nn.Parameter(torch.zeros(1, n_patches+1, config.hidden_size))
    32.         # 初始化第 0 个patch，表示分类特征 1*1*768
    33.         self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
    34.         # dropout层
    35.         self.dropout = Dropout(config.transformer["dropout_rate"])
    36.  
    37.     def forward(self, x):
    38.         print(x.shape)
    39.         B = x.shape[0]
    40.         # 拓展cls_tokens的维度:16 *1*768
    41.         cls_tokens = self.cls_token.expand(B, -1, -1)
    42.         print(cls_tokens.shape)
    43.         # 混合模型
    44.         if self.hybrid:
    45.             x = self.hybrid_model(x)
    46.         # 编码:16 * 768 * 14 * 14
    47.         x = self.patch_embeddings(x)
    48.         print(x.shape)
    49.         # 变换维度:16 * 768 * 14 * 14-->[16, 768, 196]
    50.         x = x.flatten(2)
    51.         print(x.shape)
    52.         # [16, 768, 196] --> [16, 196, 768]
    53.         x = x.transpose(-1, -2)
    54.         print(x.shape)
    55.         # 加入分类特征patch
    56.         x = torch.cat((cls_tokens, x), dim=1)
    57.         print(x.shape)
    58.  
    59.         # 加入位置编码
    60.         embeddings = x + self.position_embeddings
    61.         print(embeddings.shape)
    62.         # dropout层
    63.         embeddings = self.dropout(embeddings)
    64.         print(embeddings.shape)
    65.         return embeddings

    3.ecoder 

    多头注意力模块:

            首先构建q,k,v三个辅助向量，因为我们采用多头注意力机制(12个)，首先，我们需要将q,k,v维度从16, 197, 768转换成16, 12, 197, 64，然后获得q,k的相似性qk，因为获得的是两两之间的关系，所以维度为16, 12, 197, 197，消除量纲，经过softmax后，得到提取到的特征向量qkv，维度为16, 12, 197, 64，再将维度还原成16, 197, 768

    1. class Attention(nn.Module):
    2.     def __init__(self, config, vis):
    3.         super(Attention, self).__init__()
    4.         self.vis = vis
    5.         # heads数量
    6.         self.num_attention_heads = config.transformer["num_heads"]
    7.         # 每个head的向量维度
    8.         self.attention_head_size = int(config.hidden_size / self.num_attention_heads)
    9.         # 总head_size
    10.         self.all_head_size = self.num_attention_heads * self.attention_head_size
    11.         # query向量
    12.         self.query = Linear(config.hidden_size, self.all_head_size)
    13.         # key向量
    14.         self.key = Linear(config.hidden_size, self.all_head_size)
    15.         # value向量
    16.         self.value = Linear(config.hidden_size, self.all_head_size)
    17.         # 全连接层
    18.         self.out = Linear(config.hidden_size, config.hidden_size)
    19.         # dropout层
    20.         self.attn_dropout = Dropout(config.transformer["attention_dropout_rate"])
    21.         self.proj_dropout = Dropout(config.transformer["attention_dropout_rate"])
    22.  
    23.         self.softmax = Softmax(dim=-1)
    24.  
    25.     def transpose_for_scores(self, x):
    26.         # 维度:16, 197, 768-->16,197,12,64
    27.         new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
    28.         # print(new_x_shape)
    29.         x = x.view(*new_x_shape)
    30.         # print(x.shape)
    31.         # print(x.permute(0, 2, 1, 3).shape)
    32.         # 16,197,12,64 --> 16, 12, 197, 64
    33.         return x.permute(0, 2, 1, 3)
    34.  
    35.     def forward(self, hidden_states):
    36.         # print(hidden_states.shape)
    37.         # q,k,v:16, 197, 768
    38.         mixed_query_layer = self.query(hidden_states)
    39.         # print(mixed_query_layer.shape)
    40.         mixed_key_layer = self.key(hidden_states)
    41.         # print(mixed_key_layer.shape)
    42.         mixed_value_layer = self.value(hidden_states)
    43.         # print(mixed_value_layer.shape)
    44.         # q,k,v:16, 197, 768-->16, 12, 197, 64
    45.         query_layer = self.transpose_for_scores(mixed_query_layer)
    46.         # print(query_layer.shape)
    47.         key_layer = self.transpose_for_scores(mixed_key_layer)
    48.         # print(key_layer.shape)
    49.         value_layer = self.transpose_for_scores(mixed_value_layer)
    50.         # print(value_layer.shape)
    51.         # q,k的相似性:16, 12, 197, 197
    52.         attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
    53.         # print(attention_scores.shape)
    54.         # 消除量纲
    55.         attention_scores = attention_scores / math.sqrt(self.attention_head_size)
    56.         # print(attention_scores.shape)
    57.         attention_probs = self.softmax(attention_scores)
    58.         # print(attention_probs.shape)
    59.         weights = attention_probs if self.vis else None
    60.         attention_probs = self.attn_dropout(attention_probs)
    61.         # print(attention_probs.shape)
    62.         # print(value_layer.shape)
    63.         # 特征向量：qkv:16, 12, 197, 64
    64.         context_layer = torch.matmul(attention_probs, value_layer)
    65.         # print(context_layer.shape)
    66.         # 16, 12, 197, 64-->16, 12, 197, 64
    67.         context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
    68.         # print(context_layer.shape)
    69.         # 16, 12, 197, 64-->16, 197, 768
    70.         new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
    71.         context_layer = context_layer.view(*new_context_layer_shape)
    72.         # print(context_layer.shape)
    73.         # 全连接层:16, 197, 768
    74.         attention_output = self.out(context_layer)
    75.         # print(attention_output.shape)
    76.         # dropout层
    77.         attention_output = self.proj_dropout(attention_output)
    78.         # print(attention_output.shape)
    79.         return attention_output, weights

          transformer encoder 

            对于输入的x,首先经过层归一化后，输入多头注意力机制，对结果进行残差连接，再经过层归一化，经过两层全连接，残差连接后，得到一个模块结果，堆叠L层，输出最终结果 

     

    1. class Block(nn.Module):
    2.     def __init__(self, config, vis):
    3.         super(Block, self).__init__()
    4.         # 序列的大小:768
    5.         self.hidden_size = config.hidden_size
    6.         # 层归一化
    7.         self.attention_norm = LayerNorm(config.hidden_size, eps=1e-6)
    8.         self.ffn_norm = LayerNorm(config.hidden_size, eps=1e-6)
    9.         # MLP层
    10.         self.ffn = Mlp(config)
    11.         # 多头注意力机制
    12.         self.attn = Attention(config, vis)
    13.  
    14.     def forward(self, x):
    15.         # print(x.shape)
    16.         # 16, 197, 768
    17.         h = x
    18.         # 层归一化
    19.         x = self.attention_norm(x)
    20.         # print(x.shape)
    21.         # 多头注意力机制
    22.         x, weights = self.attn(x)
    23.         # 残差连接
    24.         x = x + h
    25.         # print(x.shape)
    26.  
    27.         h = x
    28.         # 层归一化
    29.         x = self.ffn_norm(x)
    30.         # print(x.shape)
    31.         # MLP层
    32.         x = self.ffn(x)
    33.         # print(x.shape)
    34.         # 残差连接
    35.         x = x + h
    36.         # print(x.shape)
    37.         return x, weights

    整体架构

            对于输入x，进行patch embeding和position embeding后，此时维度为16*197*768，输入encoder中，经过L层的编码模块，取出第0个patch的编码结果(表示分类特征),输入分类层，得到预测结果。

    1. class VisionTransformer(nn.Module):
    2.     def __init__(self, config, img_size=224, num_classes=21843, zero_head=False, vis=False):
    3.         super(VisionTransformer, self).__init__()
    4.         self.num_classes = num_classes
    5.         self.zero_head = zero_head
    6.         self.classifier = config.classifier
    7.  
    8.         self.transformer = Transformer(config, img_size, vis)
    9.         self.head = Linear(config.hidden_size, num_classes)
    10.  
    11.     def forward(self, x, labels=None):
    12.         x, attn_weights = self.transformer(x)
    13.         print(x.shape)
    14.         # X.shape:16, 197, 768   logits.shape:16, 10
    15.         logits = self.head(x[:, 0])
    16.         print(logits.shape)
    17.         # 交叉熵
    18.         if labels is not None:
    19.             loss_fct = CrossEntropyLoss()
    20.             loss = loss_fct(logits.view(-1, self.num_classes), labels.view(-1))
    21.             return loss
    22.         else:
    23.             return logits, attn_weights

     

     

     

             

     

     

  • 相关阅读:
    FPGA 学习笔记：Verilog 实现LED流水灯控制
    springboot民办高校科研项目管理系统-计算机毕业设计源码54009
    【AXI4 verilog】手把手带你撸AXI代码（四、AXI4接口的RAM设计）
    Redis分布式锁剖析和几种客户端的实现
    Vue中methods实现原理
    [Spring boot] Spring boot 整合RabbitMQ实现通过RabbitMQ进行项目的连接
    k8s 配置存储之 Configmap & secret
    什么是仿射变换？
    unordered_map的键值不能直接用pair；而map 可以使用 pair 作为键值，而不需要额外定义哈希函数
    Linux命令学习—Apache 服务器（上）
  • 原文地址：https://blog.csdn.net/qq_52053775/article/details/126261070