1. Vision Transformer基础概念Vision TransformerViT是一种基于Transformer架构的图像分类模型。它将图像分割成一系列图像块并将每个图像块作为序列传递给Transformer模型。ViT通过自注意力机制来捕捉图像中的全局上下文信息并使用多层感知机MLP进行特征提取和分类。ViT的核心思想是将图像转换为序列数据这使得模型能够利用Transformer强大的表达能力来处理图像。通过将图像分割成图像块并将它们展平为序列ViT能够在不需要依赖传统卷积神经网络的情况下实现图像分类任务。1.1 Transformer在视觉任务中的应用Transformer最初是为自然语言处理任务设计的但近年来在计算机视觉领域也取得了显著成功。与传统的卷积神经网络CNN不同Transformer不依赖于平移等变性和局部性等先验假设。虽然Transformer不具备这些先天优势但其核心的自注意力机制能够获取长距离信息相比之下CNN需要通过不断堆叠卷积层来获取更大的感受野尽管训练难度相对较大。ViT是Google在2020年提出的直接将Transformer应用于图像分类的模型。许多后续工作都是基于ViT进行改进的。通过这篇论文的实验最佳模型在ImageNet1K上能达到88.55%的准确率先在Google自家的JFT数据集上进行了预训练说明Transformer在CV领域确实是有效的而且效果还相当惊人。2. ViT网络架构详解不同于寻常的分类网络整个Vision Transformer可以分为两部分一部分是特征提取部分另一部分是分类部分。2.1 特征提取部分在特征提取部分ViT所做的工作是特征提取。特征提取部分在图片中的对应区域是PatchPosition Embedding和Transformer Encoder。PatchPosition Embedding的作用主要是对输入进来的图片进行分块处理每隔一定的区域大小划分图片块。然后将划分好的图片块组合成序列。在获得序列信息后传入Transformer Encoder进行特征提取这是Transformer特有的Multi-head Self-attention结构通过自注意力机制关注每个图片块的重要程度。2.1.1 Patch EmbeddingPatch的作用主要是对输入进来的图片进行分块处理每隔一定的区域大小划分图片块。然后将划分好的图片块组合成序列。该部分首先对输入进来的图片进行分块处理处理方式其实很简单使用的是现成的卷积。也就是说不是把图片分割是做了一次简单的卷积可以理解为初步特征提取或者说映射。由于卷积使用的是滑动窗口的思想我们只需要设定特定的步长就可以输入进来的图片进行分块处理了。在ViT中我们常设这个卷积的卷积核大小为16x16步长也为16x16此时卷积就会每隔16个像素点进行一次特征提取由于卷积核大小为16x16两个图片区域的特征提取过程不会有重叠。当我们输入的图片是224, 224, 3的时候我们可以得到一个14, 14, 768的特征层。在代码实现中直接通过一个卷积层来实现。以ViT-B/16为例直接使用一个卷积核大小为16x16步长为16卷积核个数为768的卷积来实现。通过卷积[224, 224, 3] - [14, 14, 768]然后把H以及W两个维度展平即可[14, 14, 768] - [196, 768]此时正好变成了一个二维矩阵正是Transformer想要的。class PatchEmbed(nn.Module): def __init__(self, input_shape[224, 224], patch_size16, in_chans3, num_features768, norm_layerNone, flattenTrue): super().__init__() self.num_patches (input_shape[0] // patch_size) * (input_shape[1] // patch_size) self.flatten flatten self.proj nn.Conv2d(in_chans, num_features, kernel_sizepatch_size, stridepatch_size) self.norm norm_layer(num_features) if norm_layer else nn.Identity() def forward(self, x): x self.proj(x) if self.flatten: x x.flatten(2).transpose(1, 2) # BCHW - BNC x self.norm(x) return x2.1.2 Position EmbeddingPosition Embedding的作用主要是对组合序列加上[class]token以及Position Embedding。除了patch embeddings模型还需要另外两个特殊的position embedding。transformer和CNN不同需要position embedding来编码tokens的位置信息这主要是因为self-attention是permutation-invariant即打乱sequence里的tokens的顺序并不会改变结果。如果不给模型提供patch的位置信息那么模型就需要通过patchs的语义来学习拼图这就额外增加了学习成本。在ViT论文中对比了几种不同的position embedding方案(如下)最终发现如果不提供positional embedding效果会差但其它各种类型的positional embedding效果都接近这主要是因为ViT的输入是相对较大的patchs而不是pixels所以学习位置信息相对容易很多。无positional embedding1-D positional embedding把2-D的patchs看成1-D序列2-D positional embedding考虑patchs的2-D位置x, yRelative positional embeddingspatchs的相对位置transformer原论文中是默认采用固定的positional embedding但ViT中默认采用学习训练的的1-D positional embedding在输入transformer的encoder之前直接将patch embeddings和positional embedding相加:# 这里多1是为了后面要说的class tokenembed_dim即patch embed_dim self.pos_embed nn.Parameter(torch.zeros(1, num_patches 1, embed_dim)) # patch emded pos_embed x x self.pos_embed对于Position Embedding作者也做了一系列对比实验在源码中默认使用的是1D Pos. Emb.对比不使用Position Embedding准确率下降了大概3个点和2D Pos. Emb.比起来没太大区别。2.1.3 Class Token除了patch tokensViT借鉴BERT还增加了一个特殊的class token。前面会说transformer的encoder输入是a sequence patch embeddings输出也是同样长度的a sequence patch features但图像分类最终需要获取image feature简单的策略是采用pooling比如求patch features的平均来获取image feature但是ViT并没有采用类似的pooling策略而是直接增加一个特殊的class token其最终输出的特征加一个linear classifier就可以实现对图像的分类ViT的pre-training时是接一个MLP head所以输入ViT的sequence长度是1。class token对应的embedding在训练时随机初始化然后通过训练得到具体实现如下# 随机初始化 self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) # Classifier head self.head nn.Linear(self.num_features, num_classes) if num_classes 0 else nn.Identity() # 具体forward过程 B x.shape[0] x self.patch_embed(x) cls_tokens self.cls_token.expand(B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks x torch.cat((cls_tokens, x), dim1) x x self.pos_embed2.2 Transformer EncoderTransformer Encoder的作用是将输入序列进行编码生成一个高维表示以便后续的任务处理。它由多个Encoder Block组成每个Encoder Block包含一个多头自注意力层和一个前馈全连接层。在编码过程中输入序列会经过多头自注意力层进行特征提取和关联性计算然后通过前馈全连接层进行非线性变换和特征融合。通过堆叠多个Encoder BlockTransformer Encoder能够捕捉输入序列中的语义信息和上下文关系生成一个更加丰富的表示。下图是太阳花的小绿本绘制的Encoder Block主要由以下几部分组成2.2.1 Layer Norm这种Normalization方法主要是针对NLP领域提出的这里是对每个token进行Norm处理之前也有讲过Layer Norm不懂的可以参考链接2.2.2 Multi-Head AttentionMulti-Head Attention是一种注意力机制它由多个独立的注意力头组成。每个注意力头都可以学习到不同的表示子空间并在不同的位置上联合关注信息。通过使用多个注意力头Multi-Head Attention能够更好地捕捉输入序列中的不同关系和特征。下面是一个示例代码演示了如何实现Multi-Head Attentionimport torch import torch.nn as nn class MultiHeadAttention(nn.Module): def __init__(self, input_dim, num_heads): super(MultiHeadAttention, self).__init__() self.num_heads num_heads self.head_dim input_dim // num_heads self.query_linear nn.Linear(input_dim, input_dim) self.key_linear nn.Linear(input_dim, input_dim) self.value_linear nn.Linear(input_dim, input_dim) self.output_linear nn.Linear(input_dim, input_dim) def forward(self, query, key, value): batch_size query.size(0) # 线性变换得到query、key、value query self.query_linear(query) key self.key_linear(key) value self.value_linear(value) # 将query、key、value分成多个头 query query.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) key key.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) value value.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2) # 计算注意力得分 scores torch.matmul(query, key.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtypetorch.float32)) attention_weights torch.softmax(scores, dim-1) # 对value进行加权求和 weighted_values torch.matmul(attention_weights, value) # 将多个头的结果拼接起来 weighted_values weighted_values.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.head_dim) # 线性变换得到最终的输出 output self.output_linear(weighted_values) return output2.2.3 MLP BlockMLP Block如图右侧所示就是全连接GELU激活函数Dropout组成也非常简单需要注意的是第一个全连接层会把输入节点个数翻4倍[197, 768] - [197, 3072]第二个全连接层会还原原节点个数[197, 3072] - [197, 768]2.3 分类部分MLP Head是指多层感知机头部它是在Transformer Encoder后面用于分类任务的一部分。MLP Head通常由几个线性层组成用于将Transformer Encoder的输出转换为最终的分类结果。在原始的MLP Head论文中它由线性层、tanh激活函数和线性层组成。但是在迁移到ImageNet1K或其它数据集时只需要使用一个线性层即可。下面是一个示例代码展示了如何使用PyTorch实现一个简单的MLP Headimport torch import torch.nn as nn class MLPHead(nn.Module): def __init__(self, input_dim, output_dim): super(MLPHead, self).__init__() self.linear nn.Linear(input_dim, output_dim) def forward(self, x): x self.linear(x) return x # 假设Transformer Encoder的输出维度为768分类任务的类别数为10 input_dim 768 output_dim 10 mlp_head MLPHead(input_dim, output_dim) # 假设Transformer Encoder的输出为transformer_output transformer_output torch.randn(197, 768) # 将transformer_output输入到MLP Head中得到分类结果 classification_result mlp_head(transformer_output) print(classification_result.shape) # 输出torch.Size([197, 10])在上述代码中我们定义了一个MLPHead类它接受一个输入维度和一个输出维度作为参数。在forward方法中我们使用一个线性层将输入转换为输出。然后我们可以将Transformer Encoder的输出输入到MLP Head中得到最终的分类结果上面ViT通过Transformer Encoder后输出的shape和输入的shape是保持不变的以ViT-B/16为例输入的是[197, 768]输出的还是[197, 768]。注意在Transformer Encoder后其实还有一个Layer Norm没有画出来前面有我自己画的ViT的模型可以看到详细结构。这里我们只是需要分类的信息所以我们只需要提取出[class]token对应的结果即可即[197, 768]中抽取出[class]token对应的[1, 768]。接着我们通过MLP Head得到我们最终的分类结果。MLP Head原论文中说在训练ImageNet21K时是由Lineartanh激活函数Linear组成。但是迁移到ImageNet1K上或者你自己的数据上时只用一个Linear即可。3. ViT完整代码实现3.1 Patch Embedding实现class PatchEmbed(nn.Module): def __init__(self, input_shape[224, 224], patch_size16, in_chans3, num_features768, norm_layerNone, flattenTrue): super().__init__() self.num_patches (input_shape[0] // patch_size) * (input_shape[1] // patch_size) self.flatten flatten self.proj nn.Conv2d(in_chans, num_features, kernel_sizepatch_size, stridepatch_size) self.norm norm_layer(num_features) if norm_layer else nn.Identity() def forward(self, x): x self.proj(x) if self.flatten: x x.flatten(2).transpose(1, 2) # BCHW - BNC x self.norm(x) return x3.2 Transformer Block实现class Mlp(nn.Module): MLP as used in Vision Transformer, MLP-Mixer and related networks def __init__(self, in_features, hidden_featuresNone, out_featuresNone, act_layerGELU, drop0.): super().__init__() out_features out_features or in_features hidden_features hidden_features or in_features drop_probs (drop, drop) self.fc1 nn.Linear(in_features, hidden_features) self.act act_layer() self.drop1 nn.Dropout(drop_probs[0]) self.fc2 nn.Linear(hidden_features, out_features) self.drop2 nn.Dropout(drop_probs[1]) def forward(self, x): x self.fc1(x) x self.act(x) x self.drop1(x) x self.fc2(x) x self.drop2(x) return x class Block(nn.Module): def __init__(self, dim, num_heads, mlp_ratio4., qkv_biasFalse, drop0., attn_drop0., drop_path0., act_layerGELU, norm_layernn.LayerNorm): super().__init__() self.norm1 norm_layer(dim) self.attn Attention(dim, num_headsnum_heads, qkv_biasqkv_bias, attn_dropattn_drop, proj_dropdrop) self.norm2 norm_layer(dim) self.mlp Mlp(in_featuresdim, hidden_featuresint(dim * mlp_ratio), act_layeract_layer, dropdrop) self.drop_path DropPath(drop_path) if drop_path 0. else nn.Identity() def forward(self, x): x x self.drop_path(self.attn(self.norm1(x))) x x self.drop_path(self.mlp(self.norm2(x))) return x3.3 完整ViT模型实现class VisionTransformer(nn.Module): def __init__( self, input_shape[224, 224], patch_size16, in_chans3, num_classes1000, num_features768, depth12, num_heads12, mlp_ratio4., qkv_biasTrue, drop_rate0.1, attn_drop_rate0.1, drop_path_rate0.1, norm_layerpartial(nn.LayerNorm, eps1e-6), act_layerGELU ): super().__init__() # 224, 224, 3 - 196, 768 self.patch_embed PatchEmbed(input_shapeinput_shape, patch_sizepatch_size, in_chansin_chans, num_featuresnum_features) num_patches (224 // patch_size) * (224 // patch_size) self.num_features num_features self.new_feature_shape [int(input_shape[0] // patch_size), int(input_shape[1] // patch_size)] self.old_feature_shape [int(224 // patch_size), int(224 // patch_size)] # classtoken部分是transformer的分类特征。用于堆叠到序列化后的图像特征中作为一个单位的序列特征进行特征提取。 # 在利用步长为16x16的卷积将输入图片划分成14x14的部分后将14x14部分的特征平铺一幅图片会存在长度为196的序列特征。 # 此时生成一个classtoken将classtoken堆叠到序列长度为196的特征上获得一个序列长度为197的特征。 # 在特征提取的过程中classtoken会与其它图像序列特征进行特征交互。最终分类时我们取出classtoken的特征利用全连接分类。 # 196, 768 - 197, 768 self.cls_token nn.Parameter(torch.zeros(1, 1, num_features)) # 为网络提取到的特征添加上位置信息。 # 以输入图片为224, 224, 3为例我们获得的序列化后的图像特征为196, 768。加上classtoken后就是197, 768 # 此时生成的pos_Embedding的shape也为197, 768代表每一个特征的位置信息。 # 197, 768 - 197, 768 self.pos_embed nn.Parameter(torch.zeros(1, num_patches 1, num_features)) self.pos_drop nn.Dropout(pdrop_rate) # 197, 768 - 197, 768 12次 dpr [x.item() for x in torch.linspace(0, drop_path_rate, depth)] self.blocks nn.Sequential( *[ Block( dimnum_features, num_headsnum_heads, mlp_ratiomlp_ratio, qkv_biasqkv_bias, dropdrop_rate, attn_dropattn_drop_rate, drop_pathdpr[i], norm_layernorm_layer, act_layeract_layer ) for i in range(depth) ] ) self.norm norm_layer(num_features) self.head nn.Linear(num_features, num_classes) if num_classes 0 else nn.Identity() def forward_features(self, x): x self.patch_embed(x) cls_token self.cls_token.expand(x.shape[0], -1, -1) x torch.cat((cls_token, x), dim1) cls_token_pe self.pos_embed[:, 0:1, :] img_token_pe self.pos_embed[:, 1:, :] img_token_pe img_token_pe.view(1, *self.old_feature_shape, -1).permute(0, 3, 1, 2) img_token_pe F.interpolate(img_token_pe, sizeself.new_feature_shape, modebicubic, align_cornersFalse) img_token_pe img_token_pe.permute(0, 2, 3, 1).flatten(1, 2) pos_embed torch.cat([cls_token_pe, img_token_pe], dim1) x self.pos_drop(x pos_embed) x self.blocks(x) x self.norm(x) return x[:, 0] def forward(self, x): x self.forward_features(x) x self.head(x) return x def freeze_backbone(self): backbone [self.patch_embed, self.cls_token, self.pos_embed, self.pos_drop, self.blocks[:8]] for module in backbone: try: for param in module.parameters(): param.requires_grad False except: module.requires_grad False def Unfreeze_backbone(self): backbone [self.patch_embed, self.cls_token, self.pos_embed, self.pos_drop, self.blocks[:8]] for module in backbone: try: for param in module.parameters(): param.requires_grad True except: module.requires_grad True4. 模型训练与评估4.1 数据准备在开始训练ViT模型之前我们需要准备合适的数据集。常用的图像分类数据集包括CIFAR-10/100、ImageNet等。这里我们以CIFAR-10为例展示如何准备数据import torchvision import torchvision.transforms as transforms # 数据预处理 transform_train transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) transform_test transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]) # 加载CIFAR-10数据集 trainset torchvision.datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform_train) trainloader torch.utils.data.DataLoader(trainset, batch_size32, shuffleTrue, num_workers2) testset torchvision.datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtransform_test) testloader torch.utils.data.DataLoader(testset, batch_size32, shuffleFalse, num_workers2) classes (plane, car, bird, cat, deer, dog, frog, horse, ship, truck)4.2 模型初始化接下来我们初始化ViT模型。对于CIFAR-10数据集我们需要将num_classes设置为10model VisionTransformer( input_shape[224, 224], patch_size16, in_chans3, num_classes10, num_features768, depth12, num_heads12, mlp_ratio4., qkv_biasTrue, drop_rate0.1, attn_drop_rate0.1, drop_path_rate0.1 )4.3 训练过程定义损失函数和优化器并开始训练模型import torch.optim as optim criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr0.001, weight_decay0.01) # 训练模型 for epoch in range(10): # 训练10个epoch running_loss 0.0 for i, data in enumerate(trainloader, 0): inputs, labels data optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() if i % 100 99: # 每100个batch打印一次 print([%d, %5d] loss: %.3f % (epoch 1, i 1, running_loss / 100)) running_loss 0.0 print(Finished Training)4.4 模型评估训练完成后我们可以在测试集上评估模型性能correct 0 total 0 with torch.no_grad(): for data in testloader: images, labels data outputs model(images) _, predicted torch.max(outputs.data, 1) total labels.size(0) correct (predicted labels).sum().item() print(Accuracy of the network on the 10000 test images: %d %% % ( 100 * correct / total))5. 实际应用中的技巧与注意事项5.1 学习率调整策略在训练ViT模型时适当的学习率调整策略可以显著提高模型性能。常用的策略包括热身(Warmup)策略在训练初期逐步提高学习率有助于模型稳定训练。余弦退火(Cosine Annealing)按照余弦函数周期性地调整学习率。线性衰减(Linear Decay)随着训练进行线性降低学习率。以下是实现学习率热身和余弦退火的示例代码from torch.optim.lr_scheduler import LambdaLR def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, num_cycles0.5, last_epoch-1): def lr_lambda(current_step): if current_step num_warmup_steps: return float(current_step) / float(max(1, num_warmup_steps)) progress float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) return max(0.0, 0.5 * (1.0 math.cos(math.pi * float(num_cycles) * 2.0 * progress))) return LambdaLR(optimizer, lr_lambda, last_epoch) # 使用示例 optimizer optim.AdamW(model.parameters(), lr0.001, weight_decay0.01) total_steps len(trainloader) * 10 # 假设训练10个epoch scheduler get_cosine_schedule_with_warmup(optimizer, num_warmup_steps500, num_training_stepstotal_steps) # 在每个训练step后调用 scheduler.step()5.2 混合精度训练使用混合精度训练可以显著减少显存占用并加快训练速度特别是对于ViT这样的大模型from torch.cuda.amp import GradScaler, autocast scaler GradScaler() for epoch in range(10): for i, data in enumerate(trainloader, 0): inputs, labels data inputs, labels inputs.cuda(), labels.cuda() optimizer.zero_grad() with autocast(): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() scheduler.step()5.3 模型微调技巧当在特定任务上微调预训练的ViT模型时可以考虑以下技巧分层学习率对模型的不同部分使用不同的学习率通常backbone使用较小的学习率分类头使用较大的学习率。渐进解冻逐步解冻模型的层从最后一层开始逐步解冻前面的层。标签平滑使用标签平滑正则化来防止模型对训练标签过度自信。# 分层学习率示例 param_groups [ {params: model.patch_embed.parameters(), lr: 0.0001}, {params: model.cls_token, lr: 0.0001}, {params: model.pos_embed, lr: 0.0001}, {params: model.blocks.parameters(), lr: 0.0001}, {params: model.head.parameters(), lr: 0.001} ] optimizer optim.AdamW(param_groups, weight_decay0.01) # 标签平滑示例 class LabelSmoothingCrossEntropy(nn.Module): def __init__(self, epsilon0.1): super().__init__() self.epsilon epsilon def forward(self, preds, target): log_probs F.log_softmax(preds, dim-1) nll_loss -log_probs.gather(dim-1, indextarget.unsqueeze(1)) nll_loss nll_loss.squeeze(1) smooth_loss -log_probs.mean(dim-1) loss (1 - self.epsilon) * nll_loss self.epsilon * smooth_loss return loss.mean() criterion LabelSmoothingCrossEntropy(epsilon0.1)5.4 可视化与解释性理解ViT模型如何做出决策非常重要以下是一些可视化方法注意力图可视化可视化不同注意力头的注意力权重了解模型关注图像的哪些部分。特征图可视化通过PCA或t-SNE等方法可视化模型学到的特征表示。遮挡测试系统地遮挡图像的不同部分观察对模型输出的影响。# 注意力图可视化示例 def visualize_attention(model, img): model.eval() with torch.no_grad(): # 获取注意力权重 outputs model.patch_embed(img) cls_token model.cls_token.expand(outputs.shape[0], -1, -1) x torch.cat((cls_token, outputs), dim1) x x model.pos_embed attention_maps [] for blk in model.blocks: x blk.norm1(x) B, N, C x.shape qkv blk.attn.qkv(x).reshape(B, N, 3, blk.attn.num_heads, C // blk.attn.num_heads).permute(2, 0, 3, 1, 4) q, k, v qkv[0], qkv[1], qkv[2] attn (q k.transpose(-2, -1)) * blk.attn.scale attn attn.softmax(dim-1) attention_maps.append(attn) # 可视化最后一个block的注意力图 last_attn attention_maps[-1][0] # 取第一个样本 # 这里可以添加具体的可视化代码 return last_attn