VIT模型个人笔记

前言

VIT模型即vision transformer，其想法是将在NLP领域的基于自注意力机制transformer模型用于图像任务中，相比于图像任务中的传统的基于卷积神经网络模型，VIT模型在大数据集上有着比卷积网络更强的效果和更节约的成本。

transformer

transformer模型是用于自然语言处理的一个基于注意力机制的模型，其图如下所示，该模型主要由解码器和编码器两部分组成。在nlp相关任务中，处理的数据对象主要是句子或句子对，因此，在训练之前，存在一个由多个token组成的字典。而输入模型的数据为形如大小为NF的向量，其中N为tokens的数量，F为表示每个token语义信息的向量长度，然后通过线性变化，加入位置信息得到大小为ND的向量，其中D为论文规定的输入给注意力层的向量大小。

注意力机制

从人认识句子的模式考虑，面对一个句子的多个单词时，我们对于不同的词组的关注度自然也存在不同，基于这个想法提出的注意力机制的抽象模型如下：query表示待处理目标，key-value表示键值对，输出attention本质上为values的加权和，而这里的权重即为注意力系数，其计算公式如下:其中，Q K分别表示query和健值的向量矩阵，dk为二者的大小。另外，还可将QK分为多个子矩阵的拼接，分别计算注意力，最后将结果拼接回去原来的大小。这种方法称为多头注意力机制

从transformer到VIT

若想将transformer模型用于对于二维图像的处理，首先需要解决的问题即是如和将二维图像转化为可输入给transformer模型的1维向量，自然想到将大小维NN的图像分为pp的小图像patch，再将每个patch展开这样得到大小维度为（NN/PP)(PP*3)的向量，3表示rgb三通道，再将该向量经过一个线性变化使其特征维度变为D，即可继续输入给transformer进行训练。其模型结构图如下：

如图所示，在做图像分类任务时，需要增加一个表示类别的,token,最后在加上位置编码信息，得到的向量作为最终的transformer的输入。另外，在vit模型中，QKV都是同样来自图像patch的同样大小的三个向量。
整个流程用公式表示如下所示其中第一步即图像的预处理，包扩图像分块，增加类别信息，位置信息，E表示将表示图像信息的向量通过线性变化进行维度转化；第二个式子为MSA部分，包括多头自注意力、跳跃连接 (Add) 和层规范化 (Norm) 三个部分，可以重复L个MSA block；第三个式子为MLP部分，包括前馈网络 (FFN)、跳跃连接 (Add) 和层规范化 (Norm) 三个部分。

代码实现

pytorch中可直接调用搭建vit模型，相关代码如下所示：

```python
import torch
import numpy as np
from vit_pytorch import ViT
import torchvision
from torchsummary import summary

#创建VIt模型实例
v=ViT(
    image_size=256,           #原始图像大小 256*256
    patch_size=32,            #图像块的大小，即将原始图像按块大小切割
    num_classes=10,           #分类数量
    dim=1024,                 #transformer隐藏变量维度。即输入给transform模型的特征维度
    depth=6,                  #transform编码器层数
    heads=6,                  #msa中多头注意力机制的头数
    mlp_dim=2048,                
    dropout=0.1,
    emb_dropout=0.1)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

如输入一个图像，能得到一个分类结果。

```python
img=torch.randn(1,3,256,256)        #batch_size*C*h*w
preds=v(img)
preds.size()
1
2
3
4

##从头搭建vit模型
通过上面的模型原理介绍，VIT模型其实是以transformer为基础的，因此需要先搭建ffn，注意力机制等组件，再将其与图像预处理，编码嵌入层等拼接起来得到一个完整的vit模型

import torch
from torch import nn , einsum
import torch.nn.functional as F
from einops import rearrange , repeat     #einops是一个处理张量的第三方库 output_tensor = rearrange(input_tensor, 't b c -> b c t')
from einops.layers.torch import Rearrange                                            # 沿着某一维复制 output_tensor = repeat(input_tensor, 'h w -> h w c', c=3)
                                            # Rearrange('b c h w -> b (c h w)'),
def pair(t):    #辅助函数，生成元组
    return t if isinstance(t,tuple) else(t,t)
#搭建layernorm层和ffn层。其中fnn主要实现向量的线性尺度变化

#规范化层封装
class preNorm(nn.Module):
    def __init__(self,dim,fn):
        super().__init__()
        self.norm=nn.LayerNorm(dim)
        self.fn=fn
    def forward(self,x,**kwargs):
        return self.fn(self.norm(x),**kwargs)
#FFM层
class FeedForward(nn.Module):
    def __init__(self,dim,hidden_dim,dropout=0.1):
        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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

注意力机制实现代码

```python
#注意力机制层
class Attention(nn.Module):
    def __init__(self,dim,heads,dim_heads,dropout=0.1):
        super().__init__()
        inner_dim=dim_heads*heads
        project_out=not(heads ==1 and dim_heads==dim)
        
        self.heads=heads
        self.scale=dim_heads**-0.5
        
        self.attend=nn.Softmax(dim=-1)
        self.to_qkv=nn.Linear(dim,inner_dim*3,bias=False)
        
        self.to_out=nn.Sequential(
         nn.Linear(inner_dim,dim),
         nn.Dropout(dropout)
        ) if project_out else nn.Identity()
        
        
        
    def forward(self,x):
        #注意力系数attention Attention(Q,K,V ) = softmax( QKT√ dk)V
        b, n, _, h = *x.shape, self.heads
        #print('x',x.size())
        qkv=self.to_qkv(x).chunk(3,dim=-1)
        #qkv=torch.tensor([item.detach().numpy() for item in qkv])
        #print('a',qkv.size())
        q,k,v=map(lambda t : rearrange(t,'b n (h d) -> b h n d  ',h=h),qkv)     #q：[batches.heads,num_patches,head_dim=dim/heads]
        #print('q',q.size())
        dots=einsum('b h i d , b h j d -> b h i j',q,k)*self.scale 
        attn=self.attend(dots)
        out=einsum('b h i j, b h j d -> b h i d',attn,v)
        print(out.size())
        out=rearrange(out,'b h n d -> b n (h d)')              #将多头注意力的各个投拼接回到原来的大小dim
        print('out',out.size())
        return self.to_out(out)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

将上面组件拼接得到transformer模型

```python
#搭建transformer
class Transformer(nn.Module):
    def __init__(self,dim,depth,heads,dim_head,mlp_dim,dropout=0.1):
        super().__init__()
        self.layers=nn.ModuleList([])
        for _ in range(depth):                                       #多头注意力
            self.layers.append(nn.ModuleList([
            preNorm(dim,Attention(dim,heads=heads,dim_heads=dim_head,dropout=dropout)),
            preNorm(dim,FeedForward(dim,mlp_dim,dropout=dropout))
            ]
            ))
    def forward(self,x):
        for attn,ff in self.layers:
            x=attn(x)+x
            x=ff(x)+x
        return x 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

搭建vit模型代码如下：

class ViT(nn.Module):
    def __init__(self,*,image_size,patch_size,num_classes,dim,depth,heads,mlp_dim,pool='cls',channels=3,dim_head,dropout=0.1,emb_dropout=0.1):
        super().__init__()
        image_height,image_width=pair(image_size)
        patch_height,patch_width=pair(patch_size)
        
        assert image_height % patch_height == 0 and patch_height % patch_width == 0
        num_patches=(image_height // patch_height)* (image_width//patch_width)
        patch_dim=channels*patch_height*patch_width                 #ji
        assert pool in {'cls', 'mean'}
        
        #定义块嵌入
        self.to_patch_embedding=nn.Sequential(
              Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c )',p1=patch_height,p2=patch_width),
              nn.Linear(patch_dim,dim),)
        #定义位置编码
        self.pos_embedding=nn.Parameter(torch.randn(1,num_patches+1,dim))
        #定义类别向量
        self.cls_token=nn.Parameter(torch.randn(1,1,dim))
        self.dropout=nn.Dropout(emb_dropout)
        
        self.transformer=Transformer(dim,depth,heads,dim_head,mlp_dim,dropout)
        
        self.pool=pool
        self.to_latent=nn.Identity()
        
        #定义MLP
        self.mlp_head=nn.Sequential(
          nn.LayerNorm(dim),
          nn.Linear(dim,num_classes)
        )
        
    def forward(self,img):
        #块嵌入
        x=self.to_patch_embedding(img)
        b,n,_=x.shape
        #给每个batch(即每张图片）追加类别向量
        cls_tokens=repeat(self.cls_token,'() n d -> b n d ',b=b)
        x=torch.cat((cls_tokens,x),dim=1)  #将类别向量与分块后向量拼接在一起
        #添加位置编码 
        x += self.pos_embedding[:,:(n+1)]  #直接是向量相加 不改变向量大小
        x=self.dropout(x)
        x=self.transformer(x)
        x=x.mean(dim=1) if self.pool == 'mean' else x[:,0]
        
        x=self.to_latent(x)
        #返回mlp后的分类结果
        return self.mlp_head(x)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

下面，我们通过输入向量来测试该vit模型的输出
假设输入向量大小为（4,3,224,224)表示batches为4的rgb三通道224*224的图片，vit模型参数如下

#测试该vit模型
images=torch.randn([4,3,224,224])
mymodel=ViT(image_size=224,patch_size=16,num_classes=10,dim=768,depth=2,heads=12,mlp_dim=3072,pool='cls',channels=3,dim_head=64,dropout=0.1,emb_dropout=0.1)

out=mymodel(images)
out.size()
1
2
3
4
5
6

其中的参数表示，输入图片大小为224，patch的大小为16，因此得到（224224/1616）（16163)为196768的向量。transformer特征维度也为768，多头注意力机制取12，depth=2表示重复两个transformer块，mlp隐藏层大小取768.该实例中做的是一个10分类问题。最后得到的输出形状为i（4*10)

最后我们通过torchsummary包来看一看该模型每层的输出大小

summary(mymodel,input_size=(3,224,224))
1