自2017年引入“Attention is All You Need”¹以来,Transformer模型已经成为自然语言处理(NLP)领域必不可少的前沿技术了。
到了2021年,“An Image is Worth 16x16 Words”²成功地将Transformer模型应用于计算机视觉任务。
自此之后,计算机视觉领域出现了多种基于Transformer的架构。
本文深入解析了“An Image is Worth 16x16 Words”²中阐述的视觉Transformer(ViT)模型。
内容不仅包括ViT的开源代码,还对其各个组件进行了概念性解释。
所有代码均使用PyTorch Python包编写。
1.视觉Transformer(ViT)简介
在“Attention is All You Need”¹一文中,Transformer作为一种利用注意力机制作为主要学习方式的机器学习模型被首次引入。
这一模型迅速成为序列到序列任务(如语言翻译)中的顶尖技术。
“An Image is Worth 16x16 Words”²成功地将[1]中提出的Transformer模型进行改造,用于解决图像分类任务,从而诞生了视觉Transformer(Vision Transformer,简称ViT)。
ViT同样基于[1]中Transformer的注意力机制,但有所不同的是,自然语言处理(NLP)任务中的Transformer包含编码器注意力分支和解码器注意力分支,而ViT仅使用编码器。编码器的输出随后被传递给一个神经网络“头”,用于做出预测。
然而,[2]中实现的ViT的一个缺点是,要达到最佳性能,它需要在大型数据集上进行预训练。
性能最优的模型是在专有的JFT-300M数据集上进行预训练的。
而在较小且开源的ImageNet-21k数据集上进行预训练的模型,其性能与最先进的卷积ResNet模型相当。
为了克服这一预训练需求,“Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet”³通过引入一种新颖的预处理方法来将输入图像转换为一系列标记(tokens),从而尝试从头开始训练ViT。
关于这一方法的更多信息,可以在此处找到。但本文将重点介绍[2]中实现的ViT。
2.模型解析
本文遵循了“An Image is Worth 16x16 Words”²中概述的模型结构。
然而,该论文的代码并未公开,不过,更新版本的“Tokens-to-Token ViT”³的代码已在GitHub上发布。
Tokens-to-Token ViT(T2T-ViT)模型在标准的ViT主干前增加了一个Tokens-to-Token(T2T)模块。
本文中的代码基于GitHub上Tokens-to-Token ViT³中的ViT组件编写。
针对本文所做的修改包括但不限于,允许非方形输入图像以及移除dropout层。
下面是ViT模型的示意图:
3.图像标记化
ViT的第一步是将输入图像转换为标记(tokens)。
Transformer模型在标记序列上操作;在自然语言处理(NLP)中,这通常是一个由单词组成的句子。
然而,在计算机视觉领域,如何将输入分割成标记并不那么直观。
ViT通过将图像转换为标记来实现,每个标记代表图像的一个局部区域——或称为块(patch)。
这涉及到将高度为H、宽度为W、通道数为C的图像重新塑形成N个大小为P的块标记:
让我们以Luis Zuno (@ansimuz)的像素艺术《黄昏之山》为例,来演示块标记化的过程。
原始作品已被裁剪并转换为单通道图像,这意味着每个像素的值在0到1之间。
单通道图像通常以灰度显示,但为了更清晰地展示,我们将使用紫色色调。
请注意,块标记化的代码并未包含在[3]的关联代码中。
mountains = np.load(os.path.join(figure_path, 'mountains.npy'))
H = mountains.shape[0]
W = mountains.shape[1]
print('Mountain at Dusk is H =', H, 'and W =', W, 'pixels.')
print('\n')
fig = plt.figure(figsize=(10,6))
plt.imshow(mountains, cmap='Purples_r')
plt.xticks(np.arange(-0.5, W+1, 10), labels=np.arange(0, W+1, 10))
plt.yticks(np.arange(-0.5, H+1, 10), labels=np.arange(0, H+1, 10))
plt.clim([0,1])
cbar_ax = fig.add_axes([0.95, .11, 0.05, 0.77])
plt.clim([0, 1])
plt.colorbar(cax=cbar_ax);
#plt.savefig(os.path.join(figure_path, 'mountains.png'))
Mountain at Dusk is H = 60 and W = 100 pixels.
这张图像的高度H=60,宽度W=100。我们将设置P=20,因为它可以整除H和W。
P = 20
N = int((H*W)/(P**2))
print('There will be', N, 'patches, each', P, 'by', str(P)+'.')
print('\n')
fig = plt.figure(figsize=(10,6))
plt.imshow(mountains, cmap='Purples_r')
plt.hlines(np.arange(P, H, P)-0.5, -0.5, W-0.5, color='w')
plt.vlines(np.arange(P, W, P)-0.5, -0.5, H-0.5, color='w')
plt.xticks(np.arange(-0.5, W+1, 10), labels=np.arange(0, W+1, 10))
plt.yticks(np.arange(-0.5, H+1, 10), labels=np.arange(0, H+1, 10))
x_text = np.tile(np.arange(9.5, W, P), 3)
y_text = np.repeat(np.arange(9.5, H, P), 5)
for i in range(1, N+1):
plt.text(x_text[i-1], y_text[i-1], str(i), color='w', fontsize='xx-large', ha='center')
plt.text(x_text[2], y_text[2], str(3), color='k', fontsize='xx-large', ha='center');
#plt.savefig(os.path.join(figure_path, 'mountain_patches.png'), bbox_inches='tight'
通过将这些块展平,我们得到了结果标记。以第12个块为例,因为它包含四种不同的色调。
print('Each patch will make a token of length', str(P**2)+'.')
print('\n')
patch12 = mountains[40:60, 20:40]
token12 = patch12.reshape(1, P**2)
fig = plt.figure(figsize=(10,1))
plt.imshow(token12, aspect=10, cmap='Purples_r')
plt.clim([0,1])
plt.xticks(np.arange(-0.5, 401, 50), labels=np.arange(0, 401, 50))
plt.yticks([]);
#plt.savefig(os.path.join(figure_path, 'mountain_token12.png'), bbox_inches='tight')
Each patch will make a token of length 400.
从图像中提取标记后,通常会使用线性投影来改变标记的长度。这通过可学习的线性层来实现。
标记的新长度被称为潜在维度²、通道维度³或标记长度。
投影后,这些标记不再能从视觉上识别为原始图像中的块。
现在,我们理解了这一概念,接下来我们将查看如何在代码中实现块标记化。
class Patch_Tokenization(nn.Module):
def __init__(self,
img_size: tuple[int, int, int]=(1, 1, 60, 100),
patch_size: int=50,
token_len: int=768):
""" Patch Tokenization Module
Args:
img_size (tuple[int, int, int]): size of input (channels, height, width)
patch_size (int): the side length of a square patch
token_len (int): desired length of an output token
"""
super().__init__()
## Defining Parameters
self.img_size = img_size
C, H, W = self.img_size
self.patch_size = patch_size
self.token_len = token_len
assert H % self.patch_size == 0, 'Height of image must be evenly divisible by patch size.'
assert W % self.patch_size == 0, 'Width of image must be evenly divisible by patch size.'
self.num_tokens = (H / self.patch_size) * (W / self.patch_size)
## Defining Layers
self.split = nn.Unfold(kernel_size=self.patch_size, stride=self.patch_size, padding=0)
self.project = nn.Linear((self.patch_size**2)*C, token_len)
def forward(self, x):
x = self.split(x).transpose(1,0)
x = self.project(x)
return x
请注意,这里有两个断言语句,用于确保图像尺寸能够被块大小整除。
实际的块分割是通过torch.nn.Unfold⁵层实现的。
我们将使用裁剪后的单通道版本《黄昏之山》⁴来运行此代码示例。
我们应该能看到标记数量和初始标记大小的值,就像我们之前看到的那样。
我们将使用token_len=768作为投影后的长度,这是ViT基础变体的尺寸²。
以下代码块的第一行是将《黄昏之山》⁴的数据类型从NumPy数组更改为Torch张量。
我们还需要使用unsqueeze⁶方法来为张量创建一个通道维度和一个批量大小维度。
如上所述,我们有一个通道。由于只有一张图像,所以批量大小为1。
x = torch.from_numpy(mountains).unsqueeze(0).unsqueeze(0).to(torch.float32)
token_len = 768
print('Input dimensions are\n\tbatchsize:', x.shape[0], '\n\tnumber of input channels:', x.shape[1], '\n\timage size:', (x.shape[2], x.shape[3]))
# Define the Module
patch_tokens = Patch_Tokenization(img_size=(x.shape[1], x.shape[2], x.shape[3]),
patch_size = P,
token_len = token_len)
Input dimensions are
batchsize: 1
number of input channels: 1
image size: (60, 100)Input dimensions are
batchsize: 1
number of input channels: 1
image size: (60, 100)
现在,我们将把图像分割成标记。
x = patch_tokens.split(x).transpose(2,1)
print('After patch tokenization, dimensions are\n\tbatchsize:', x.shape[0], '\n\tnumber of tokens:', x.shape[1], '\n\ttoken length:', x.shape[2])
After patch tokenization, dimensions are
batchsize: 1
number of tokens: 15
token length: 400
正如我们在示例中看到的,有N=15个标记,每个标记的长度为400。
最后,我们将标记投影到token_len的长度。
def get_sinusoid_encoding(num_tokens, token_len):
""" Make Sinusoid Encoding Table
Args:
num_tokens (int): number of tokens
token_len (int): length of a token
Returns:
(torch.FloatTensor) sinusoidal position encoding table
"""
def get_position_angle_vec(i):
return [i / np.power(10000, 2 * (j // 2) / token_len) for j in range(token_len)]
sinusoid_table = np.array([get_position_angle_vec(i) for i in range(num_tokens)])
sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])
sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])
return torch.FloatTensor(sinusoid_table).unsqueeze(0)
PE = get_sinusoid_encoding(num_tokens+1, token_len)
print('Position embedding dimensions are\n\tnumber of tokens:', PE.shape[1], '\n\ttoken length:', PE.shape[2])
x = x + PE
print('Dimensions with Position Embedding are\n\tbatchsize:', x.shape[0], '\n\tnumber of tokens:', x.shape[1], '\n\ttoken length:', x.shape[2])
现在,我们的标记已经准备好进入编码块了。
4.编码块
编码块是模型实际从图像标记中学习的地方,编码块的数量是一个由用户设置的超参数。以下是编码块的图示。
下面是编码块的代码。
class Encoding(nn.Module):
def __init__(self,
dim: int,
num_heads: int=1,
hidden_chan_mul: float=4.,
qkv_bias: bool=False,
qk_scale: NoneFloat=None,
act_layer=nn.GELU,
norm_layer=nn.LayerNorm):
""" Encoding Block
Args:
dim (int): size of a single token
num_heads(int): number of attention heads in MSA
hidden_chan_mul (float): multiplier to determine the number of hidden channels (features) in the NeuralNet component
qkv_bias (bool): determines if the qkv layer learns an addative bias
qk_scale (NoneFloat): value to scale the queries and keys by;
if None, queries and keys are scaled by ``head_dim ** -0.5``
act_layer(nn.modules.activation): torch neural network layer class to use as activation
norm_layer(nn.modules.normalization): torch neural network layer class to use as normalization
"""
super().__init__()
## Define Layers
self.norm1 = norm_layer(dim)
self.attn = Attention(dim=dim,
chan=dim,
num_heads=num_heads,
qkv_bias=qkv_bias,
qk_scale=qk_scale)
self.norm2 = norm_layer(dim)
self.neuralnet = NeuralNet(in_chan=dim,
hidden_chan=int(dim*hidden_chan_mul),
out_chan=dim,
act_layer=act_layer)
def forward(self, x):
x = x + self.attn(self.norm1(x))
x = x + self.neuralnet(self.norm2(x))
return x
num_heads、qkv_bias和qk_scale参数定义了Attention模块的组件。
hidden_chan_mul和act_layer参数定义了神经网络模块的组件。
激活层可以是任何torch.nn.modules.activation⁷层。我们稍后会更详细地讨论神经网络模块。
norm_layer可以从任何torch.nn.modules.normalization⁸层中选择。
现在,我们将逐一讲解图示中的每个蓝色块及其对应的代码。
我们将使用176个长度为768的标记,由于13是一个质数,不会与其他任何参数混淆,我们将批量大小设置为13。
我们将使用4个注意力头,因为它们可以均匀地划分标记长度;
然而,在编码块中你不会看到注意力头的维度。
# Define an Input
num_tokens = 176
token_len = 768
batch = 13
heads = 4
x = torch.rand(batch, num_tokens, token_len)
print('Input dimensions are\n\tbatchsize:', x.shape[0], '\n\tnumber of tokens:', x.shape[1], '\n\ttoken length:', x.shape[2])
# Define the Module
E = Encoding(dim=token_len, num_heads=heads, hidden_chan_mul=1.5, qkv_bias=False, qk_scale=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm)
E.eval();
Input dimensions are
batchsize: 13
number of tokens: 176
token length: 768
接下来,我们将通过一个归一化层和一个Attention模块。编码块中的Attention模块是参数化的,以便它不会改变标记的长度。
在Attention模块之后,我们实现了第一个残差连接。
y = E.norm1(x)
print('After norm, dimensions are\n\tbatchsize:', y.shape[0], '\n\tnumber of tokens:', y.shape[1], '\n\ttoken size:', y.shape[2])
y = E.attn(y)
print('After attention, dimensions are\n\tbatchsize:', y.shape[0], '\n\tnumber of tokens:', y.shape[1], '\n\ttoken size:', y.shape[2])
y = y + x
print('After split connection, dimensions are\n\tbatchsize:', y.shape[0], '\n\tnumber of tokens:', y.shape[1], '\n\ttoken size:', y.shape[2])
After norm, dimensions are
batchsize: 13
number of tokens: 176
token size: 768
After attention, dimensions are
batchsize: 13
number of tokens: 176
token size: 768
After split connection, dimensions are
batchsize: 13
number of tokens: 176
token size: 768
然后,我们再通过一个归一化层,接着是神经网络模块。最后,我们完成第二个残差连接。
z = E.norm2(y)
print('After norm, dimensions are\n\tbatchsize:', z.shape[0], '\n\tnumber of tokens:', z.shape[1], '\n\ttoken size:', z.shape[2])
z = E.neuralnet(z)
print('After neural net, dimensions are\n\tbatchsize:', z.shape[0], '\n\tnumber of tokens:', z.shape[1], '\n\ttoken size:', z.shape[2])
z = z + y
print('After split connection, dimensions are\n\tbatchsize:', z.shape[0], '\n\tnumber of tokens:', z.shape[1], '\n\ttoken size:', z.shape[2])
After norm, dimensions are
batchsize: 13
number of tokens: 176
token size: 768
After neural net, dimensions are
batchsize: 13
number of tokens: 176
token size: 768
After split connection, dimensions are
batchsize: 13
number of tokens: 176
token size: 768
以上就是一个编码块的全部内容!由于最终维度与初始维度相同,因此模型可以轻松地根据深度超参数的设置,将标记传递通过多个编码块。
5.神经网络模块
神经网络(NN)模块是编码块的一个子组件,它结构非常简洁,由全连接层、激活层以及另一个全连接层组成。
激活层可以是任何torch.nn.modules.activation⁷层,作为模块的输入传递。
神经网络模块可以根据需要配置为改变输入的形状或保持其形状不变。
由于神经网络在机器学习中非常普遍,且不是本文的重点,因此我们将不深入代码细节。
以下是神经网络模块的代码示例。
class NeuralNet(nn.Module):
def __init__(self,
in_chan: int,
hidden_chan: NoneFloat=None,
out_chan: NoneFloat=None,
act_layer = nn.GELU):
""" Neural Network Module
Args:
in_chan (int): number of channels (features) at input
hidden_chan (NoneFloat): number of channels (features) in the hidden layer;
if None, number of channels in hidden layer is the same as the number of input channels
out_chan (NoneFloat): number of channels (features) at output;
if None, number of output channels is same as the number of input channels
act_layer(nn.modules.activation): torch neural network layer class to use as activation
"""
super().__init__()
## Define Number of Channels
hidden_chan = hidden_chan or in_chan
out_chan = out_chan or in_chan
## Define Layers
self.fc1 = nn.Linear(in_chan, hidden_chan)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_chan, out_chan)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.fc2(x)
return x
6.简预测处理
在通过编码块之后,模型需要做的最后一件事就是进行预测。
ViT(Vision Transformer)图示中的“预测处理”组件如下所示。
我们将逐一探讨这一过程的每个步骤,继续以176个长度为768的标记为例,我们将使用批量大小为1来演示如何生成单个预测。
如果批量大小大于1,则表示并行计算这些预测。
# Define an Input
num_tokens = 176
token_len = 768
batch = 1
x = torch.rand(batch, num_tokens, token_len)
print('Input dimensions are\n\tbatchsize:', x.shape[0], '\n\tnumber of tokens:', x.shape[1], '\n\ttoken length:', x.shape[2])
Input dimensions are
batchsize: 1
number of tokens: 176
token length: 768
首先,所有标记都会通过一个归一化层。
norm = nn.LayerNorm(token_len)
x = norm(x)
print('After norm, dimensions are\n\tbatchsize:', x.shape[0], '\n\tnumber of tokens:', x.shape[1], '\n\ttoken size:', x.shape[2])
After norm, dimensions are
batchsize: 1
number of tokens: 1001
token size: 768
接下来,我们从剩余的标记中分离出预测标记,在编码块的处理过程中,预测标记已经变得非零,并获得了关于输入图像的信息。
pred_token = x[:, 0]
print('Length of prediction token:', pred_token.shape[-1])
我们将仅使用此预测标记来做出最终预测。
Length of prediction token: 768
最后预测标记通过“头”(head)部分来生成预测结果,
“头”部分通常是某种形式的神经网络,具体取决于模型设计。
在《An Image is Worth 16x16 Words²》中,他们在预训练阶段使用了一个带有隐藏层的多层感知机(MLP),而在微调阶段则使用了一个单线性层。
而在《Tokens-to-Token ViT³》中,他们则使用了单线性层作为“头”,本例将遵循使用单线性层的做法。
请注意,头的输出形状是根据学习问题的参数来设置的。
对于分类问题,它通常是一个长度等于类别数的一维向量,采用独热编码方式。对于回归问题,它将是任意数量的预测参数(整数)。本例中,我们将使用长度为1的输出形状来代表单一的回归估计值。
head = nn.Linear(token_len, 1)
pred = head(pred_token)
print('Length of prediction:', (pred.shape[0], pred.shape[1]))
print('Prediction:', float(pred))
Length of prediction: (1, 1)
Prediction: -0.5474240779876709
到这里,模型已经完成了预测!