Vision transformer（ViT）输入后维度变化情况 + 代码实现，如何理解ViT？

大家好，欢迎来到IT知识分享网。

11.1 Vision Transformer(vit)网络详解_哔哩哔哩_bilibili

文章地址：Vision Transformer详解_霹雳吧啦Wz-CSDN博客

一句话来解释vit：

将图像的这些东西（C，H，W）变为NLP处理的方式（T, D），所以就消除了CV中的通道C，宽高H\W的概念。

从而使C就是D，H和W变为了T。那么原来对CV的某些实现，例如IN, LN..就可以通过这种方式进行替换套用

整体运行：

其中关键的图：

图像的维度变化：

1、输入 –>（N，C，H，W）

（N，C，H，W）维度的图像（4，3，32，32）, N代表批次

2、经过timm.PatchEmbed() –> （N, T, D）

PatchEmbed()的patch size 为 (8, 8)，input_dim设置为1024：

使用 patch size 为 (8, 8) 将输入图像划分为图像块，得到（32//8）*（32//8）=16个图像块，对应image_sequence_length，后面的维度就是定义的input_dim，所以此时的维度是（4，16，1024）

添加 Class Token，将维度变为（4，17，1024）

此时这里你也可以加入位置编码，维度不会发生变化（4，17，1024）

self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, input_dim)

3、多头自注意力机制（Multi-Head Self-Attention）–> （N, T, D）

如果默认的头数是 h（通常为 12），并且头维度（head dim）是 d_h（通常为 64），则每个头产生的输出维度为（4，17，64）。embed_dim对应d_h*h = 768

对patch后的token进行线性变换得到 Q, K, V –> （N，T，d_h*h​）

对输入（N, T, D）进行三个线性变换，得到Q,K,V，每个变换的权重矩阵维度为 (d_model​, d_h*h​)，即(1024->12×64=768)。h为头的个数

Q 的维度为 (4,17,d_h*h​)，K 的维度为 (4,17,d_h*h​)，V 的维度为 (4,17,d_h*h​)。

self.W_q = nn.Linear(input_dim, embed_dim) self.W_k = nn.Linear(input_dim, embed_dim) self.W_v = nn.Linear(input_dim, embed_dim)

将Q, K, V拆分为多头：–> （N，T，h，d_h​）–> （N，h，T，d_h​）

将得到的 Q,K,V 分别拆分为 h 个头，每个头的维度为 (4,17,d_h)。

这里的拆分是在最后一维上进行的。

# 应用线性变换生成Q, K, V Q = self.W_q(x).view(batch_size, image_sequence_length, self.num_heads, self.head_dim).transpose(1, 2) # 实际为(batch_size, self.num_heads, image_sequence_length, self.head_dim) K = self.W_k(x).view(batch_size, image_sequence_length, self.num_heads, self.head_dim).transpose(1, 2) V = self.W_v(x).view(batch_size, image_sequence_length, self.num_heads, self.head_dim).transpose(1, 2) 

对每个头进行注意力计算：–> （N，h，T，h_d​）

对 Q 和 K 进行点积，然后进行缩放操作，最后应用 softmax 函数。

将注意力分数乘以 V，得到每个头的注意力输出(4,17,d_h)。

# 得到(b,num_heads,image_sequence_length,image_sequence_length) attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim 0.5) # 应用Softmax获取注意力权重 # 计算块与块之间的注意力 attn_weights = torch.softmax(attn_scores, dim=-1) # 使用注意力权重加权V attn_output = torch.matmul(attn_weights, V)

合并多头输出：–> （N，T，D​）

将每个头的注意力输出按最后一维度进行连接（concatenate），得到多头注意力的输出。多头注意力的输出的维度为 (4,17,768)。这里768=h*d_h

映射为（4，17，1024），加残差，维度不变（4，17，1024）

attn_output = attn_output.view(batch_size, image_sequence_length, self.embed_dim)

4、LayerNormalization–> （N, T, D）

对于每个位置，进行 LayerNormalization 维度不变（4，17，1024）。

class PreNorm(nn.Module): def __init__(self,dim,fn):#fn指norm之后的函数 如上图，可以接mha或者mlp super().__init__() self.norm = nn.LayerNorm(dim) self.fn = fn def forward(self,x): return self.fn(self.norm(x))

5、Feedforward 网络 –> （N, T, D）

如果默认的 Feedforward 网络的中间维度（为 2048）：
输入维度：（4，17，1024）
变化为：（4，17，2048）
输出维度：（4，17，1024）
加残差，维度不变（4，17，1024）

#ffn也就是mlp class FeedForward(nn.Module): def __init__(self,dim,hidden_dim,drop_out=0.3): super().__init__() self.net = nn.Sequential( nn.Linear(dim,hidden_dim), nn.GELU(), nn.Dropout(drop_out), nn.Linear(hidden_dim,dim), nn.Dropout()) def forward(self,x): return self.net(x)

6、MLP head –> （N,  D）

（4，17，1024）全局池化或者只取cls token（4，1024）mlp输出1000类->（4，1000）

保姆级VIT详解输入维度变化以及对应代码实现_vit的输入-CSDN博客