Vision Transformer(ViT)模型原理及PyTorch逐行实现

Vision Transformer(ViT)模型原理及PyTorch逐行实现

一、TRM模型结构

1.Encoder

1. Position Embedding 注入位置信息
2. Multi-head Self-attention 对各个位置的embedding融合（空间融合）
3. LayerNorm & Residual
4. Feedforward Neural Network 对每个位置上单独仿射变换（通道融合）
   • Linear1(large)
   • Linear2(d_model)
5. LayerNorm & Residual

2.Decoder

1. Position Embedding
2. Casual Multi-head Self-attention
3. LayerNorm & Residual
4. Memory-base Multi-head Cross-attention
5. LayerNorm & Residual
6. Feedforward Neural Network
   • Linear1(large)
   • Linear2(d_model)
7. LayerNorm & Residual

二、TRM使用类型

1. Encoder only 【 ViT 所使用的】
   • BERT、分类任务、非流式任务
2. Decoder only
   • GPT系列、语言建模、自回归生成任务、流式任务
3. Encoder-Decoder
   • 机器翻译、语音识别

三、TRM特点

1. 无先验假设（例如：局部关联性、有序建模性）
2. 核心计算在于自注意力机制，平方复杂度
3. 数据量的要求与归纳偏置【人类通过归纳法得到的经验，把这些经验带入到模型中，很多事物的共性】的引入成反比

四、Vision Transformer(ViT)

1. DNN perspective 图像的信息量主要还是聚集在一块区域上
   • image2patch 将图片切分成很多个块
   • patch2embedding 将每个块转换为向量
2. CNN perspective 从卷积的角度得到向量
   • 2D convolution over image 二维卷积
   • flatten the output feature map 把输出的卷积图拉直
3. class token embedding 占位符
4. position embedding
   • interpolation when inference
5. Transformer Encoder 只使用的Encoder
6. classification head 最后分类

五、ViT论文讲解

​ 首先将一副图片分为很多个块，每个块的大小都是不会变化的，图片即使大一点，只是序列更长一点。先左到右，再上到下，把图片拉直成一个序列的形状。把每个块中的像素点进行归一化，范围变为0到1之间，再把块里面的所有值通过一个线性变换映射到模型的维度，得到patchembedding，得到以后，我们为了做分类任务，还需要在序列的开头加上一个可训练的embedding，这个是随机初始化的。这样就构造出了一个n+1长度的序列，然后我们再加入position embedding，加上后的这个序列的表征就可以送入到TRM的encoder当中，最后取出结果中的我们加入的可训练的embedding位置上的值（输出状态），经过一个MLP，得到各个类别的概率分布，再通过一个交叉熵函数算出分类的loss，这样就完成了一个ViT模型的搭建。

六、代码实现

1.convert image to embedding vector sequence

1.通过DNN实现

import torch
import torch.nn as nn
import torch.nn.functional as F

def image2emb_naive(image,patch_size,weight):
    # image shape: bs*channel*h*w
    patch = F.unfold(image,kernel_size=patch_size,stride=patch_size).transpose(-1,-2)
    patch_embedding = patch @ weight
    return patch_embedding

# test code for image2emb
bs,ic,image_h,image_w=1,3,8,8
patch_size=4 # 每个块的大小为4*4（自定义）
model_dim=8 #将每个块映射成长度为8的向量（自定义）
patch_depth=patch_size*patch_size*ic
image=torch.randn(bs,ic,image_h,image_w) #初始化
weight=torch.randn(patch_depth,model_dim)#初始化

patch_embedding_navie=image2emb_navie(image,patch_size,weight)
print(patch_embedding_naive.shape) # [1,4,8],分成四块了，每块对应一个长度为8的向量 

2.通过CNN实现

import torch
import torch.nn as nn
import torch.nn.functional as F

def image2emb_conv(image,kernel,stride):
    conv_output=F.conv2d(image,kernel,stride=stride) # bs*oc*oh*ow
    bs,oc,oh,ow=conv_output.shape
    patch_embedding=conv_output.reshape((bs,oc,oh*ow)).transpose(-1,-2)
    return patch_embedding

# test code for image2emb
bs,ic,image_h,image_w=1,3,8,8
patch_size=4
model_dim=8
patch_depth=patch_size*patch_size*ic
image=torch.randn(bs,ic,image_h,image_w)
weight=torch.randn(patch_depth,model_dim) #model_dim是输出通道数目，patch_depth是卷积核的面积乘以输入通道数

kernel=weight.transpose(0,1).reshape((-1,ic,patch_size,patch_size)) # oc*ic*kh*kw
patch_embedding_conv=image2emb_conv(image,kernel,patch_size) # 二维卷积的方法得到embedding

2.prepend CLS token embedding

cls_token_embedding = torch.randn(1,model_dim,requires_grad=True)
token_embedding = torch.cat([[bs,cls_token_embedding],patch_embedding_conv],dim=1)

​ 提问：本身cls_token_embedding没有和任何样本矩阵有乘法联系，最后训练出来的也是一张确定的表，在做inference的时候，完全是一个常数的作用。送入transformer后，又与其他矩阵做了MHA，没搞懂用意何在啊？

​ 答：有联系啊，就是与其他时刻的sample做MHSA。这个token其实是取代了avg pool的作用，也就是说，你可以用avg pool得到分类的logits，也可以用采用cls token来得到分类的logits

​ 注意：cls_token_embedding作为batch_size中每一个序列的开始，应该对于每一个序列的开始都torch.cat同样的一个cls_token_embedding，然后都是对这同一个cls_token_embedding进行训练，所以这里的cls token embedding应该是二维的，1*model_dim，与batchsize无关。

3.add position embedding

max_num_token=16 #自定义
position_embedding_table = torch.randn(max_num_token,model_dim,requires_grad=True)
seq_len=token_embedding.shape[1] # 刚刚的1+4
position_embedding=torch.tile(position_embedding_table[:seq_len],[token_embedding.shape[0],1,1]) # 5,bs,1,1
token_embedding += position_embedding

4.pass embedding to Transformer Encoder

encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim,nhead=8)
transformer_encoder=nn.TransformerEncoder(encoder_layer,num_layers=6)
encoder_output=transformer_encoder(token_embedding)

5.do classification

cls_token_output=encoder_output[:,0,:] #拿到TRM的输出值
num_classes=10 # 自定义的类别数目
label=torch.randint(10,(bs,)) # 自定义的生成的label
linear_layer = nn.Linear(model_dim,num_classes) 
logits = linear_layer(cls_token_output)
loss_fn=nn.CrossEntropyLoss()
loss=loss_fn(logits,label)
print(loss)