Vision Transformer图像分类(MindSpore实现)
Vision Transformer(ViT)简介
近些年,随着基于自注意(Self-Attention)结构的模型的发展,特别是Transformer模型的提出,极大的促进了自然语言处理模型的发展。由于Transformers的计算效率和可扩展性,它已经能够训练具有超过100B参数的空前规模的模型。
ViT则是自然语言处理和计算机视觉两个领域的融合结晶。在不依赖卷积操作的情况下,依然可以在图像分类任务上达到很好的效果。
模型结构
ViT模型的主体结构是基于Transformer模型的Encoder部分(部分结构顺序有调整,如:normalization的位置与标准Transformer不同),其结构图如下:
模型特点
ViT模型是应用于图像分类领域。因此,其模型结构相较于传统的Transformer有以下几个特点:
- 数据集的原图像被划分为多个patch后,将二维patch(不考虑channel)转换为一维向量,再加上类别向量与位置向量作为模型输入。
- 模型主体的Block基于Transformer的Encoder部分,但是调整了normaliztion的位置,其中,最主要的结构依然是Multi-head Attention结构。
- 模型在Blocks堆叠后接全连接层接受类别向量输出用于分类。通常情况下,我们将最后的全连接层称为Head,Transformer Encoder部分为backbone。
下面将通过代码实例来详细解释基于ViT实现ImageNet分类任务。
环境准备与数据读取
本案例基于MindSpore-GPU版本,在单GPU卡上完成模型训练和验证。
首先导入相关模块,配置相关超参数并读取数据集,该部分代码在Vision套件中都有API可直接调用,详情可以参考以下链接:https://gitee.com/mindspore/vision 。
可通过:http://image-net.org/ 进行数据集下载。
加载前先定义数据集路径,请确保你的数据集路径如以下结构。
.ImageNet/
├── ILSVRC2012_devkit_t12.tar.gz
├── train/
└── val/
from mindspore import context
from mindvision.classification.dataset import ImageNet
context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
data_url = './ImageNet/'
resize = 224
batch_size = 16
dataset_train = ImageNet(data_url,
split="train",
shuffle=True,
resize=resize,
batch_size=batch_size,
repeat_num=1,
num_parallel_workers=1).run()
模型解析
下面将通过代码来细致剖析ViT模型的内部结构。
Transformer基本原理
Transformer模型源于2017年的一篇文章[2]。在这篇文章中提出的基于Attention机制的编码器-解码器型结构在自然语言处理领域获得了巨大的成功。模型结构如下图所示:
其主要结构为多个Encoder和Decoder模块所组成,其中Encoder和Decoder的详细结构如下图所示:
Encoder与Decoder由许多结构组成,如:多头注意力(Multi-Head Attention)层,Feed Forward层,Normaliztion层,甚至残差连接(Residual Connection,图中的“add”)。不过,其中最重要的结构是多头注意力(Multi-Head Attention)结构,该结构基于自注意力(Self-Attention)机制,是多个Self-Attention的并行组成。
所以,理解了Self-Attention就抓住了Transformer的核心。
Attention模块
以下是Self-Attention的解释,其核心内容是为输入向量的每个单词学习一个权重。通过给定一个任务相关的查询向量Query向量,计算Query和各个Key的相似性或者相关性得到注意力分布,即得到每个Key对应Value的权重系数,然后对Value进行加权求和得到最终的Attention数值。
在Self-Attention中:
- 最初的输入向量首先会经过Embedding层映射成Q(Query),K(Key),V(Value)三个向量,由于是并行操作,所以代码中是映射成为dim x 3的向量然后进行分割,换言之,如果你的输入向量为一个向量序列( 1x1 1x1, 2x2 2x2, 3x3 3x3),其中的 1x1 1x1, 2x2 2x2, 3x3 3x3都是一维向量,那么每一个一维向量都会经过Embedding层映射出Q,K,V三个向量,只是Embedding矩阵不同,矩阵参数也是通过学习得到的。这里大家可以认为,Q,K,V三个矩阵是发现向量之间关联信息的一种手段,需要经过学习得到,至于为什么是Q,K,V三个,主要是因为需要两个向量点乘以获得权重,又需要另一个向量来承载权重向加的结果,所以,最少需要3个矩阵,也是论文作者经过实验得出的结论。
2. 自注意力机制的自注意主要体现在它的Q,K,V都来源于其自身,也就是该过程是在提取输入的不同顺序的向量的联系与特征,最终通过不同顺序向量之间的联系紧密性(Q与K乘积经过softmax的结果)来表现出来。Q,K,V得到后就需要获取向量间权重,需要对Q和K进行点乘并除以维度的平方根 ⎯⎯√d ⎯⎯√d,对所有向量的结果进行Softmax处理,通过公式(2)的操作,我们获得了向量之间的关系权重。
3.其最终输出则是通过V这个映射后的向量与QK经过Softmax结果进行weight sum获得,这个过程可以理解为在全局上进行自注意表示。每一组QKV最后都有一个V输出,这是Self-Attention得到的最终结果,是当前向量在结合了它与其他向量关联权重后得到的结果。
通过下图可以整体把握Self-Attention的全部过程。
多头注意力机制就是将原本self-Attention处理的向量分割为多个Head进行处理,这一点也可以从代码中体现,这也是attention结构可以进行并行加速的一个方面。
总结来说,多头注意力机制在保持参数总量不变的情况下,将同样的query, key和value映射到原来的高维空间(Q,K,V)的不同子空间(Q_0,K_0,V_0)中进行自注意力的计算,最后再合并不同子空间中的注意力信息。
所以,对于同一个输入向量,多个注意力机制可以同时对其进行处理,即利用并行计算加速处理过程,又在处理的时候更充分的分析和利用了向量特征。下图展示了多头注意力机制,其并行能力的主要体现在下图中的$a_1$和$a_2$是同一个向量进行分割获得的。
以下是vision套件中的Multi-Head Attention代码,结合上文的解释,代码清晰的展现了这一过程。
import mindspore.nn as nn
class Attention(nn.Cell):
def __init__(self,
dim: int,
num_heads: int = 8,
keep_prob: float = 1.0,
attention_keep_prob: float = 1.0):
super(Attention, self).__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = Tensor(head_dim ** -0.5)
self.qkv = nn.Dense(dim, dim * 3)
self.attn_drop = nn.Dropout(attention_keep_prob)
self.out = nn.Dense(dim, dim)
self.out_drop = nn.Dropout(keep_prob)
self.mul = P.Mul()
self.reshape = P.Reshape()
self.transpose = P.Transpose()
self.unstack = P.Unstack(axis=0)
self.attn_matmul_v = P.BatchMatMul()
self.q_matmul_k = P.BatchMatMul(transpose_b=True)
self.softmax = nn.Softmax(axis=-1)
def construct(self, x):
"""Attention construct."""
b, n, c = x.shape
# 最初的输入向量首先会经过Embedding层映射成Q(Query),K(Key),V(Value)三个向量
# 由于是并行操作,所以代码中是映射成为dim*3的向量然后进行分割
qkv = self.qkv(x)
#多头注意力机制就是将原本self-Attention处理的向量分割为多个Head进行处理
qkv = self.reshape(qkv, (b, n, 3, self.num_heads, c // self.num_heads))
qkv = self.transpose(qkv,