详解视觉Transformers

文摘 2024-06-22 17:44 新加坡

一、什么是Vision Transformers？

从2017年在《Attention is All You Need》中首次提出以来，Transformer模型已经成为自然语言处理（NLP）领域的最新技术。在2021年，论文《An Image is Worth 16x16 Words》成功地将Transformer应用于计算机视觉任务。从那时起，基于Transformer的各种架构陆续被提出用于计算机视觉。

Transformer是一种利用注意力机制作为主要学习机制的机器学习模型，它很快成为了序列到序列任务（如语言翻译）的最新技术。论文《An Image is Worth 16x16 Words》成功地改造了《Attention is All You Need》中中提出的Transformer，用于解决图像分类任务，创建了Vision Transformer（ViT）。ViT基于与《Attention is All You Need》中中的Transformer相同的注意力机制。然而，与用于NLP任务的Transformers包含编码器和解码器两个注意力分支不同（当然现在主流的NLP如GPT只使用Decoder），ViT仅使用编码器。编码器的输出随后传递给一个神经网络“头”，进行预测。

《An Image is Worth 16x16 Words》中ViT的缺点是其最佳性能需要在大规模数据集上进行预训练。表现最好的模型是在专有的JFT-300M数据集上预训练的。在较小的开源数据集ImageNet-21k上预训练的模型，其性能与最新的卷积神经网络ResNet模型相当。《Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet》尝试通过引入一种新的预处理方法，将输入图像转换为一系列token，从而消除这种预训练需求。在本文中，我们将重点介绍《An Image is Worth 16x16 Words》中实现的ViT。

二、模型详解

本文遵循《An Image is Worth 16x16 Words》中描述的模型结构。然而，这篇论文中的代码没有公开。最新的Tokens-to-Token ViT的代码可以在GitHub上找到。Tokens-to-Token ViT（T2T-ViT）模型在一个基础的ViT骨干网络前加上了一个Tokens-to-Token（T2T）模块。本文中的代码基于Tokens-to-Token ViT GitHub代码中的ViT组件。本文所做的修改包括但不限于：允许输入非正方形图像，并去除了dropout层。下面是ViT模型的示意图。

Vision Transformer (ViT) 工作流程（嵌入维度768）：

1. Input Image（输入图像）：

输入图像的尺寸是 60 x 100 像素。

2. Image Tokens（图像Token）：

将输入图像分割成多个 20 x 20 像素的块。图像的高度被分割成 60 / 20 = 3 行块。图像的宽度被分割成 100 / 20 = 5 列块。因此，总的块数为 3 * 5 = 15 个块。每个块被展平成一个长度为 20 * 20 = 400 的向量，所以我们有15个长度为400的图像Token。

3. Linearly Embed Patches（线性嵌入块）：

将每个长度为400的图像Token通过一个线性变换映射到一个嵌入空间，生成长度为768的嵌入向量。768是经过Token化后的每个Token的维度或长度。在Vision Transformer (ViT) 中，每个图像块（Patch）在被展平并通过线性变换后，都会被映射到一个长度为768的嵌入向量。这意味着每个Token的表示是在768维的嵌入空间中。

这样，每个块就被表示为一个长度为768的嵌入向量。

4. Prediction Token（预测Token）：

初始化一个全零的预测Token，长度为768（与每个图像Token的长度相同）。

5. Concatenate（拼接）：

将预测Token与15个图像Token拼接在一起，形成一个包含16个Token的序列。序列的形式为：[预测Token, 图像Token1, 图像Token2, ..., 图像Token15]，每个Token的长度为768。

6. Position Embedding and All Tokens（位置嵌入和所有Token）：

为16个Token（包括预测Token和图像Token）添加位置嵌入。位置嵌入是一个与Token序列长度相同的向量，用于表示每个Token在序列中的位置。位置嵌入的维度也是768。添加位置嵌入后的序列仍然是16个Token，每个Token的长度为768。这一步确保所有Token（包括预测Token和图像Token）都已经添加了位置信息，并且准备好进行下一步的编码块处理。

7. Encoding Block（编码块）：

这些Token序列通过一系列的编码块进行处理。在你的代码示例中，编码块的数量是12（由depth参数指定）。每个编码块包含以下组件：

Norm（规范化层）：对输入进行规范化，确保输入数据在不同维度上的分布更加均匀。
Attention（注意力层）：应用多头自注意力机制来捕捉Token之间的关系。注意力机制允许每个Token与其他Token交换信息，从而捕捉图像的全局特征。
Norm（规范化层）：再次进行规范化，确保在经过注意力层之后的数据继续保持均匀分布。
NN（神经网络）：一个小型的前馈神经网络，由一个全连接层、一个激活层和另一个全连接层组成。前馈神经网络进一步处理和转换Token信息。

经过12个编码块的处理后，每个Token（包括预测Token）都被充分处理和转换，捕捉到了输入图像的丰富特征信息。

8. Seperate（分离）：

从编码块输出的16个Token序列中分离出预测Token和图像Token。此时，预测Token已经通过多个编码块获取了输入图像的信息。分离操作通常是通过索引来实现的，例如：

# 从编码块输出的Token序列中提取预测Token  prediction_token = encoded_tokens[:, 0]  # 提取图像Token  image_tokens = encoded_tokens[:, 1:]

通过这种方式，我们可以得到一个包含预测信息的预测Token和15个图像Token。

9. Prediction Token（预测Token）：

提取后的预测Token，已经从编码块中获取了输入图像的信息。这个Token包含了用于最终预测的所有必要信息。由于预测Token在多个编码块中与其他图像Token进行了交互，它已经聚合了整个图像的全局特征。

10. Linear Layer（线性层）：

预测Token通过一个线性层来生成最终的预测结果。这个线性层将预测Token映射到输出空间，输出形状根据具体的学习任务进行设置。例如：

对于分类任务，输出是一个长度为类数的向量，每个元素表示对应类别的概率。

对于回归任务，输出是一个或多个预测值。

# 假设我们有一个线性层用于分类任务  linear_layer = nn.Linear(768, num_classes)  prediction = linear_layer(prediction_token)

11. Prediction（预测结果）：

最终的预测结果。例如，在图像分类任务中，预测结果可以是一个包含所有类别概率的向量，然后选择概率最高的类别作为最终分类结果。

# 对于分类任务，选择概率最高的类别作为最终预测结果  predicted_class = torch.argmax(prediction, dim=-1)

解释

图像Token：每个图像块被展平成一个Token，然后通过线性变换映射到嵌入空间，生成长度为768的嵌入向量。
预测Token：初始时是全零，但在经过多个编码块的处理后，它会逐渐聚合图像中的信息，最终用于预测。
注意力机制：在编码块中，注意力机制允许每个Token（包括预测Token）与其他Token交互，从而捕捉图像的全局特征。
嵌入维度768：所有Token（包括图像Token和预测Token）在嵌入空间中的表示长度都是768，这使得它们能够在编码块中进行有效的交互和处理。

通过这个过程，ViT模型不仅提取图像的局部特征（通过图像Token），还通过预测Token聚合这些特征并用于最终的预测。

那为什么图像token是15个而预测token只是一个？预测token中最终会包含什么信息？

让我们详细解释为什么图像Token有15个，而预测Token只有一个，以及预测Token最终会包含什么信息。

1. 图像Token的数量

图像Token的数量是由图像分割成的块（patch）的数量决定的。在你的例子中：

输入图像的尺寸是 60 x 100 像素。
每个块的尺寸是 20 x 20 像素。
图像的高度被分割成 60 / 20 = 3 行块。
图像的宽度被分割成 100 / 20 = 5 列块。
因此，总的块数是 3 * 5 = 15 个块。

每个块被展平成一个长度为 20 * 20 = 400 的向量，所以我们有15个长度为400的图像Token。

2. 预测Token的数量

预测Token只有一个，这是设计ViT模型的一部分。这个单一的预测Token初始时是全零的，长度为400（与每个图像Token的长度相同）。

3. 预测Token的作用

预测Token的主要作用是通过注意力机制和编码块来聚合所有图像Token中的信息，最终用于生成预测结果。具体来说：

初始阶段：预测Token是一个全零的向量，长度为400。
编码块处理：在每个编码块中，预测Token与所有图像Token进行交互（通过注意力机制）。注意力机制允许每个Token与其他Token交换信息。
信息聚合：通过多个编码块的处理，预测Token逐渐聚合了所有图像Token的信息。这意味着预测Token在多个编码块处理之后，包含了整个图像的全局特征。

4. 预测Token最终包含的信息

经过多个编码块的处理后，预测Token会包含以下信息：

全局特征：预测Token聚合了所有图像Token的信息，表示整个图像的全局特征。
上下文信息：通过注意力机制，预测Token能够捕捉到图像的上下文信息和不同部分之间的关系。

5. 预测Token如何用于最终预测

最终的预测步骤如下：

提取预测Token：从编码块输出的Token序列中提取预测Token。
线性层：预测Token通过一个线性层生成最终的预测结果。这个线性层的输出形状根据具体的学习任务设置。例如：

对于分类任务，输出是一个长度为类数的向量。
对于回归任务，输出是一个或多个预测值。

预测结果：最终的预测结果。例如，在图像分类任务中，预测结果可以是一个包含所有类别概率的向量，然后选择概率最高的类别作为最终分类结果。

总结起来，单一的预测Token通过多个编码块的处理，聚合了所有图像Token的信息，包含了整个图像的全局特征和上下文信息。这个预测Token最终用于生成模型的预测结果。

接下来，我们详细介绍上述关键步骤以及代码实现。

三、图像标记化 Image Tokenization

ViT的第一步是从输入图像创建tokens。Transformers在一系列tokens上操作；在NLP中，这通常是一句话中的单词。对于计算机视觉来说，ViT将图像转换为tokens，使得每个token代表图像的一个局部区域——或称为patch。他们描述了将高度为H、宽度为W、通道数为C的图像重新排列为N个大小为P的patch tokens：

每个token的长度为 P²∗C。

让我们来看一个patch tokenization的示例，这里使用的是Luis Zuno (@ansimuz)的像素艺术作品《Mountain at Dusk》。原始作品已经被裁剪并转换为单通道图像。这意味着每个像素的值在零到一之间。单通道图像通常以灰度显示；然而，我们将以紫色配色方案显示它，因为这样更容易看清。

这幅图像的高度 H=60，宽度 W=100。我们将 P 设置为 20，因为它能够整除 H 和 W。

表示“patch size”，即图像分割成小块（patch）的尺寸。在Vision Transformer (ViT)模型中，图像被分割成大小相同的patch，然后每个patch被转换成一个token。每个token将被输入到Transformer模型中进行处理。

具体来说：

H：图像的高度。
W：图像的宽度。
P：patch的尺寸，即每个patch的高度和宽度。

当你将图像分割成大小为P×P的patch时，图像将被划分成多个小块，每个小块都是一个P×P的子图像。对于每个patch，都可以将其展平并转换为一个向量，这个向量就是一个token。然后，这些token将作为输入传递给Transformer模型。

在这个示例中，图像的高度H是60，宽度W是100。选择P=20，因为20能够整除60和100，这样可以确保图像被均匀地分割成不重叠的patch。具体来说，图像将被分割成（60 / 20）×（100 / 20）= 3 × 5 = 15个patch。

通过将这些小块展平，我们可以看到生成的tokens。我们以第12个小块为例进行观察，因为它包含了四种不同的色调。

每个 patch将产生一个长度为 400 的令牌。

从图像中提取tokens之后，通常会使用线性投影来改变这些tokens的长度。这通过一个可学习的线性层来实现。tokens的新长度被称为潜在维度、通道维度或者token长度。经过投影后，这些tokens不再能从视觉上识别为原始图像中的某个小块。现在我们已经理解了这个概念，我们可以来看一下在代码中如何实现小块token化。

在代码中，Patch Tokenization（块标记化）是通过一个自定义的PyTorch模块实现的，该模块首先将图像分割成小块，然后将每个小块展平并通过线性变换映射到指定的嵌入空间。让我们详细解析一下代码中是如何实现Patch Tokenization的。

请注意，这里的两个断言语句确保了图像的尺寸可以被小块的大小整除。实际将图像分割成小块的操作是通过 torch.nn.Unfold 层实现的。

我们将使用裁剪后的、单通道版本的《Mountain at Dusk》运行这个代码示例。我们应该能看到上面提到的token数量和初始token大小的值。我们将使用 token_len=768 作为投影后的长度，这是ViT基础版本的尺寸。在Vision Transformer (ViT) 模型中，出现的数字768通常是指Token的嵌入维度（embedding dimension）。

下面代码块的第一行将《Mountain at Dusk》从NumPy数组转换为Torch张量。我们还需要对张量进行 unsqueeze 操作，以创建一个通道维度和一个批次大小维度。如上所述，我们只有一个通道。由于只有一张图像，批次大小 batchsize 为1。

x = torch.from_numpy(mountains).unsqueeze(0).unsqueeze(0).to(torch.float32)token_len = 768print('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 Modulepatch_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)

现在，我们将把图像分割成tokens。

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])

正如我们在示例中看到的那样，每个长度为400的token有N=15个。最后，我们将这些tokens投影到指定的 token_len。

x = patch_tokens.project(x)print('After projection, dimensions are\n\tbatchsize:', x.shape[0], '\n\tnumber of tokens:', x.shape[1], '\n\ttoken length:', x.shape[2])

现在我们已经得到了tokens，接下来我们可以继续进行ViT的处理。

四、Token处理 Token Processing

接下来我们使用另外的例子，而不再是上面《Mountain at Dusk》的例子。

在编码块之前的ViT的接下来的两个步骤，我们称之为“Token Processing”。下图展示了ViT图示中的Token Processing部分。

第一步是在图像tokens前添加一个空白token，称为预测token。这个token将在编码块的输出中用于做出预测。它一开始是空白的，这样它可以从其他图像tokens中获取信息。

我们将从175个tokens开始。每个token的长度为768，这是ViT基础版本的尺寸。我们使用批次大小为13，因为它是质数，不会与其他参数混淆。

# Define an Inputnum_tokens = 175token_len = 768batch = 13x = 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])
# Append a Prediction Tokenpred_token = torch.zeros(1, 1, token_len).expand(batch, -1, -1)print('Prediction Token dimensions are\n\tbatchsize:', pred_token.shape[0], '\n\tnumber of tokens:', pred_token.shape[1], '\n\ttoken length:', pred_token.shape[2])
x = torch.cat((pred_token, x), dim=1)print('Dimensions with Prediction Token are\n\tbatchsize:', x.shape[0], '\n\tnumber of tokens:', x.shape[1], '\n\ttoken length:', x.shape[2])

现在，我们为我们的tokens添加位置嵌入。位置嵌入允许transformer理解图像tokens的顺序。请注意，这里是相加操作，而不是连接操作。位置嵌入的具体细节是一个较为复杂的话题，留待以后再详细讨论。

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 + PEprint('Dimensions with Position Embedding are\n\tbatchsize:', x.shape[0], '\n\tnumber of tokens:', x.shape[1], '\n\ttoken length:', x.shape[2])

现在，我们的tokens已经准备好进入编码块。

五、编码块 Encoding Block Prerequisite Code

编码块是模型实际从图像tokens中学习的地方。编码块的数量是由用户设置的超参数。下面是编码块的示意图。

编码块的代码如下所示。

classEncoding(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)
    defforward(self, x):        x = x + self.attn(self.norm1(x))        x = x + self.neuralnet(self.norm2(x))        return x

num_heads、qkv_bias 和 qk_scale 参数定义了注意力模块的组件。

hidden_chan_mul 和 act_layer 参数定义了神经网络模块的组件。激活层可以是任何 torch.nn.modules.activation 层。

norm_layer 可以从任何 torch.nn.modules.normalization 层中选择。

现在我们将逐步解析图示中的每个蓝色块及其对应的代码。我们将使用长度为768的176个tokens。我们将使用批次大小为13，因为它是质数，不会与其他参数混淆。我们将使用4个注意力头，因为它可以整除token的长度；不过，在编码块中你不会看到注意力头的维度。

# Define an Inputnum_tokens = 176token_len = 768batch = 13heads = 4x = 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 ModuleE = 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

现在，我们将通过一个规范化层和一个注意力模块。编码块中的注意力模块经过参数化处理，因此不会改变token的长度。在注意力模块之后，我们实现了第一个分离连接。

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 + xprint('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 + yprint('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

这就是一个编码块的全部内容。由于最终的维度与初始维度相同，模型可以轻松地将tokens传递通过多个编码块，这由深度超参数设置。

六、神经网络模块

神经网络（NN）模块是编码块的一个子组件。NN模块非常简单，由一个全连接层、一个激活层和另一个全连接层组成。激活层可以是任何 torch.nn.modules.activation 层，并作为输入传递给模块。NN模块可以配置为改变输入的形状，或者保持相同的形状。我们不会逐步解析这段代码，因为神经网络在机器学习中非常常见，并不是本文的重点。然而，下面是NN模块的代码。

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

预测处理

在通过编码块之后，模型必须做的最后一件事就是进行预测。下面展示的是ViT图中的“预测处理”组件。

我们将查看此过程的每一步。我们将继续使用长度为768的176个tokens。为了说明单个预测是如何进行的，我们将使用批次大小为1。批次大小大于1时，则会并行计算这些预测。

# Define an Inputnum_tokens = 176token_len = 768batch = 1x = 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

首先，所有的token都通过一个norm layer.。

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

接下来，我们将预测token与其余的tokens分离。在整个编码块过程中，预测token变得非零并获取了关于输入图像的信息。我们将仅使用这个预测token来进行最终预测。

pred_token = x[:, 0]print('Length of prediction token:', pred_token.shape[-1])

Length of prediction token: 768

最后，预测token会通过head来进行预测。head通常是某种神经网络，其形式取决于具体模型。在《An Image is Worth 16x16 Words》中，他们在预训练期间使用了一个带有隐藏层的多层感知器（MLP），而在微调期间使用了一个单一的线性层。在《Tokens-to-Token ViT》³中，他们使用了一个单一的线性层作为head。这个示例中也将使用一个单一的线性层。

需要注意的是，head的输出形状是根据学习问题的参数设置的。对于分类任务，输出通常是一个长度为类数的one-hot编码向量。对于回归任务，输出则是任意整数个预测参数。在这个示例中，将使用输出形状为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

参考文章：

https://towardsdatascience.com/vision-transformers-explained-a9d07147e4c8

https://github.com/lanl/vision_transformers_explained/blob/main/notebooks/VisionTransformersExplained.ipynb

http://mp.weixin.qq.com/s?__biz=MzAwMDc2NjQ4Nw==&mid=2663561046&idx=1&sn=2d34816f6743061b9ecb3fbf71034f73

大魏分享

https://github.com/davidsajare/david-share.git