VIT（AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE）论文网络结构简析

源码来源作者：Adenialzz（本文仅对此源码进行理解，仅为个人观点，如有错误，请指正！）

论文地址：https://arxiv.org/abs/2010.11929

上述源码输入案例参数（depth=6, img_size=256, patch_size=32, num_classes=1000, dim=1024, heads=16, mlp_dim=2048, dropout=0.1, embedding_dropout=0.1）

model_vit = ViT(
     
        depth = 6,#使用6个transformer encoder模块   
        image_size = 256,#使用图像尺寸为256*256
        patch_size = 32,#将图像分为32*32份
        num_classes = 1000,#分类目标类数
        dim = 1024,#线性映射目标维度尺寸
        heads = 16,#多头注意力机制中的头数
        mlp_dim = 2048,#mlp中间层的维度（一般为第一层输入的4倍）
        dropout = 0.1,#dropout参数
        emb_dropout = 0.1#emb中dropout参数
    )

本文根据代码运行顺序进行讲解。

代码流程：

1. 输入图像（图像尺寸(16,3,256,256)）16为batchsize，3为channels 

img = torch.randn(16, 3, 256, 256)
preds = model_vit(img) 

2. 对图像进行分割。

尺寸变换：(6,3,256,256)---->(16,8*8,32*32*3)=(16,64,3072)

尺寸变换方法：可参考：einops和einsum：直接操作张量的利器_Adenialzz的博客-CSDN博客

尺寸变换目的：满足patchsize的输入要求（满足步骤3线性层的输入要求）

Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_height, p2=patch_width)

3. 将图像通过线性层（3072，1024），进行线性映射。作用：将像素值转换为特征向量

尺寸变换：(16,64,3072)---->(16,64,1024)

nn.Linear(patch_dim, dim)#patch_dim=3072,dim=1024

4. 定义一个可学习变量cls_token(1,1,1024),并将其复制为与batchsize相匹配大小。

cls_token尺寸变换：(1,1,1024)---->(16,1,1024)

cls_token变量目的：根据论文的结论，作者认为此变量能够对其他64个输入块进行总结，综合各输入块的输出。

self.cls_token = nn.Parameter(torch.randn(1, 1, dim))#dim=1024
cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b)

5. 将cls_token拼接进特征向量中。（ cat[ (16,1,1024) , (16,64,1024) ] = (16,65,1024) = X ）

x = torch.cat((cls_tokens, x), dim=1)

6. 定义一个可学习变量pos_embedding (位置信息)

pos_embedding尺寸：(16,65,1024)

定义pos_embedding目的：作者认为，图像分块后，每一块的相对位置信息对图像识别是有好处的，根据论文实验所得，定义一维位置信息与定义二维位置信息差别不大，一维位置信息可学习到二维位置的相对位置信息。（论文图7Center图）

所以将位置信息加入X中,X尺寸不变(16,65,1024)

self.pos_embedding = nn.Parameter(torch.randn(1, num_patches+1, dim))
x += self.pos_embedding[:, :(n+1)]

至此，上述步骤为图像解析与相关特征信息嵌入，以下部分为encoder部分

7. 将X输入layernorm层，对维度=dim=1024的那层进行layernorm归一化操作。

X尺寸信息：(16,65,1024)

PreNorm(dim, Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout)),

self.norm = nn.LayerNorm(dim)
self.norm(x)

8. 将X放入线性层（1024，3072）

X-->qkv 尺寸变换：(16,65,1024)---->(16,65,3072)---->(3,16,65,1024)

此线性层目的：自注意力机制，按照原理而言q=k=v，但此处通过线性层变换，再拆分为q,k,v三部分，q!=k!=v。此种做法可再一定程度上提升模型的拟合能力。

self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
qkv = self.to_qkv(x).chunk(3, dim=-1)

9. 将q,k,v三部分通过qkv拆出，b=16,h=16,n=65,d=64.

q.shape=k.shape=v.shape=(16,16, 65, 64)

拆分为此尺寸目的：b为batchsize不变，h为heads满足多头注意力输入要求尺寸，n为有65个输入(64块图片加上cls_token变量)，d为每张图片的总像素：8*8=64(imgsize256/patchsize32=8)

q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), qkv)

10. 计算q，k的乘积（计算矩阵向量相似度）

similarity / dots 尺寸：(16,16, 65, 65)

self.scale = dim_head ** -0.5
dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale

11. 将dots 通过softmax函数，在最后一个维度计算权重attn

attn 尺寸：(16,16, 65, 65)

self.attend = nn.Softmax(dim=-1)
attn = self.attend(dots)

12. 将权重（attn）与v相乘，得到输出out

out 尺寸：(16,16, 65, 64)

out = einsum('b h i j, b h j d -> b h i d', attn, v)

13. 多头机制结束，重新将输出拼接为原尺寸

out 尺寸：(16,16, 65, 64)---->(16, 65, 64*16)=(16, 65, 1024)

out = rearrange(out, 'b h n d -> b n (h d)')

14. 将out输入线性层（1024，1024）

out 尺寸：(16, 65, 1024)---->(16, 65, 1024)

self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout),
        )

return self.to_out(out)

15. 残缺连接（第14步结果out+第7步的结果X）

X 尺寸：(16, 65, 1024)+(16, 65, 1024)=(16, 65, 1024)

x = attn(x) + x

16. 再将X输入layernorm层

PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout))

self.norm = nn.LayerNorm(dim)
self.norm(x)

17. 再将X输入feedforward层中

feedforward层=线性层1（1024，2048）+GELU线性激活函数+线性层2（2048，1024）

X 尺寸：(16, 65, 1024)--线性层1-->(16, 65, 2048)--GELU-->--线性层2-->(16, 65, 1024)

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim), 
            nn.Dropout(dropout)
        )
    def forward(self, x):
        return self.net(x)

PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout))

18.残缺连接（第15步结果X+第17步的X结果）

x = ff(x) + x

至此，上述步骤为encoder部分，根据参数depth将步骤8-18步循环depth次

以下部分为分类器部分。

19. 选择训练目标：64个输出的mean值(无cls_token情况)或extra learning position值（cls_token）

X 尺寸：(16, 65, 1024)---->(16,1024)

x = x.mean(dim=1) if self.pool == 'mean' else x[:, 0]  

20. 将X 输入layernorm层,而后将X 输入线性层(1024,num_classes)

X 尺寸：(16,1024)---->(16,num_classes)

self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_classes)
        )

return self.mlp_head(x)   

结束正向传播过程。