跳至主要内容

视觉 Transformer(ViT)教程:架构与代码示例

了解视觉 Transformer(ViT)如何利用补丁嵌入和自注意力,在现代图像分类中胜过 CNN。本深入教程拆解 ViT 架构,提供分步 Python 代码,并讲解何时在真实计算机视觉项目中选择 ViT。
· 15分钟 阅读

使用 AI 探索

我们都知道,Transformer 已成为自然语言处理的基石,如今也在计算机视觉领域迅速发展。视觉 Transformer(ViT)通过采用与文本模型相同的原理,为图像理解提供了一种全新的思路;大型语言模型就是一个典型代表。

ViT 最早由谷歌研究院于 2020 年提出,论文题为“An Image is Worth 16×16 Words”,展示了在大型图像数据集上训练的 Transformer 可以媲美或超越 CNN。此后,诸如DeiT 和 Swin 等实用变体已将其应用拓展到众多视觉任务。

在本教程中,我将带您了解视觉 Transformer(原始/vanilla 版本)的核心概念、架构组件与实际应用。最后,您不仅能理解 ViT 的工作原理,还能在自己的计算机视觉项目中实现它。我也建议您查看我关于多头注意力的另一篇指南。  

什么是视觉 Transformer(ViT)?

先从一个温和的方式介绍视觉 Transformer。 

这些模型将图像视为一个序列数据,类似语言模型处理句子中的词语。图像会被划分为固定大小的补丁,每个补丁被展平并嵌入为一个向量。 

随后,这些向量被送入标准的 Transformer 编码器,通过注意力机制在全局范围内建模补丁之间的关系。 

与卷积神经网络的比较

这与卷积神经网络(CNN)的根本思路不同,后者侧重局部特征,并依赖分层特征提取。

既然提到了 CNN,我们就来谈谈它们与 ViT 的差异:

特性

卷积神经网络(CNN)

视觉 Transformer(ViT)

架构

使用卷积 + 池化层

使用带自注意力的 Transformer 编码器

感受野

局部感受野,随深度扩大

从一开始就是全局

位置信息感知

通过结构隐式获得

通过位置编码显式注入

数据需求

较小数据集也能有效

需要大规模数据集

归纳偏置

强

较弱但更灵活

并行潜力

中等

非常高

模型可扩展性

向大模型扩展性较差

高度可扩展

可解释性

可通过池化层分析,但仍有难度

可借助注意力可视化增强,但也较困难

显然我们尚未全面覆盖视觉 Transformer,但先梳理一些起初可能不直观的要点是很有帮助的。 

例如,归纳偏置和感受野是很重要但您可能未曾听过的术语。 

归纳偏置解释

归纳偏置是指模型为帮助其从数据中学习而做出的假设。这个概念看似抽象,但通过下面的类比就容易理解了。

想象我们一起拼一幅拼图。我们会认为拼图块基于形状和颜色能够拼接在一起。这个假设能帮助我们更快地完成拼图。这就像 CNN 的工作方式。 

它通过假设图像中相邻区域相关来进行“聪明的猜测”,就像彼此靠近的拼图块往往属于图像的同一部分。它还假设将一个形状移动位置后,它仍是同一形状。这些有益的“先验”被称为强归纳偏置。

再想象我们拿到一大堆拼图块,但完全不知道成品图像是什么样,也不知道该看形状、颜色,还是别的什么。于是我们开始尝试一切方法,慢慢摸索出可行的策略。

这就是Transformer 模型的做法。它不做强假设,而是通过大量数据来学习模式。这被称为最小归纳偏置。这种方式给模型更大的自由度,但需要更多的数据和更长的学习时间。

感受野解释

上图展示了 CNN 模型感受野如何变化。请记住,CNN 中有多层卷积,随着网络加深,每一层的神经元受到的输入区域影响会逐步扩大。 

这种不断扩大的感受野使得越深的层能够检测更复杂、更全局的特征,如图所示(从耳朵到脸再到整只狗)。 

这与视觉 Transformer 的感受野有明显不同:

Receptive Field of Vision Transformer - Image by Author.

在上图中可以看到,(vanilla)ViT 的感受野是全局的,因此所有补丁都会以不同的权重彼此关注(由箭头粗细表示)。 

为简明起见,我画出了绿色补丁被关注的情况,但实际上所有补丁都会彼此关注(画出来会太乱!)

何时使用 CNN 与视觉 Transformer

明确何时使用 CNN、何时使用视觉 Transformer,对我们在计算机视觉流程中做出有效的架构选择非常重要。 

两类模型各有优劣,它们在不同条件下表现更佳,这取决于数据规模、算力资源以及我们试图建模的视觉模式的复杂度。

我会参考以下几点来做出选择:

在以下情况下应使用 CNN:

  1. 数据有限时——我们已经指出,CNN 具有如局部性等强归纳偏置,使其即便在较小数据集上也能很好泛化。这使得它在难以收集大量标注数据的任务中是可靠的选择。
  2. 需要实时性能时——CNN 的推理延迟通常更低,并且在标准硬件上计算更高效,尤其是移动端与边缘设备。其架构在较少参数与更低内存占用下实现性能优化。
  3. 算力预算有限时——从零训练一个 CNN 或微调预训练 CNN 相比 ViT 需要的资源显著更少。我们可以用单张 GPU 和适中的内存训练出不错的 CNN,而基于 Transformer 的模型并非总是如此。

在以下情况下应使用视觉 Transformer:

  1. 具备大规模标注数据与强大算力——正如对比表所述,视觉 Transformer 很“吃数据”,特别是 vanilla 形态下需要大量训练时间与内存。有足够数据与算力时,它们能在许多基准上超越 CNN。
  2. 需要捕获长程空间关系时——与以局部处理为主的 CNN 不同(回想感受野一节),ViT 通过自注意力建模所有图像补丁之间的关系,尤其适用于需要整个图像空间上下文的任务。
  3. 希望使用预训练模型与迁移学习时——我们可以获取预训练的 ViT(例如通过HuggingFace或 timm),因此微调会更为实用。在这种情况下,即便是中等规模的数据集也能在不从零训练的前提下获得很好的效果。

总之,如果您处于资源受限的环境,或在有限数据下需要快速、稳定的结果,CNN 更安全也更高效。 

相反,若您的项目需要在视觉空间中进行深度的全局推理,且具备足够算力支持,那么视觉 Transformer 会是更好的选择。 

视觉 Transformer 的组成部分

现在让我们深入 ViT 的实际架构。由于其包含多个阶段(确切地说是五个),我们逐一展开。 

阶段一:图像分块(patching)

这是第一步转换,我们需要将图像变为 Transformer 能理解的形式。

假设我们有一张 224×224 的图像。不是整体输入,而是将其划分为更小且固定大小的补丁。如果每个补丁是 16×16(与原论文一致),则图像被分为 14×14 = 196 个补丁,因为(224/16 = 14)。

一个直观的类比是把照片裁成完全等大的方形瓷砖。每个补丁包含局部信息,并将像句子中的“词”那样发挥作用。

在上图中,您可以看到狼的图像被分割为更小且相等的方形补丁——确切地说是 196 个。

不过您可能会有一个问题:

“图像是 RGB 的,那3 个颜色通道怎么办?”

这是一个需要注意的点:每个补丁实际上是一个数据块,在我们的例子中其形状为 16×16×3。16×16 是高和宽,3 是颜色通道,因此每个补丁是一个三维张量。 

这使得每个补丁都紧凑地表示了局部区域的空间与颜色信息。

这种分块操作将图像的二维空间结构转化为一维序列的多维数据,适配 Transformer 模型。

阶段二:补丁嵌入(patch embedding)

接下来,我们需要将每个图像补丁变为 Transformer 可用的数值表示。

每个 16×16 的 RGB 补丁会被展平成一个向量。在我们的例子中,16×16×3 = 768 个值。实质上,我们把所有像素与颜色值“展开”为一串数字。

这些原始向量虽包含信息,但尚不适合学习。因此我们将每个向量送入一个可训练的线性投影层。可以把它看作一个全连接层,将每个 768 维输入映射到大小为D 的新嵌入空间(常见为 768,但不总是)。

为什么?

这一投影很重要,因为它将所有补丁 token 对齐为编码器预期的、一致且可学习的格式。到目前为止,我们已经得到了一串固定长度向量的序列,其中每个向量代表图像的一个补丁。

这个过程将富含空间与颜色信息的补丁转化为统一的“token”序列(类似于NLP 中的词嵌入),从而让 Transformer 将图像视为一组结构化、可学习的实体序列。

阶段三:位置编码

在把这些 token 送入编码器之前,还有一件事要做。如果您了解过 LLM,可能已经猜到了!

Transformer 天生不具备位置信息。

这意味着它将输入序列视为无序集合,没有内置机制去理解 token 的顺序或位置。而在图像分析中,空间结构与布局至关重要!

我们详细谈谈这一点。 

在图像语境中,每个补丁代表原图的特定区域。如果我们不给模型说明每个补丁来自哪里,它就不知道某个补丁是来自左上角还是右下角。 

这种缺乏空间感知会限制模型理解模式、形状与结构等的能力。

为了解决这个问题,我们为每个补丁嵌入添加位置编码。这是一种类似“空间标签”的向量,用于编码补丁在原图中的位置。在进入 Transformer 之前,这些位置向量会直接与补丁嵌入相加。

常见的两类位置编码:

  • 固定(正弦)位置编码:使用不同频率的正弦和余弦函数来编码位置信息。它是确定性的,无需学习。这也是著名论文“Attention is all you need”中采用的方法。
  • 可学习的位置编码:这些是可训练参数,模型在训练中学习。序列中的每个位置都有一个可学习向量。

如果您此刻有些困惑,可以把位置编码想象成补丁的 GPS 坐标。 

它们告诉 Transformer token 的原始位置。加入这些信息后,我们就允许自注意力机制同时建模每个补丁的内容以及其在图像中的相对或绝对位置。

这一步至关重要——没有它,无论注意力机制多强,ViT 都无法重建图像结构。

阶段四:Transformer 编码器块

接下来就是深度学习真正发生的地方——模式发现、上下文建模与特征融合。

通常会将若干编码器块堆叠,构成完整的 ViT 编码器。每个编码器块包含如下组件:

  • 多头自注意力(MHSA):允许每个补丁与所有其他补丁交互。它帮助模型判断图像中哪些部分彼此最相关。
    • MHSA 不再孤立地看补丁,而是赋予其全局上下文。
    • 多头注意力关注关系的不同方面,从而获得更丰富的图像解释。 
  • 前馈神经网络(FFN):这是一个小型的多层感知机,在注意力之后独立处理每个 token。这样做是为了提升模型对复杂非线性模式的学习能力。
  • 残差连接与层归一化:用于稳定训练,并帮助梯度更好地在网络中传播。

总体而言,这些编码器块按层堆叠(如 ViT-Base 为 12 层),每层都会学习到更深层次的图像表示。

阶段五:[CLS] token 与 MLP 头

快到终点了!现在进入(vanilla)ViT 的最后阶段。 

通过所有编码器层之后,我们需要一种方式来汇总全部信息并做出预测。

因此我们在补丁序列的开头引入一个特殊的[CLS] token。这是一个可学习的嵌入,它会在编码过程中与所有其他 token 交互。

随着其通过各层 Transformer,它会聚合全局信息。最后,我们提取该 token,并将其送入另一个多层感知机(MLP)头,通常由若干全连接层和一个 softmax(用于分类)组成。

最后一步会给出预测标签,比如基于[CLS] token 汇聚的信息说“这是一只狗”。务必注意:只有 [CLS] token 会被送入最终的 MLP 头并用于分类,其他 token 不会!

值得关注的视觉 Transformer 架构

我们来看看一些具有里程碑意义的视觉 Transformer 架构。 

  • ViT(Google Brain)——最早的视觉 Transformer(也是本文一直跟随的版本),提出了基于补丁的 Transformer 模型,用于在 ImageNet-21k 等大型数据集上进行图像分类。 
  • DeiT(高数据效率的图像 Transformer)——由 Facebook AI 提出,通过引入知识蒸馏技术提升数据效率,使得在较小数据集上也能训练。
  • Swin Transformer——采用分层架构与移位窗口,在保持精度的同时降低计算量。尤其在目标检测等密集预测任务中表现出色。
  • 其他变体——如 BEiT、CvT、MobileViT 等在训练效率、模型紧凑性或推理速度上进行了多种设计改进。

使用预训练的视觉 Transformer 

很多时候我们不会从零训练 ViT!因此来看看如何使用主流深度学习库中的预训练 ViT 模型。

Hugging Face Transformers 库

使用 Hugging Face Transformers 库是调用视觉 Transformer 最简单可靠的方法之一。它提供预训练的 ViT 模型,只需几行代码即可使用。 

我在这里写了些代码来下载并使用原始的 ViT Base 模型:

# Imports neededfrom transformers import ViTImageProcessor, ViTForImageClassificationfrom PIL import Imageimport torchimport requests# Here I am loading and preprocessing an image - but you can use any image of your choiceimage = Image.open(requests.get(https://huggingface.co/datasets/huggingface/cats-image?image-viewer=ED6CB286C90CF5F1FAF278077DF091477DF05416', stream=True).raw)# Here, we load and process the ViT modelprocessor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')# Preprocessing the image so it gets it into the necessary formatinputs = processor(images=image, return_tensors='pt')# Forward pass - with no backpropagationwith torch.no_grad():    outputs = model(**inputs)# Here we get the predicted labels (since we are performing classification)logits = outputs.logits	predicted_class = logits.argmax(-1) # Selects the index of the class with the highest logit score (highest predicted probability)print(f"Predicted class index: {predicted_class.item()}")

我们稍微深入讲讲这段代码。首先,您可能会问为什么要下载两个组件。它们相互独立,但彼此配合:

  1. ViTImageProcessor —— 负责在输入模型前预处理图像。它处理大小调整、归一化,并将图像转换为 ViT 期望的格式。最好的类比是:它是视觉领域中 NLP“分词器”的对应物!
  2. ViTForImageClassification —— 实际执行推理的模型。它期望特定格式的张量输入,并输出预测结果(如 logits 或类别概率)。

您可能还有一个关于 output.logits 的疑问,更具体地说,logits 是什么。 

在机器学习模型——尤其是分类模型——中,logits 是指模型最后一层在应用 softmax 之类的激活函数之前产生的原始输出值。它们代表每个类别的未归一化得分。

可以把它们理解为模型对各类别的内部“信心”程度。这些值可以为正或负,且不要求总和为 1。

要将 logits 转为概率,需要应用softmax 函数,把它们缩放为概率分布。

因此,当我们使用 argmax(logits, -1) 时,我们只是选择具有最高原始得分的索引(类别),通常也就对应 softmax 后概率最高的类别。

使用 timm 库

另一个使用 ViT 的优秀且知名的库是 timm(由 Ross Wightman 开发)。它同样提供大量预训练的 ViT 变体,适合快速试验或微调。以下演示了如何加载并在推理中使用 ViT:

# Relevant importsimport timmimport torchfrom torchvision import transformsfrom PIL import Imageimport requests# Loading an imageurl = 'https://huggingface.co/datasets/huggingface/cats-image?image-viewer=ED6CB286C90CF5F1FAF278077DF091477DF05416'image = Image.open(requests.get(url, stream=True).raw)# Here we have defined the transformation we are going to do on the imagetransform = transforms.Compose([    transforms.Resize((224, 224)),    transforms.ToTensor(),    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])image_tensor = transform(image).unsqueeze(0)  # Adding batch dimension# Loading the ViT modelmodel = timm.create_model('vit_base_patch16_224', pretrained=True)model.eval() # Setting it to evaluation mode# Inference componentwith torch.no_grad():    outputs = model(image_tensor)    predicted_class = outputs.argmax(dim=1) # Obtaining the class with the highest probabilityprint(f"Predicted class index: {predicted_class.item()}")

您现在应该能看出 timm 与 transformers 库之间的相似性。两者都提供预训练的 ViT 模型,都支持图像处理(尽管在 timm 中我们使用了标准的 PyTorch 代码),最终都得到最高 logit 对应的类别。

结语

总体而言,视觉 Transformer 通过把图像看作“由补丁组成的句子”,并从第一层就使用自注意力捕获全局上下文,从而重塑了计算机视觉。在本教程中,我们学习了:

  1. ViT 与 CNN 的差异——最小归纳偏置、全局感受野、强可扩展性等
  2. 如何选择各自模型——CNN 在数据有限与时延严格场景中更优;而在数据充足、硬件强劲时,ViT 表现更佳。
  3. ViT 的底层架构——我们覆盖了从图像分块、可学习位置编码,到多头自注意力、残差路径以及用于分类的 [CLS] token 各阶段内容。
  4. 实践落地——如何从 Hugging Face 或 timm 加载预训练模型、预处理图像并获取预测结果。

想继续学习,建议您阅读Transformer 的工作原理一文,了解 Transformer 背后的理论。随后,我建议您再探索如何从零实现一个 Transformer。  

如果您希望在上手实践的同时同步学习,那么我们的课程Transformer Models with PyTorch是不二之选。

视觉 Transformer 常见问答

为什么 ViT 要把图像切成 16 × 16 的补丁,而不是直接读取原始像素?

分块将 2D 图像变成 1D 的“token 句子”,这样 Transformer 就能应用驱动语言模型的自注意力技巧。

ViT 一定需要数百万张图像才能表现良好吗?

不一定。我们可以从预训练 ViT 出发做微调。为此,通常只需几千张有标注的图像就能取得很好的效果。

ViT 如何“知道”每个补丁来自哪里?

它们会为每个补丁 token 添加位置编码,从而让模型具备空间顺序感。

什么时候我仍应选择 CNN 而非 ViT?

当数据稀缺或需要在边缘设备上进行超快且省内存的推理时,应选择 CNN。

我能可视化 ViT 关注的内容吗?

可以。我们可以提取注意力图,放大后以热力图叠加,显示哪些补丁驱动了预测。

Vaibhav Mehra

作者

Vaibhav Mehra

主题