ViT结构

Vision Transformer

图像输入尺寸为$[N,C,H,W]$，$C$通常为3，为了构建为$Transformer$需要的输入，将输入图像切分为$p_h\ast p_w\ast C$尺寸的$n$个小图块，合计切出$h\ast w$个小图块。

# reshape and flatten
[N, C, H, W] => [N, h*w, p_h * p_w * C] => [N, h*w, dim] # h = H // p_h, w = W // p_w, input flattened feature to nn.Linear, map into dim dimenstion.
# concat cls_tokens and add positional embedding
cls_token = nn.Parameter(torch.randn(1, 1, dim))
cls_token = repeat(cls_token, '() n d -> b n d', b=b)
pose_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim)
[N, n, dim] => [N, n + 1, dim] => [N, n + 1, dim] # n = h * w, cls_tokens -> positional embedding.

经过$n$个$encodinglayers$构建成的$Transformer$提取特征后，输入到$MLPhead$ 模块

[N, n + 1, dim] => [N, num_classes]

$Transformer$的$encodinglayer$模块的结构如下：

encoding layer = MSA + MLP
MSA: Multi-headed Self-Attention
MLP: Multi-Layer Perceptron

注意力模块如下：

多层注意力由多个单一的注意力模块提取信息后，concat到一起。