如何在基于Transformer的模型帮助下对图像进行分类
自2017年推出以来，Transformer已被广泛认为是一个强大的编码器-解码器模型，可以解决几乎所有的语言建模任务。
BERT、RoBERTa和XLM-RoBERTa是语言处理中最先进的模型的几个例子，它们在架构中使用Transformer编码器的堆栈作为骨干。ChatGPT和GPT系列也使用Transformer的解码器部分来生成文本。可以说，几乎所有最先进的自然语言处理模型都将Transformer纳入其架构中。
Transformer的性能如此之好，以至于不把它用于自然语言处理以外的任务，例如计算机视觉，似乎是一种浪费。然而，最大的问题是：我们是否能将其用于计算机视觉任务？
事实证明，Transformer也有很好的潜力可以应用于计算机视觉任务。2020年，谷歌大脑团队推出了一个基于Transformer的模型，可以用来解决一个名为Vision Transformer（ViT）的图像分类任务。与传统的CNN相比，它在几个图像分类基准上的表现非常有竞争力。
因此，在这篇文章中，我们将讨论这个模型。具体来说，我们将讨论ViT模型是如何工作的，以及我们如何在HuggingFace库的帮助下在自己的自定义数据集上对其进行微调，以完成图像分类任务。
所以，作为第一步，让我们从本文中要使用的数据集开始。
关于数据集
我们将使用一个小吃数据集，你可以很容易地从HuggingFace的dataset 库中获取。这个数据集被列为具有CC-BY 2.0许可证，这意味着你可以自由地分享和使用它，只要你在工作中引用数据集的来源。
让我们来偷看一下这个数据集：
我们只需要几行代码来加载数据集，你可以在下面看到：
!pip install -q datasets
from datasets import load_dataset
# Load dataset
dataset = load_dataset("Matthijs/snacks")
print(dataset)
# Output
'''
DatasetDict({
train: Dataset({
features: ['image', 'label'],
num_rows: 4838
})
test: Dataset({
features: ['image', 'label'],
num_rows: 952
})
validation: Dataset({
features: ['image', 'label'],
num_rows: 955
})
})'''
该数据集是一个字典对象，由4898张训练图像、955张验证图像和952张测试图像组成。
每张图片都有一个标签，它属于20个小吃类中的一个。我们可以用下面的代码检查这20个不同的类：
print(dataset["train"].features['label'].names)
# Output
'''
['apple','banana','cake','candy','carrot','cookie','doughnut','grape',
'hot dog', 'ice cream','juice','muffin','orange','pineapple','popcorn',
'pretzel','salad','strawberry','waffle','watermelon']'''
并且让我们在每个标签和其相应的索引之间建立一个映射。
# Mapping from label to index and vice versa
labels = dataset["train"].features["label"].names
num_labels = len(dataset["train"].features["label"].names)
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
label2id[label] = i
id2label[i] = label
print(label2id)
print(id2label)
# Output
'''
{'apple': 0, 'banana': 1, 'cake': 2, 'candy': 3, 'carrot': 4, 'cookie': 5, 'doughnut': 6, 'grape': 7, 'hot dog': 8, 'ice cream': 9, 'juice': 10, 'muffin': 11, 'orange': 12, 'pineapple': 13, 'popcorn': 14, 'pretzel': 15, 'salad': 16, 'strawberry': 17, 'waffle': 18, 'watermelon': 19}
{0: 'apple', 1: 'banana', 2: 'cake', 3: 'candy', 4: 'carrot', 5: 'cookie', 6: 'doughnut', 7: 'grape', 8: 'hot dog', 9: 'ice cream', 10: 'juice', 11: 'muffin', 12: 'orange', 13: 'pineapple', 14: 'popcorn', 15: 'pretzel', 16: 'salad', 17: 'strawberry', 18: 'waffle', 19: 'watermelon'}
'''
在我们继续前进之前，我们需要知道的一件重要事情是，每张图片都有不同的尺寸。因此，在将图像送入模型之前，我们需要执行一些图像预处理步骤，以达到微调的目的。
现在我们知道了我们正在使用的数据集，让我们仔细看看ViT的架构。
ViT是如何工作的
在引入ViT之前，Transformer模型依赖于自我注意机制的事实给我们将其用于计算机视觉任务带来了很大的挑战。
自我注意机制是基于Transformer的模型能够区分一个词在不同语境中的语义的原因。例如，由于自我注意机制，BERT模型可以区分 " 他们把车停在地下室 " 和*"她在公园里遛狗 "这两个句子中"公园*"这个词的含义*。*
然而，自我注意有一个问题：它是一个计算成本很高的操作，因为它要求每个标记都要关注序列中的其他标记。
现在，如果我们在图像数据上使用自我关注机制，那么图像中的每个像素都需要关注并与其他每个像素进行比较。问题是，如果我们把像素值增加一个，那么计算成本就会成四倍增加。如果我们有一个具有相当大分辨率的图像，这根本不可行。
为了克服这个问题，ViT引入了将输入图像分割成斑块的概念。每个补丁的尺寸为16 x 16像素。假设我们有一张尺寸为48 x 48像素的图像，那么我们的图像的补丁将看起来像这样：
在其应用中，ViT有两个选项，即如何将我们的图像分割成补丁：
- 将大小为
height x width x channel的输入图像重塑为一连串大小为no.of patches x (patch_size^2.channel)的扁平化的二维图像斑块。然后，我们将扁平化的斑块投射到一个基本的线性层中，得到每个斑块的嵌入。 - 将我们的输入图像投射到卷积层中，卷积层的核大小和跨度等于补丁的大小。然后，我们对该卷积层的输出进行平移。
在几个数据集上测试了模型的性能后，发现第二种方法导致了更好的性能。因此，在这篇文章中，我们将使用第二种方法。
让我们用一个玩具例子来演示用卷积层将输入图像分割成补丁的过程。
import torch
import torch.nn as nn
# Create toy image with dim (batch x channel x width x height)
toy_img = torch.rand(1, 3, 48, 48)
# Define conv layer parameters
num_channels = 3
hidden_size = 768 #or emb_dimension
patch_size = 16
# Conv 2D layer
projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size,
stride=patch_size)
# Forward pass toy img
out_projection = projection(toy_img)
print(f'Original image size: {toy_img.size()}')
print(f'Size after projection: {out_projection.size()}')
# Output
'''
Original image size: torch.Size([1, 3, 48, 48])
Size after projection: torch.Size([1, 768, 3, 3])
'''
该模型接下来要做的事情是将这些斑块压平，并按顺序放好，如下图所示：
我们可以用下面的代码来完成扁平化过程：
# Flatten the output after projection with Conv2D layer
patch_embeddings = out_projection.flatten(2).transpose(1, 2)
print(f'Patch embedding size: {patch_embeddings.size()}')
# Output
'''
Patch embedding size: torch.Size([1, 9, 768]) #[batch, no. of patches, emb_dim]
'''
在扁平化过程之后，我们得到的基本上是每个补丁的向量嵌入。这与许多基于Transformer的语言模型中的令牌嵌入类似。
接下来，与BERT类似，ViT将为我们补丁序列中第一个位置的**[CLS]**标记添加一个特殊的向量嵌入。
# Define [CLS] token embedding with the same emb dimension as the patches
batch_size = 1
cls_token = nn.Parameter(torch.randn(1, 1, hidden_size))
cls_tokens = cls_token.expand(batch_size, -1, -1)
# Prepend [CLS] token in the beginning of patch embedding
patch_embeddings = torch.cat((cls_tokens, patch_embeddings), dim=1)
print(f'Patch embedding size: {patch_embeddings.size()}')
# Output
'''
Patch embedding size: torch.Size([1, 10, 768]) #[batch, no. of patches+1, emb_dim]
'''
正如你所看到的，通过在我们的补丁嵌入的开头预加**[CLS]**标记嵌入，序列的长度增加了一个。之后的最后一步是将位置嵌入添加到我们的补丁序列中。这一步很重要，这样我们的ViT模型就可以学习我们补丁的序列顺序。
这个位置嵌入是一个可学习的参数，在训练过程中会被模型更新。
# Define position embedding with the same dimension as the patch embedding
position_embeddings = nn.Parameter(torch.randn(batch_size, 10, hidden_size))
# Add position embedding into patch embedding
input_embeddings = patch_embeddings + position_embeddings
print(f'Input embedding size: {input_embeddings.size()}')
# Output
'''
Input embedding size: torch.Size([1, 10, 768]) #[batch, no. of patches+1, emb_dim]
'''
现在，每个补丁的位置嵌入和矢量嵌入将成为一堆变形器编码器的输入。Transformer编码器的数量取决于你使用的ViT模型的类型。总的来说，有三种类型的ViT模型：
- **ViT-base：**它有12层，隐藏大小为768，总共有86M个参数。
- **ViT-large：**它有24层，隐藏大小为1024，总参数为307M。
- **ViT-huge：**它有32层，隐藏大小为1280，总参数为632M。
在下面的代码片断中，我们假设要使用Vit-base。这意味着我们有12层的Transformer编码器：
## Define parameters for ViT-base (example)
num_heads = 12
num_layers = 12
# Define Transformer encoders' stack
transformer_encoder_layer = nn.TransformerEncoderLayer(
d_model=hidden_size, nhead=num_heads,
dim_feedforward=int(hidden_size * 4),
dropout=0.1)
transformer_encoder = nn.TransformerEncoder(
encoder_layer=transformer_encoder_layer,
num_layers=num_layers)
# Forward pass
output_embeddings = transformer_encoder(input_embeddings)
print(f' Output embedding size: {output_embeddings.size()}')
# Output
'''
Output embedding size: torch.Size([1, 10, 768])
'''
最后，Transformer编码器的堆栈将输出每个图像补丁的最终向量表示。最终向量的维度与我们使用的ViT模型的隐藏尺寸相对应。
基本上就是这样了。
我们当然可以从头开始建立和训练自己的ViT模型。然而，和其他基于Transformer的模型一样，ViT需要在大量的图像数据（14M-300M的图像）上进行训练，以使它们在未见过的数据上有良好的概括性。
如果我们想在一个自定义的数据集上使用ViT，最常见的方法是对预训练的模型进行微调。最简单的方法是利用HuggingFace库。我们所要做的就是调用ViTModel.from_pretrained() 方法，并将我们的预训练模型的路径作为一个参数。HuggingFace的VitModel()类也将作为我们上面讨论的所有步骤的封装器。
!pip install transformers
from transformers import ViTModel
# Load pretrained model
model_checkpoint = 'google/vit-base-patch16-224-in21k'
model = ViTModel.from_pretrained(model_checkpoint, add_pooling_layer=False)
# Example input image
input_img = torch.rand(batch_size, num_channels, 224, 224)
# Forward pass input image
output_embedding = model(input_img)
print(output_embedding)
print(f"Ouput embedding size: {output_embedding['last_hidden_state'].size()}")
# Output
'''
BaseModelOutputWithPooling(last_hidden_state=tensor([[[ 0.0985, -0.2080, 0.0727, ..., 0.2035, 0.0443, -0.3266],
[ 0.1899, -0.0641, 0.0996, ..., -0.0209, 0.1514, -0.3397],
[ 0.0646, -0.3392, 0.0881, ..., -0.0044, 0.2018, -0.3038],
...,
[-0.0708, -0.2932, -0.1839, ..., 0.1035, 0.0922, -0.3241],
[ 0.0070, -0.3093, -0.0217, ..., 0.0666, 0.1672, -0.4103],
[ 0.1723, -0.1037, 0.0317, ..., -0.0571, 0.0746, -0.2483]]],
grad_fn=<NativeLayerNormBackward0>), pooler_output=None, hidden_states=None, attentions=None)
Output embedding size: torch.Size([1, 197, 768])
'''
完整的ViT模型的输出是一个向量嵌入，代表每个图像补丁和**[CLS]**标记。它的维度是[batch_size, image_patches+1, hidden_size] 。
为了执行图像分类任务，我们遵循与BERT模型相同的方法。我们提取**[CLS]**标记的输出向量嵌入并通过最后的线性层来确定图像的类别。
num_labels = 20
# Define linear classifier layer
classifier = nn.Linear(hidden_size, num_labels)
# Forward pass on the output embedding of [CLS] token
output_classification = classifier(output_embedding['last_hidden_state'][:, 0, :])
print(f"Output embedding size: {output_classification.size()}")
# Output
'''
Output embedding size: torch.Size([1, 20]) #[batch, no. of labels]
'''
微调的实现
在本节中，我们将对一个ViT-base模型进行微调，该模型是在ImageNet-21K数据集上预训练的，该数据集由大约1400万张图像和21843个类别组成。数据集中的每张图片的尺寸为224 x 224像素。
首先，我们需要为预训练的模型定义检查点路径并加载必要的库。
import numpy as np
import torch
import cv2
import torch.nn as nn
from transformers import ViTModel, ViTConfig
from torchvision import transforms
from torch.optim import Adam
from torch.utils.data import DataLoader
from tqdm import tqdm
#Pretrained model checkpoint
model_checkpoint = 'google/vit-base-patch16-224-in21k'
图像数据加载器
如前所述，ViT-base模型已经在一个由尺寸为224 x 224像素的图像组成的数据集上进行了预训练。这些图像也已经根据其每个颜色通道的特定平均值和标准偏差进行了归一化处理。
因此，在我们将自己的数据集输入ViT模型进行微调之前，我们必须首先对图像进行预处理。这包括将每张图像转化为张量，调整其大小至适当的维度，然后使用与模型预训练的数据集相同的平均值和标准差对其进行标准化处理。
class ImageDataset(torch.utils.data.Dataset):
def __init__(self, input_data):
self.input_data = input_data
# Transform input data
self.transform = transforms.Compose([
transforms.ToTensor(),
transforms.Resize((224, 224), antialias=True),
transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
])
def __len__(self):
return len(self.input_data)
def get_images(self, idx):
return self.transform(self.input_data[idx]['image'])
def get_labels(self, idx):
return self.input_data[idx]['label']
def __getitem__(self, idx):
# Get input data in a batch
train_images = self.get_images(idx)
train_labels = self.get_labels(idx)
return train_images, train_labels
从上面的图像数据加载器中，我们将得到一批经过预处理的图像和它们相应的标签。在微调过程中，我们可以使用上述图像数据加载器的输出作为我们模型的输入。
模型的定义
我们的ViT模型的架构是简单明了的。由于我们将对一个预训练的模型进行微调，我们可以使用VitModel.from_pretrained() 方法，并提供模型的检查点作为参数。
我们还需要在最后添加一个线性层，它将作为最终的分类器。这一层的输出应该等于我们数据集中不同标签的数量。
class ViT(nn.Module):
def __init__(self, config=ViTConfig(), num_labels=20,
model_checkpoint='google/vit-base-patch16-224-in21k'):
super(ViT, self).__init__()
self.vit = ViTModel.from_pretrained(model_checkpoint, add_pooling_layer=False)
self.classifier = (
nn.Linear(config.hidden_size, num_labels)
)
def forward(self, x):
x = self.vit(x)['last_hidden_state']
# Use the embedding of [CLS] token
output = self.classifier(x[:, 0, :])
return output
上述ViT模型为每个图像补丁加上**[CLS]标记生成最终的向量嵌入。为了对图像进行分类，正如你所看到的，我们提取[CLS]**标记的最终向量嵌入并将其传递给最终的线性层以获得最终的类别预测。
模型微调
现在我们已经定义了模型架构，并准备好了用于批处理的输入图像，我们可以开始微调我们的ViT模型。训练脚本是一个标准的Pytorch训练脚本，你可以在下面看到：
def model_train(dataset, epochs, learning_rate, bs):
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
# Load nodel, loss function, and optimizer
model = ViT().to(device)
criterion = nn.CrossEntropyLoss().to(device)
optimizer = Adam(model.parameters(), lr=learning_rate)
# Load batch image
train_dataset = ImageDataset(dataset)
train_dataloader = DataLoader(train_dataset, num_workers=1, batch_size=bs, shuffle=True)
# Fine tuning loop
for i in range(epochs):
total_acc_train = 0
total_loss_train = 0.0
for train_image, train_label in tqdm(train_dataloader):
output = model(train_image.to(device))
loss = criterion(output, train_label.to(device))
acc = (output.argmax(dim=1) == train_label.to(device)).sum().item()
total_acc_train += acc
total_loss_train += loss.item()
loss.backward()
optimizer.step()
optimizer.zero_grad()
print(f'Epochs: {i + 1} | Loss: {total_loss_train / len(train_dataset): .3f} | Accuracy: {total_acc_train / len(train_dataset): .3f}')
return model
Hyperparameters
EPOCHS = 10 LEARNING_RATE = 1e-4 BATCH_SIZE = 8
Train the model
trained_model = model_train(dataset['train'], EPOCHS, LEARNING_RATE, BATCH_SIZE)
由于我们的小吃数据集有20个不同的类，那么我们要处理的是一个多类的分类问题。因此，CrossEntropyLoss() 将是合适的损失函数。在上面的例子中，我们训练了10个epochs的模型，学习率被设置为1e-4，批次大小为8。 你可以通过这些超参数来调整模型的性能。
在你训练完模型后，你会得到一个与下面类似的输出：
模型预测
既然我们已经微调了我们的模型，自然我们要用它来对测试数据进行预测。要做到这一点，首先让我们创建一个函数，将所有必要的图像预处理步骤和模型推理过程封装起来。
def predict(img):
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Resize((224, 224)),
transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
])
img = transform(img)
output = trained_model(img.unsqueeze(0).to(device))
prediction = output.argmax(dim=1).item()
return id2label[prediction]
正如你在上面看到的，推理过程中的图像预处理步骤与我们在训练数据上做的步骤完全相同。然后，我们将转换后的图像作为我们训练过的模型的输入，最后我们将其预测结果映射到相应的标签上。
如果我们想在测试数据上预测一个特定的图像，我们可以直接调用上面的函数，之后我们就会得到预测结果。让我们来试试吧。
print(predict(dataset['test'][900]['image']))# Output: waffle
我们的模型正确预测了我们的测试图像。让我们再试试另一张。
print(predict(dataset['test'][250]['image']))# Output: cookie
我们的模型再次正确预测了测试数据。通过对ViT模型的微调，我们可以在自定义数据集上获得良好的性能。你也可以在图像分类任务中对任何自定义数据集做同样的处理。
总结
在这篇文章中，我们看到了Transformer不仅可以用于语言建模任务，还可以用于计算机视觉任务，在这个例子中就是图像分类。
为此，首先将输入的图像分解成大小为16×16像素的斑块。然后，视觉转化器模型利用转化器编码器的堆栈来学习每个图像补丁的向量表示。最后，我们可以使用在图像补丁序列开始时预置的**[CLS]** 标记的最终向量表示来预测我们输入图像的标签。
我希望这篇文章对你开始使用Vision Transformer模型是有用的。一如既往，你可以在这篇文章中的代码实现中找到 这个笔记本.
