How to Train Your StateoftheArt Vision Transformer? Li
作者：禅与计算机程序设计艺术
1.简介
在自然语言处理与计算机视觉等领域的研究者们均取得了显著成就，在深度学习方面尤其表现突出。其中，在图像分类任务中一种高效且具有创新性的方法正是基于transformer架构的设计理念。该方法通过将输入图像划分为多个区域并随后将每个像素点映射至特征空间从而提取相应的特征向量随后采用序列化的方式对这些特征进行训练最终展现出超越传统方法的优势。在同一月，在提出该研究团队基础上开发出了基于transformer结构的ViT（Vision Transformer）模型在同一数据集上实现了更优的表现
ViT模型包含三种核心组件：编码器、Transformer架构和解码器模块。每个组件都具备特定的功能与作用。其中编码器能够从输入图像中提取出有用且丰富的特征信息；而Transformer通过自注意力机制捕捉全局语义信息，并帮助模型实现从局部到整体的理解能力；最后的解码器则通过将编码层的输出与经过自注意力处理后的结果进行融合处理，在此基础上生成最终的输出结果。
在进行ViT模型的训练过程中，在图像处理阶段需要将原始图像划分为不同尺寸的小块，在这一系列操作中这些小块会被传递给对应的Transformer模块进行处理。具体而言，在这一过程当中 Transformer 会利用其独特的自注意力机制，在此过程中生成反映各子块特征的空间表示。与此同时，在整个训练过程中 需要注意的是 随机缩放和裁剪的方式不仅能够提升图像多样性 还能有效增加可学习样本的数量
本文中，作者对其结构、原理以及训练流程进行了详尽的介绍，并分享了若干实用经验，旨在帮助读者更深入地掌握ViT模型的相关知识。
2.基本概念术语说明
2.1 深度学习
深度学习（Deep Learning）是一种机器学习的技术手段。它指通过模拟人类大脑结构和功能的神经网络算法来进行数据处理与分析。利用海量数据对 neural networks 进行训练后，在复杂的数据中能够自主识别出潜在的模式结构。其主要分支包括 neural networks、convolutional neural networks (CNN) 以及 recurrent neural networks (RNN)。其中，在图像识别任务中广泛应用的是 CNN；而 RNN 则在自然语言处理任务中表现出色。
深度学习中的关键技术包括：
- 神经网络架构
- 优化策略（例如梯度下降法）
- 激活函数模块（包括sigmoid函数、ReLU激活以及softmax分类器等）
- 数据增强技术（包括图像翻转、裁剪和旋转操作等）
2.2 Transformer
Transformer是一种基于注意力机制的序列到序列(Sequence-to-sequence, Seq2Seq)体系，在捕捉全局信息以及处理长距离依赖关系方面具有显著优势。其架构设计参考了LSTM和GRU的思路，并未遵循标准门控网络中的残差连接这一常规路径。相较于传统的LSTM与GRU架构，在计算效率方面表现更为出色。
2.3 ViT
Vision Transformer（ViT）是一种基于无监督学习的图像分类模型。它属于深度学习领域中的一个模型，并主要由编码器模块、Transformer模块和解码器模块这三个关键组件构成。
ViT模型的结构如下图所示：
- 编码器 用于 从输入图像中提取特征向量，并分为三种不同的架构模块：
- 第一种变体采用 Vision Transformer（ViT）架构；
- 第二种变体基于 ResNet 模型设计；
- 第三种变体则采用了 EfficientNet 结构。
- Transformer 通过 自注意力机制实现对整体图像信息的捕捉，并生成相应的特征图。
- 在解码过程中，在解码过程中，在解码过程中，在解码过程中，在解码过程中，在解码过程中，在解码过程中，在解码过程中，在解码过程中，在解码过程中，在解码过程中，
在解码过程中，
将编码器提取的信息与 Transformer 产生的特征图进行融合，
并最终得到预测结果。
在具体实施时, 将图像分割成若干个小块, 定义为patch, 并将其输入到Transformer模型中进行训练. 由于Transformer架构采用了自注意力机制的设计, 因此, 在Transformer模型中, 并不需要显式地输入原始图像的空间尺寸信息. 即使如此, 通过动态地划分不同尺寸的空间分辨率区域的方式, 在训练阶段自动完成相应调整.
2.4 小块Attention Mask
在训练Transformer模型时, 必须固定每一个位置的attention mask, 即每个像素是否能够关注周围的其他像素. 因为Transformer架构完全依赖于注意力机制, 所以每个位置都会对所有的先驱位置执行关注操作, 这将导致显存消耗过高且计算开销较大. 为此, 作者提出了分块注意力掩码的设计方案.
该方法的核心机制在于专注于图像中邻近像素的关注机制；即每个像素仅限于对其周围区域内的其他像素进行attention操作而非全局范围内的全部图像信息。例如，在一个规模为N\times N的图像中；其尺寸设定为k\times k单位；这意味着对于任何一个特定的像素而言；只有与其邻近设置在k\times k范围内的所有 pixels 能够在其attention范围内进行操作。
通过这种方式设计的机制，在实现于固定内存环境下的训练过程能够得以完成，并且其性能相较于传统方法有显著提升
2.5 小批量随机梯度下降
在深度学习模型的训练过程中，常见挑战包括梯度爆炸与梯度消失。具体而言，在神经网络的学习过程中，当神经元参数显著增大时（即出现较大的权重值），会导致更新步长异常庞大（即导数值过大），这将使模型难以继续有效优化（即陷入难以恢复的状态）。相反地，在神经元参数显著减小时（即权重值接近零），会导致更新步长异常微小（即导数值过小），这会显著降低模型性能（即学习效果差）。
为了有效解决这一问题，在机器学习算法中通常会采用mini-batch梯度下降方法来实现对这一问题的有效解决。在每一次迭代过程中，算法会选取一批样本进行梯度计算，并非一次性计算所有样本的梯度以避免运算量过大带来的潜在问题。通过这种方法，在每一步参数更新时的变化幅度都会比较小，并且能够有效防止出现极端情况下的模型行为偏差或不稳定现象发生
随着网络规模不断扩大时
3.核心算法原理及具体操作步骤
ViT模型的训练主要由以下几个步骤组成：
Patch Embedding技术：对输入图像进行分割处理后，每个分割块（patch）被转化为向量形式。
基于Transformer架构的设计理念，在图像分割任务中引入注意力机制，在每个patch上附加位置编码信息。
通过Positional Encoding方法，在每个patch上附加位置编码信息。
基于MLP Head设计，在特征图上建立非线性映射关系以生成预测结果。
下面结合具体的代码展示一下ViT模型的训练过程。
3.1 Patch Embedding
在ViT模型中，Patch Embedding技术被视为核心步骤之一。这一过程通过将输入图像分割成多个小块来处理细节信息，并对每个小块进行向量化处理。
下面是一个例子：
import torch
from torchvision import transforms as T
def patch_embedding():
# 模型初始化
model = ViT()
# 创建输入数据
imgs = torch.rand((1, 3, 224, 224))
# 执行Patch Embedding操作
x = model.patch_embed(imgs)
print(x.shape) # [1, 768, 14, 14]
return x
代码解读Patch Embedding通过将输入图像进行分割得到196个尺寸为768像素的小块。输出的空间维度为B×C×H×W，在此表示法中B代表批次大小,C是每个小块的通道数量,H与W分别为小块的高度和宽度。
3.2 Attention
在ViT架构中，Attention模块被定位为关键组件，在其运行过程中完成核心功能。具体而言，在这一阶段中，通过Transformer机制完成注意力运算。随后系统将通过注意力机制捕获全局信息并生成相应的特征图作为后续处理的基础数据来源
下面是一个例子：
import torch
from torchvision import transforms as T
from transformers import ViTFeatureExtractor, ViTModel, ViTPreTrainedModel
class ViTWithAttentionHead(ViTPreTrainedModel):
def __init__(self, config):
super().__init__(config)
self.vit = ViTModel(config)
self.head = nn.Linear(config.hidden_size, 10)
def forward(self, input_ids, attention_mask, position_ids=None):
outputs = self.vit(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids)
pooled_output = outputs[1] # 获取Transformer的输出
logits = self.head(pooled_output) # 添加一个线性层，生成最终的预测结果
return logits
model = ViTWithAttentionHead.from_pretrained('google/vit-base-patch16-224')
feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224', do_resize=True, image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5])
def attention():
# 初始化模型
inputs = {"inputs_ids": None, "attention_mask": None}
# 执行Attention操作
with torch.no_grad():
inputs["inputs_ids"] = model.vit.embeddings.patch_embeddings(imgs).flatten(start_dim=2).transpose(1, 2) # 执行Patch Embedding操作
attn_mask = torch.ones_like(inputs["inputs_ids"]).bool().reshape(*inputs["inputs_ids"].shape[:-1]).repeat(1, 1, 1, model.vit.encoder.num_heads)
inputs["attention_mask"] = attn_mask * (inputs["inputs_ids"] > 0).unsqueeze(-1).repeat(1, 1, 1, model.vit.encoder.num_heads) # 添加小块attention mask
pooled_output = model.vit.encoder(inputs["inputs_ids"], attention_mask=inputs["attention_mask"])[-1][-1].view(model.vit.encoder.num_layers, -1, model.vit.encoder.all_head_size).permute([1, 2, 0, 3]).contiguous().view([-1, model.vit.encoder.all_head_size])
output = model.head(pooled_output)
return output
output = attention()
print(output.shape) # [1, 10]
代码解读首先由ViTModel模块执行操作，在给定input_ids、attention_mask和position_ids的情况下得到其输出结果；接着通过计算pooled_output并应用一个线性层来获得最终预测结果。
Positional Encoding属于ViT模型的关键步骤之一，在其架构中起着重要的作用。该步骤通过向每个token注入其在序列中的位置信息，在向量表示中嵌入了位置信息。从而使得每个token能够被区分为序列中的不同位置。
import torch
from torchvision import transforms as T
def positional_encoding():
# 模型初始化
model = ViT()
# 创建输入数据
imgs = torch.rand((1, 3, 224, 224))
# 执行Positional Encoding操作
pos_enc = model.positional_encoding(imgs.shape)
print(pos_enc.shape) # [1, 768, 224, 224]
return pos_enc
代码解读Positional Encoding的实现比较简单，直接调用一个可学习的矩阵即可。
3.3 Training Strategy
为了加速训练过程，作者设计了两种训练策略：
- Layer-wise Dropout: 在每个Transformer层后插入一个Dropout层以避免模型过拟合。
- Stochastic Depth: 以一定概率随机省略部分Transformers。
下面是一个例子：
import torch
import numpy as np
from torchvision import transforms as T
from transformers import AdamW, get_linear_schedule_with_warmup
def train_step():
# 模型初始化
model = ViT()
optimizer = AdamW(params=model.parameters(), lr=3e-5)
scheduler = get_linear_schedule_with_warmup(optimizer=optimizer, num_warmup_steps=500, num_training_steps=10000)
criterion = nn.CrossEntropyLoss()
for step in range(10000):
batch_data = generate_train_data()
imgs = batch_data["imgs"]
labels = batch_data["labels"]
features = model.forward(imgs)
loss = criterion(features, labels)
if step % 5 == 0:
print("loss:", float(loss))
model.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
代码解读Layer Dropout其实现极为简便，在每一个Transformer后方紧跟一个Dropout层。
除了基础模型外，在每一块拼接块之后添加一个跳跃连接，并将两个分支的输出进行加权求和
4.具体代码实例和解释说明
4.1 安装环境
ViT模型的训练基于PyTorch和transformers库。在Python环境中使用pip命令安装：
!pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 -f https://download.pytorch.org/whl/torch_stable.html
!pip install transformers==4.5.1
代码解读建议您采用非当前版本的pytorch或transformers时，请注意其版本号的一致性。
4.2 数据准备
ViT模型的训练过程通常会采用ImageNet数据集作为基准。该库提供了一个便捷的方法来获取并加载这一广泛使用的图像分类基准数据集。
import os
import urllib.request
import torch
import torchvision.transforms as transforms
import torchvision.datasets as datasets
def load_imagenet_dataset():
data_dir = '/tmp/imagenet'
valdir = os.path.join(data_dir, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
dataset = datasets.ImageFolder(valdir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
loader = torch.utils.data.DataLoader(
dataset,
batch_size=128,
shuffle=False,
pin_memory=True,
num_workers=4,
)
return loader
loader = load_imagenet_dataset()
for i, batch in enumerate(loader):
images, target = batch
print(images.shape) # [128, 3, 224, 224]
break
代码解读这里我们获取并导入了来自ImageNet的数据作为验证集合。你可以对这个脚本进行调整以导入你自己的训练数据。
4.3 数据加载器
ViT模型的训练依赖于数据加载器，它负责读取数据并封装成Tensor。
import torch
import torch.nn as nn
import numpy as np
from PIL import Image
from torchvision import transforms
class Dataset(torch.utils.data.Dataset):
"""ImageNet数据集"""
def __init__(self, root_dir, transform=None):
self.root_dir = root_dir
self.transform = transform
self.classes = sorted(os.listdir(root_dir))
class_to_idx = {cls_name: idx for idx, cls_name in enumerate(self.classes)}
samples = make_dataset(self.root_dir, class_to_idx)
self.samples = samples
def __len__(self):
return len(self.samples)
def __getitem__(self, index):
path, label = self.samples[index]
image = Image.open(path).convert('RGB')
if self.transform is not None:
image = self.transform(image)
return image, int(label)
def make_dataset(directory, class_to_idx):
"""读取目录下的文件名和标签"""
instances = []
directory = os.path.expanduser(directory)
directories = [os.path.join(directory, d) for d in os.listdir(directory)]
for idx, directory in enumerate(directories):
if os.path.isdir(directory):
classes = os.listdir(directory)
for subdir in classes:
subpath = os.path.join(directory, subdir)
if os.path.isdir(subpath):
for filename in os.listdir(subpath):
filepath = os.path.join(subpath, filename)
if os.path.isfile(filepath):
item = (filepath, idx)
instances.append(item)
return instances
dataset = Dataset('/tmp/imagenet/val/', transform=transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]))
loader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True, num_workers=4)
代码解读这里我们创建了一个定制化的Dataset用于获取根目录下的所有文件名及对应的标签并经过特定流程进行图像预处理
然后，我们创建了一个数据加载器，以shuffle的方式对数据进行批处理。
4.4 模型构建
ViT架构包含三个核心组件：编码器模块、Transformer主体以及解码器模块。在实际应用中，我们可以通过调用预训练的官方库实现高效的模型构建
from transformers import ViTConfig, ViTForImageClassification
config = ViTConfig(
image_size=224,
patch_size=16,
num_channels=3,
hidden_size=768,
num_hidden_layers=12,
num_attention_heads=12,
intermediate_size=3072,
hidden_act='gelu',
dropout=0.1,
attention_dropout=0.1,
classifier_dropout=0.1
)
model = ViTForImageClassification(config)
代码解读这里，我们构建了一个ViT模型，其中ViTConfig用来定义模型参数。
4.5 损失函数和优化器
ViT模型的核心目标通常是交叉熵损失。通过调用PyTorch提供的相应接口及其功能模块，我们能够实现损失函数的定义与优化器的配置。
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.AdamW(model.parameters())
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)
代码解读在本研究中，我们定义了一个交叉熵损失函数（Cross Entropy Loss）为目标函数服务，并通过AdamW优化器实现参数更新。为了进一步提升训练效果，在每一步迭代中结合学习率衰减策略（StepLR）来动态调整学习速率。
4.6 训练循环
ViT模型的训练循环一般分为以下四个步骤：
- 模型初始化
- 模型训练
- 模型评估
- 保存最优模型
from tqdm import trange
best_acc1 = 0.
def train(epoch):
global best_acc1
epoch_loss = 0.
total = 0
correct = 0
model.train()
with trange(len(loader)) as t:
for batch_idx, (inputs, targets) in enumerate(loader):
inputs = inputs.cuda()
targets = targets.cuda()
optimizer.zero_grad()
outputs = model(inputs)["logits"]
_, predicted = torch.max(outputs.data, 1)
total += targets.size(0)
correct += (predicted == targets).sum().item()
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
t.set_description("Train Epoch {:>2}: [{:.2%} | Loss:{:.3f}]".format(epoch, batch_idx / len(loader), epoch_loss / (batch_idx + 1)))
acc1 = 100.*correct/total
if acc1 > best_acc1:
best_acc1 = acc1
save_checkpoint({
'state_dict': model.state_dict(),
'acc1': best_acc1,
'epoch': epoch,
}, is_best=True)
return epoch_loss / len(loader)
def validate(epoch):
global best_acc1
epoch_loss = 0.
total = 0
correct = 0
model.eval()
with torch.no_grad():
with trange(len(valid_loader)) as t:
for batch_idx, (inputs, targets) in enumerate(valid_loader):
inputs = inputs.cuda()
targets = targets.cuda()
outputs = model(inputs)["logits"]
_, predicted = torch.max(outputs.data, 1)
total += targets.size(0)
correct += (predicted == targets).sum().item()
loss = criterion(outputs, targets)
epoch_loss += loss.item()
t.set_description("Valid Epoch {:>2}: [{:.2%} | Loss:{:.3f}]".format(epoch, batch_idx / len(valid_loader), epoch_loss / (batch_idx + 1)))
acc1 = 100.*correct/total
if acc1 > best_acc1:
best_acc1 = acc1
save_checkpoint({
'state_dict': model.state_dict(),
'acc1': best_acc1,
'epoch': epoch,
}, is_best=True)
return epoch_loss / len(valid_loader)
def save_checkpoint(state, filename='checkpoint.pth.tar', is_best=False):
torch.save(state, filename)
if is_best:
shutil.copyfile(filename,'model_best.pth.tar')
代码解读这里，我们创建了训练函数，它在每一个epoch结束后返回训练损失值。
我们还开发了一个验证函数，在每次 epoch 结束时输出验证损失值和准确率值。
最后，我们创建了一个保存检查点的函数，用于保存模型状态和最优模型。
4.7 启动训练
for epoch in range(epochs):
train_loss = train(epoch)
valid_loss = validate(epoch)
print("\nEpoch: {}/{}, Train Loss: {:.3f}, Valid Loss: {:.3f}\n".format(epoch + 1, epochs, train_loss, valid_loss))
代码解读在当前环境中初始化模型参数后，在每个 epoch 开始时执行训练循环，并记录并输出每个 epoch 的训练损失和验证损失数据
4.8 总结
以上部分阐述了如何运用ViT模型进行图像分类任务的训练。在具体实施过程中，我们利用transformers库中的相关类和接口完成了模型构建以及数据加载器、损失函数和优化器等组件的配置设置。
训练循环的具体流程则包含模型训练、评估与保存环节以及参数优化阶段。我们也可以采用多层dropout策略和随机深度混合技术等手段来增强模型的鲁棒性。