PyTorch学习笔记:Vision Transformer(ViT)模型原理及PyTorch逐行实现_pytorch中的vit模型,如何使用

前言

首先再开始之前,我想问一下在座的各位实现过ViT吗,不会的扣1,会的小伙伴们扣脚指头,嘿嘿,开玩笑的,我也不知道,那么接下来我们一起学习如何写吧!

首先,本期还是 一样,推荐观看原视频进行学习,以便更加痛彻(顶部附有思维导图)

链接地址:28、Vision Transformer(ViT)模型原理及PyTorch逐行实现_哔哩哔哩_bilibili

---

 以上为ViT脑图(全局结构)

我们有多个角度去理解transformer,例如DNN,首先把图片切割成很多块,就是image2patch步骤;然后对很多个patch经过仿射变换得到一个新的向量(patch2embedding);又比如CNN角度,可以把图片得到embedding的过程理解成(类似于)卷积神经网络,我们首先会用一个二维的卷积,且kernel_size=stride,然后把输出卷积图拉直得到embedding(flatting过程)

为了做分类任务,ViT使用了多个embedding,发现可训练更好(即position embedding)

---

认识ViT

为了更加了解ViT,我们接下来欣赏一篇论文:2010.11929.pdf (arxiv.org)

为了让小伙伴们更加清楚,我找了一篇中译解读该论文的文章:论文解读：AN IMAGE IS WORTH 16X16 WORDS:TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE - 知乎 (zhihu.com)

对于该论文标题使用:An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

引用'一图胜千言'这个话改编为一图胜16*16,相当于16*16个像素点看成一个整体(embedding)来进行解读更加好

本文中只讲到了识别,后面还有分割,检测等论文需要读者们多找找,多读读,更有利于知识框架的建立

下面我们对 摘要进行解读,大致意思是:

虽然变压器架构已成为自然的事实标准语言处理任务，其对计算机视觉的应用仍然有限。在 视觉，注意力要么与卷积网络结合使用，要么用于替换卷积网络的某些组件，同时保持其整体结构到位。我们表明，这种对CNN的依赖是不必要的直接应用于图像补丁序列的纯转换器可以执行在图像分类任务上做得很好。当预先训练大量数据并传输到多个中型或小型图像识别基准（ImageNet，CIFAR-100，VTAB等），Vision Transformer（ViT）获得卓越, 结果与最先进的卷积网络相比，同时需要sub stantially(经验)训练的计算资源更少

总结一下大致是:

CV领域被CNN占据，NLP领域Transformer成为标配，近几年Transformer跨界迁移到CV领域的文章也有很多，大多基于两个思路：

（1）注意力机制与CNN结合；

（2）在整体结构不变的情况下注意力机制替换CNN某些结构；

但是,这些特殊的注意力机制无法实现硬件层面加速，所以本文讲述的是在不依赖CNN结构的情况下，如何尽可能地讲NLP领域的标配——Transformer不做修改的迁移到CV领域。

 然后下图是ViT的结构,如图一所示

 图 1：模型概述。我们将图像拆分为固定大小的补丁，线性嵌入每个补丁，添加位置嵌入，并将生成的向量序列馈送到标准转换器编码器。为了进行分类，我们使用标准方法添加额外的可学习性序列的“classification token”。

“classification token”可以理解为为了做好分类任务而做的收集信息等任务

左图示例中,图形会被分成很多块,图形大小会变化,但是每个块的大小不会变化(在同一个模型中),然后从左到右,从上到下把块拉直,然后进行归一化,再把块中的值进行线性变化映射到这个模型的维度,得到一个patch embedding,然后还需要在开头增加一个可训练的embedding(也是可初始化的embedding),构成新的常用embedding(位置编码),然后送入transformer encoder中,加到多余的未知状态,经过MLP Head(多层感知机) 经过交叉熵完成ViT模型的搭建

---

PyTorch搭建ViT

根据下图我们知道首先需要把一幅图变成embedding

由思维导图可知,有俩种方式

step1 convert image to embedding vector sequence

所以我们分别用俩个函数来实现:

1. naive实现

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)
    print(patch.shape)
    
 

在定义image2emb_naive函数中,我们没有使用for循环,而是使用了f.unfold函数,这个函数简单一点讲就是拿出这个卷积的区

讲到这,相信有些宝子还是迷糊的,那么还是老样子,先附上函数的官方讲解:torch.nn.functional.unfold — PyTorch 2.0 documentation

然后再附上其他博主的详细讲解:PyTorch中torch.nn.functional.unfold函数使用详解_咆哮的阿杰的博客-CSDN博客

为了测试我们所写的函数是否正确,我们接下来定义一些常量 来测试

# test code for image2emb (定义常量)
bs, ic, image_h, image_w = 1, 3, 8, 8 #ic是input channel
patch_size = 4 #4*4为一个patch
model_dim = 8  #在模型汇总,patch_embedding大小跟模型大小是一致的
patch_depth = patch_size * patch_size * ic
image = torch.randn(bs, ic, image_h, image_w) #得到一张图片
weight = torch.randn(patch_depth, model_dim) #patch to embedding的乘法矩阵,是个二维张量,张量的第一维度应该是张量大小
image2emb_naive(image, patch_size, weight)

打印出patch.shape是

 出现1*4*48的原因是(bs,  num_patch,  patch_depth(patch_size*patch_size*ic))

详细解释是1是batch_size,4是因为图片是 8*8的面积,patch_size是4*4,一个8*8的图片经过4*4的处理后就是4块8是patch_size*patch_size*input channel(4*4*3)

同时我们打印出weight.shape是

 所以patch_embedding= patch @ weight(矩阵相乘),即

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)
    print(patch.shape)
    #patch_embedding = patch @ weight
    #return patch_embedding

# test code for image2emb (定义常量)
bs, ic, image_h, image_w = 1, 3, 8, 8 #ic是input channel
patch_size = 4 #4*4为一个patch
model_dim = 8  #在模型汇总,patch_embedding大小跟模型大小是一致的
patch_depth = patch_size * patch_size * ic
image = torch.randn(bs, ic, image_h, image_w) #得到一张图片
weight = torch.randn(patch_depth, model_dim) #patch to embedding的乘法矩阵,是个二维张量,张量的第一维度应该是张量大小
patch_embedding_naive  = image2emb_naive(image, patch_size, weight)
print(patch_embedding_naive.shape)

打印出patch_embedding_naive.shape为

 把3*8*8的图片变成了embedding的形式,每个大小是4*8

上面是naive的版本,接下来我们使用卷积实现该版本

2.卷积实现

def image2emb_conv(image, kernel, stride): #定义卷积三要素:输入,kernel(特征提取器),步长
    conv_output = F.conv2d(image, kernel, stride=stride)  # 大小是:bs*oc*oh*ow(batch_szie*output_channel*output_height*output_weight)
    #一般我们卷积过的宽度和高度会拉成一个序列
    bs, oc, oh, ow = conv_output.shape
    patch_embedding = conv_output.reshape(bs, oc, oh * ow).transpose(-1, -2)#拉直,且把序列长度放中间
    return patch_embedding

接下来我们将定义最关键的kernel,那么kernel该如何定义呢

kernel = weight.transpose(0,1).reshape((-1, ic, patch_size,  patch_size))#形状是oc*ic*kh*kw(output channel*input channel*kernel height*kernel weight
#先使用transpose,将通道数放到前面,然后再进行reshhape操作,然后调用到image2emb_conv

然后使用以下代码即可得到

patch_embedding_conv = image2emb_conv(image, kernel, patch_size)

我们来对比以下naive和conv俩种方法打印出来的

import torch
import torch.nn as nn
import torch.nn.functional as F
 
 
def image2emb_naive(image, patch_size, weight): #patch_size:块的大小
    # image shape: bs * channel * h * w
    patch = F.unfold(image, kernel_size=patch_size, stride=patch_size).transpose(-1, -2)
    #因为可知图像分块没有交叠,所以stride不会为1,stride=kernel_size会分块
    patch_embedding = patch @ weight
    return patch_embedding
 
 
def image2emb_conv(image, kernel, stride): #定义卷积三要素:输入,kernel(特征提取器),步长
    conv_output = F.conv2d(image, kernel, stride=stride)  # 大小是:bs*oc*oh*ow(batch_szie*output_channel*output_height*output_weight)
    #一般我们卷积过的宽度和高度会拉成一个序列
    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 #ic是input channel
patch_size = 4 #4*4为一个patch
model_dim = 8  #在模型汇总,patch_embedding大小跟模型大小是一致的
patch_depth = patch_size * patch_size * ic
image = torch.randn(bs, ic, image_h, image_w) #得到一张图片
# weight = torch.randn(patch_depth, model_dim) #naive使用-->>patch to embedding的乘法矩阵,是个二维张量,张量的第一维度应该是张量大小
weight = torch.randn(patch_depth, model_dim) #卷积使用-->>model_dim是输出通道数目,patch_size是卷积核的面积*输入通道数
patch_embedding_naive  = image2emb_naive(image, patch_size, weight)#---分块方法得到embedding----
kernel = weight.transpose(0,1).reshape((-1, ic, patch_size,  patch_size))#形状是oc*ic*kh*kw(output channel*input channel*kernel height*kernel weight
#先使用transpose,将通道数放到前面,然后再进行reshhape操作,然后调用到image2emb_conv
 
patch_embedding_conv = image2emb_conv(image, kernel, patch_size)#---二维卷积的方法得到embedding---
print(patch_embedding_naive.shape)
print(patch_embedding_conv.shape)

打印出结果为,可以看到基本上一样的

 上述操作得到了图片的embedding

---

step2 prepend CLS token embedding

#step2 prepend CLS token embedding(在模型开头增加一个embedding)
cls_token_embedding = torch.randn(bs, 1, model_dim, requires_grad=True)
token_embedding = torch.cat([cls_token_embedding, patch_embedding_conv], dim=1)

---

step3 add position embedding

首先定义位置数量max_num_token 

max_num_token = 16

然后进行 下述操作 

# step3 add position embedding
position_embedding_table = torch.randn(max_num_token, model_dim, requires_grad=True)
seq_len = token_embedding.shape[1]
position_embedding = torch.tile(position_embedding_table[:seq_len], [token_embedding.shape[0], 1, 1])
token_embedding += position_embedding

---

step4 pass embedding to Transform Encoder

接下来我们进行transformencord,为了让大家了解更清楚有什么用,我找了该代码的官方解释:TransformerEncoder — PyTorch 2.0 documentation

模仿官方代码中的

 写出

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)

---

step5 do classfication

首先定义位置数目num_classes以及实例化一个label

num_classes = 10
label = torch.randint(10, (bs,))

cls_token_output = encoder_output[:, 0, :] #三维为 bs 位置 通道数目,得到一维通道输出
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)

最后得到

 

---

最后给出总体代码

import torch
import torch.nn as nn
import torch.nn.functional as F
 
#step1 convert image to embedding vector sequence
def image2emb_naive(image, patch_size, weight): #patch_size:块的大小
    # image shape: bs * channel * h * w
    patch = F.unfold(image, kernel_size=patch_size, stride=patch_size).transpose(-1, -2)
    #因为可知图像分块没有交叠,所以stride不会为1,stride=kernel_size会分块
    patch_embedding = patch @ weight
    return patch_embedding
 
 
def image2emb_conv(image, kernel, stride): #定义卷积三要素:输入,kernel(特征提取器),步长
    conv_output = F.conv2d(image, kernel, stride=stride)  # 大小是:bs*oc*oh*ow(batch_szie*output_channel*output_height*output_weight)
    #一般我们卷积过的宽度和高度会拉成一个序列
    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 #ic是input channel
patch_size = 4 #4*4为一个patch
model_dim = 8  #在模型汇总,patch_embedding大小跟模型大小是一致的
max_num_token = 16
num_classes = 10
label = torch.randint(10, (bs,))
patch_depth = patch_size * patch_size * ic
image = torch.randn(bs, ic, image_h, image_w) #得到一张图片
# weight = torch.randn(patch_depth, model_dim) #naive使用-->>patch to embedding的乘法矩阵,是个二维张量,张量的第一维度应该是张量大小
weight = torch.randn(patch_depth, model_dim) #卷积使用-->>model_dim是输出通道数目,patch_size是卷积核的面积*输入通道数
patch_embedding_naive  = image2emb_naive(image, patch_size, weight)#---分块方法得到embedding----
kernel = weight.transpose(0,1).reshape((-1, ic, patch_size,  patch_size))#形状是oc*ic*kh*kw(output channel*input channel*kernel height*kernel weight
#先使用transpose,将通道数放到前面,然后再进行reshhape操作,然后调用到image2emb_conv
 
patch_embedding_conv = image2emb_conv(image, kernel, patch_size)#---二维卷积的方法得到embedding---
print(patch_embedding_naive.shape)
print(patch_embedding_conv.shape)
print(patch_embedding_naive)
print(patch_embedding_conv)
 
#step2 prepend CLS token embedding(在模型开头增加一个embedding)
cls_token_embedding = torch.randn(bs, 1, model_dim, requires_grad=True)#增加参数requires_grad,因为是可训练的
token_embedding = torch.cat([cls_token_embedding, patch_embedding_conv], dim=1) #在一位置上(中间维度)去拼接
 
# step3 add position embedding
position_embedding_table = torch.randn(max_num_token, model_dim, requires_grad=True)
seq_len = token_embedding.shape[1]
position_embedding = torch.tile(position_embedding_table[:seq_len], [token_embedding.shape[0], 1, 1])#position_embedding_table[:seq_len]是复制成batch_size的步数,[token_embedding.shape[0]指复制这么多份
token_embedding += position_embedding
 
# step4 pass embedding to Transform 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)
 
# step5 do classfication
cls_token_output = encoder_output[:, 0, :] #三维为 bs 位置 通道数目,得到一维通道输出
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)

ViT结构简单,一般 都是用于图像识别,但是成本很高,需要大量图形,所以可以去了解更多模型

最后推荐记录的视频:28、Vision Transformer(ViT)模型原理及PyTorch逐行实现_哔哩哔哩_bilibili

有疑问的小伙伴可以查看原视频解读