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)

拆分为此尺寸目的:bbatchsize不变,hheads满足多头注意力输入要求尺寸,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. 计算qk的乘积(计算矩阵向量相似度)

similarity = \frac{(q*k^{T})}{\sqrt{dim\_head}}

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)   

结束正向传播过程。 

你可能感兴趣的:(pytorch,计算机视觉,transformer,深度学习)