ViT 原理解析 (Transformers for Image Recognition at Scale)

ViT 原理解析 (Transformers for Image Recognition at Scale)

原创 小白 小白研究室 2024年06月10日 21:09 北京

如何将 transformer 应用到图像领域

Transformer模型最开始是用于自然语言处理(NLP)领域的，NLP主要处理的是文本、句子、段落等，即序列数据。

视觉领域处理的是图像数据，因此将Transformer模型应用到图像数据上面临着诸多挑战，理由如下：

• 与单词、句子、段落等文本数据不同，图像中包含更多的信息，并且是以像素值的形式呈现。

• 如果按照处理文本的方式来处理图像，即逐像素处理的话，复杂度较高，硬件难以实现。

• Transformer缺少CNNs的归纳偏差，比如平移不变性和局部受限感受野。

• CNNs是通过相似的卷积操作来提取特征，随着模型层数的加深，感受野也会逐步增加。但是由于Transformer的本质，其在计算量上会比CNNs更大。

• Transformer无法直接用于处理基于网格的数据，比如图像数据。

为了解决上述问题，Google的研究团队提出了ViT模型，它的本质其实也很简单，既然Transformer只能处理序列数据，那么我们就把图像数据转换成序列数据就可以了呗。下面来看下ViT是如何做的。

基本结构

另外，从网上也看到有人绘制了比较详细的算法结构图，对于理解 ViT 也是有比较大的帮助，就复用粘贴在这里供大家学习：

(结构图来自https://blog.csdn.net/weixin_42118657/article/details/121789116)

模块细节

将图片转换成 patches 序列

对于图像数据而言，其数据格式为[H, W, C]，是三维矩阵，明显不是Transformer想要的。所以需要先通过一个Embedding层来对数据做个变换。

首先将一张图片按给定大小分成一堆Patches。

以ViT-B/16为例，将输入图片(224x224)按照16x16大小的Patch进行划分，划分后会得到196个Patches。接着通过线性映射将每个Patch映射到一维向量中，以ViT-B/16为例，每个Patche数据shape为[16, 16, 3]通过映射得到一个长度为768的向量（后面都直接称为token）。[16, 16, 3] -> [768]

在代码实现中，直接通过一个卷积层来实现。以ViT-B/16为例，直接使用一个卷积核大小为16x16，步距为16，卷积核个数为768的卷积来实现。通过卷积[224, 224, 3] -> [14, 14, 768]，然后把H以及W两个维度展平即可[14, 14, 768] -> [196, 768]，此时正好变成了一个二维矩阵，正是Transformer想要的。

# https://github.com/lucidrains/vit-pytorch/blob/90be7233a3f55c29692a72da6ee4dcb5aab267d4/vit_pytorch/twins_svt.py#L59
class PatchEmbedding(nn.Module):
def __init__(self, *, dim, dim_out, patch_size):
super().__init__()
self.dim = dim
self.dim_out = dim_out
self.patch_size = patch_size

self.proj = nn.Sequential(
LayerNorm(patch_size ** 2 * dim),
nn.Conv2d(patch_size ** 2 * dim, dim_out, 1),
LayerNorm(dim_out)
)

def forward(self, fmap):
p = self.patch_size
fmap = rearrange(fmap, 'b c (h p1) (w p2) -> b (c p1 p2) h w', p1 = p, p2 = p)
return self.proj(fmap)

这里增加了 class token，class token的维度是[1,768]，然后将其与第1步得到的tokens进行拼接，即Cat([1, 768], [196, 768]) -> [197, 768]。

在传统CNN分类任务中，会对最后卷积输出的 feature map 进行一个 global average pooling 操作，用以进行最后的类别预测；在 vision transformer 里面能否进行相同的操作呢，即把 16 个 patch 的 token 进行一个 average pooling 来替代 class token。作者消融实验下来验证是可以的，但是要验证使用不同的学习率。论文中作者是为了尽可能的和 transformer 结构保持一致、所以才默认使用了 class token (In order to stay as close as possible to the original Transformer model)

添加 Position embedding

从公式可以看出，其实一个词语的位置编码是由不同频率的余弦函数函数组成的，从低位到高位，余弦函数对应的频率由 1 降低到了 110000 ，按照论文中的说法，也就是，波长从 2𝜋 增加到了 10000⋅2𝜋 。

class PositionalEncoding(nn.Module):

def __init__(self, d_model, max_len=5000):
super(PositionalEncoding, self).__init__()
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0).transpose(0, 1)
#pe.requires_grad = False
self.register_buffer('pe', pe)

def forward(self, x):
return x + self.pe[:x.size(0), :]

Transformer Encoder

将patch embedding和class token拼接起来输入标准的Transformer Encoder中。

Transformer的架构图：

Transformer是由一堆encoder和decoder形成的，那encoder一般的架构图如下：

Transformer Encoder其实就是重复堆叠Encoder Block L次，实现细节如下面代码所示（layer normalization -> multi-head attention -> drop path -> layer normalization -> mlp -> drop path）

class Block(nn.Module):
def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
super().__init__()
# 将每个样本的每个通道的特征向量做归一化
# 也就是说每个特征向量是独立做归一化的
# 我们这里虽然是图片数据，但图片被切割成了patch，用的是语义的逻辑
self.norm1 = norm_layer(dim)
self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop)
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
self.norm2 = norm_layer(dim)
mlp_hidden_dim = int(dim * mlp_ratio)
# 全连接，激励，drop，全连接，drop,若out_features没填，那么输出维度不变。
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
def forward(self, x):
# 最后一维归一化，multi-head attention, drop_path
# (B, N, C) -> (B, N, C)
x = x + self.drop_path(self.attn(self.norm1(x)))
# (B, N, C) -> (B, N, C)
x = x + self.drop_path(self.mlp(self.norm2(x)))

return x

分类头

注意Transformer Encoder的输出其实也是一个序列，但是在ViT模型中只使用了class token的输出，将其送入MLP模块中，去输出最终的分类结果。

参考资料

• https://github.com/google-research/vision_transformer

• https://blog.csdn.net/weixin_42118657/article/details/121789116

• https://blog.csdn.net/lsb2002/article/details/135320751

• https://zhuanlan.zhihu.com/p/338592312

• 如何理解Transformer论文中的positional encoding，和三角函数有什么关系？- 猛猿的回答 - https://www.zhihu.com/question/347678607/answer/2301693596

• https://zhuanlan.zhihu.com/p/427388113

• 【ViT论文逐段精读【论文精读】】 https://www.bilibili.com/video/BV15P4y137jb/?share_source=copy_web&vd_source=fe71fe6eb662a29ef2a5866a5fa7b14a