基于pytorch的视觉变换器-Vision Transformer(ViT)的介绍与应用

近年���，计算机视觉领域因变换器模型的出现而发生了革命性变化。最初为自然语言处理任务设计的变换器，在捕捉视觉数据的空间依赖性方面也显示出了惊人的能力。视觉变换器（Vision Transformer，简称ViT）就是这种变革的一个典型例子，它提出了一种新颖的架构，在各种图像分类任务上实现了最先进的性能。

在这篇文章中，我们将一起构建我们自己的视觉变换器模型，使用PyTorch进行实现。通过逐步分解实现过程，我们旨在提供对ViT架构的全面理解，并使您能够清晰地掌握其内部工作原理。当然，我们可以直接使用PyTorch内置的视觉变换器模型，但那样就没有乐趣了。

我们将从设置必要的依赖和库开始，确保整个项目的顺利进行。接下来，我们将进入数据获取阶段，获取适合训练我们的视觉变换器模型的数据集。

为了准备训练数据，我们将定义必要的转换操作，用于增强和标准化输入图像。随着数据转换的到位，我们将创建自定义数据集和数据加载器，为训练模型搭建舞台。

要理解视觉变换器架构，从零开始构建它至关重要。在后续部分，我们将逐一解析ViT模型的每个组件，并解释其作用。我们将从负责将输入图像划分为更小块并将其嵌入向量格式的Patch嵌入层开始。紧接着，我们将探索多头自注意力模块，它允许模型捕捉图像块之间的全局和局部关系。

此外，我们还将深入了解机器学习感知器模块，这是一个关键组件，使模型能够捕捉输入数据的层次性表示。通过组合这些组件，我们将构建变换器块，它构成了视觉变换器的核心构建块。

最后，我们将通过使用我们精心打造的组件来创建ViT模型。随着模型的完成，我们可以对其进行实验，微调，并在各种计算机视觉任务上发挥其潜力。

到本文结束时，您将获得对视觉变换器架构及其在PyTorch中的实现的扎实理解。有了这些知识，您将能够根据自己的具体需求修改和扩展模型，甚至在此基础上构建更高级的计算机视觉应用。让我们开始吧。

快速部署

第一步：安装并导入依赖项

!pip install -q torchinfo

import torchfrom torch import nnfrom torchinfo import summary

第二步：获取数据集

import requestsfrom pathlib import Pathimport osfrom zipfile import ZipFile# 定义zip文件的URLurl = "https://github.com/mrdbourke/pytorch-deep-learning/raw/main/data/pizza_steak_sushi.zip"# 发送GET请求以下载文件response = requests.get(url)# 定义数据目录的路径data_path = Path("data")# 定义图像目录的路径image_path = data_path / "pizza_steak_sushi"# 检查图像目录是否已经存在if image_path.is_dir():    print(f"{image_path} 目录已存在。")else:    print(f"未找到 {image_path} 目录，正在创建...")    image_path.mkdir(parents=True, exist_ok=True)# 将下载的内容写入一个zip文件with open(data_path / "pizza_steak_sushi.zip", "wb") as f:    f.write(response.content)# 解压zip文件到图像目录with ZipFile(data_path / "pizza_steak_sushi.zip", "r") as zipref:    zipref.extractall(image_path)# 删除下载的zip文件os.remove(data_path / 'pizza_steak_sushi.zip')

第三步：定义转化器

1.调正图像大小

2.运用ToTensor()函数转化格式

from torchvision.transforms import Resize, Compose, ToTensor# 使用Compose定义train_transformtrain_transform = Compose([    Resize((224, 224)),  # 将图像调整为224x224像素    ToTensor()           # 将图像转换为Tensor格式])# 使用Compose定义test_transformtest_transform = Compose([    Resize((224, 224)),  # 将图像调整为224x224像素    ToTensor()           # 将图像转换为Tensor格式])

第四步：创建数据集和数据加载器

我们可以使用PyTorch的ImageFolder数据集库来创建我们的数据集。

为了使ImageFolder工作，您的数据文件夹需要按照这种方式进行结构化。

data└── pizza_steak_sushi    ├── test    │   ├── pizza    │   ├── steak    │   └── sushi    └── train        ├── pizza        ├── steak        └── sushi

所有的比萨图片将位于train和test子文件夹下的pizza文件夹中，其他类别的图片也是如此。

在创建的training_dataset和test_dataset上，有两个有用的方法可以调用：

1. training_dataset.classes，它会返回['pizza', 'steak', 'sushi']
2. training_dataset.class_to_idx，它会返回{'pizza': 0, 'steak': 1, 'sushi': 2}

from torchvision.datasets import ImageFolderfrom torch.utils.data import DataLoaderBATCH_SIZE = 32# 定义数据目录data_dir = Path("data/pizza_steak_sushi")# 使用ImageFolder创建训练数据集training_dataset = ImageFolder(root=data_dir / "train", transform=train_transform)# 使用ImageFolder创建测试数据集test_dataset = ImageFolder(root=data_dir / "test", transform=test_transform)# 使用DataLoader创建训练数据加载器training_dataloader = DataLoader(    dataset=training_dataset,    shuffle=True,                # 开启随机打乱    batch_size=BATCH_SIZE,       # 批量大小设置为32    num_workers=2                # 使用2个工作进程)# 使用DataLoader创建测试数据加载器test_dataloader = DataLoader(    dataset=test_dataset,    shuffle=False,               # 关闭随机打乱    batch_size=BATCH_SIZE,       # 批量大小设置为32    num_workers=2                # 使用2个工作进程)

我们可以可视化一些训练数据集图像，并查看它们的标签

import matplotlib.pyplot as pltimport randomnum_rows = 5num_cols = num_rows# 创建一个包含多个子图的图形fig, axs = plt.subplots(num_rows, num_cols, figsize=(10, 10))# 遍历子图并展示训练数据集中的随机图片for i in range(num_rows):    for j in range(num_cols):        # 从训练数据集中随机选择一个索引        image_index = random.randrange(len(training_dataset))        # 在子图中显示图片        axs[i, j].imshow(training_dataset[image_index][0].permute((1, 2, 0)))        # 设置子图的标题为对应的类名        axs[i, j].set_title(training_dataset.classes[training_dataset[image_index][1]], color="white")        # 禁用坐标轴以获得更好的视觉效果        axs[i, j].axis(False)# 设置图形的超标题fig.suptitle(f"训练数据集中随机{num_rows * num_cols}张图片", fontsize=16, color="white")# 将图形的背景色设置为黑色fig.set_facecolor(color='black')# 展示绘图plt.show()

理解视觉变换器架构

让我们现在花些时间来理解视觉变换器架构。这是原始视觉变换器论文的链接：https://arxiv.org/abs/2010.11929。

下面，您可以看到论文中提出的架构图。

视觉变换器（ViT）是一种为图像处理任务设计的变换器架构类型。与传统的变换器不同，后者主要处理词嵌入序列，ViT处理的是图像嵌入序列。换句话说，它将输入图像分割成多个小块，并将这些小块视为一系列可学习的嵌入。

ViT的基本操作步骤包括：

1. 创建补丁嵌入：它将图像分割成补丁，并将每个补丁转换成嵌入向量。
2. 通过变换器块传递嵌入：
   • 补丁嵌入连同一个分类令牌一起，被送入多个变换器块。
   • 每个变换器块由一个多头自注意力模块（MultiHead Self-Attention Block，简称MSA块）和一个多层感知器块（Multi-Layer Perceptron Block，简称MLP块）组成。
   • 在变换器块的输入和MSA块的输入之间，以及MLP块的输入和MLP块的输出之间建立跳过连接，以缓解梯度消失问题。
   • 执行分类：
     • 变换器块的最终输出通过一个MLP块进行处理。
     • 使用其中的分类令牌，该令牌包含有关输入图像类别的信息，来进行预测。

接下来，我们将详细探讨每个步骤，首先从创建补丁嵌入的关键过程开始。

第五步：创建补丁嵌入层

根据ViT论文，我们需要对图像执行以下操作，然后才能传递给多头自注意力变换器层：

1. 将图像转换为16x16像素大小的补丁。
2. 将每个补丁嵌入768维空间。因此，每个补丁变成一个[1 x 768]向量。每张图像将有KaTeX parse error: Can't use function '\(' in math mode at position 1: \̲(̲N = H \times W …个这样的补丁。这样会得到形状为[14 x 14 x 768]的图像。
3. 将图像沿着一个向量展平，得到一个[196 x 768]矩阵，即我们的图像嵌入序列。
4. 在上述输出前加上类别令牌嵌入。
5. 将位置嵌入添加到类别令牌和图像嵌入中。

PATCH_SIZE = 16IMAGE_WIDTH = 224IMAGE_HEIGHT = IMAGE_WIDTHIMAGE_CHANNELS = 3EMBEDDING_DIMS = IMAGE_CHANNELS * PATCH_SIZE**2NUM_OF_PATCHES = int((IMAGE_WIDTH * IMAGE_HEIGHT) / PATCH_SIZE**2)# 图像宽度和高度应该能被补丁大小整除。这是一个检查以确保这一点。assert IMAGE_WIDTH % PATCH_SIZE == 0 and IMAGE_HEIGHT % PATCH_SIZE == 0, print("图像宽度不是补丁大小的整数倍")

4.1 将图像转换为16x16大小的补丁，并为每个大小为768的补丁创建一个嵌入向量。

这可以通过使用一个Conv2D层来实现，其kernel_size（核大小）等于patch_size（补丁大小），stride（步长）也等于patch_size。

conv_layer = nn.Conv2d(in_channels = IMAGE_CHANNELS, out_channels = EMBEDDING_DIMS, kernel_size = PATCH_SIZE, stride = PATCH_SIZE)

我们可以将随机图像传递到卷积层中，看看会发生什么

random_images, random_labels = next(iter(training_dataloader))random_image = random_images[0]# 创建一个新的图形fig = plt.figure(1)# 展示随机选取的图像plt.imshow(random_image.permute((1, 2, 0)))# 禁用坐标轴以获得更好的视觉效果plt.axis(False)# 设置图像的标题plt.title(training_dataset.classes[random_labels[0]], color="white")# 将图形的背景色设置为黑色fig.set_facecolor(color="black")

我们需要将形状更改为[1，14，14，768]，并将输出扁平化为[1，196，768]

# 将图像通过卷积层image_through_conv = conv_layer(random_image.unsqueeze(0))print(f'通过卷积层的嵌入形状 -> {list(image_through_conv.shape)}  {list(image_through_conv_and_flatten.shape)}  {list(class_token_embeddings.shape)}  {list(embedded_image_with_class_token_embeddings.shape)}  {list(position_embeddings.shape)}  {list(final_embeddings.shape)}  {list(patch_embeddings.shape)}  {list(multihead_self_attention_block(patch_embeddings).shape)}  {list(patch_embeddings.shape)}  {list(transformer_block(patch_embeddings).shape)}