Vision Transformer (ViT) 是将纯Transformer架构应用于计算机视觉任务的里程碑式工作，其核心思想是将图像视为序列，从而利用NLP领域成熟的Transformer架构[^1]。 VIT首次将transformer模型运用到CV领域并且取得了不错的分类效果，该模型只用了transformer模型的编码器部分，并未涉及解码器。可将其拆成三个部分，分别是图像特征嵌入模块、transformer编码器模块和MLP分类模块，这样能更容易捋顺它的结构[^2]。 ViT模型利用Transformer模型在处理上下文语义信息的优势，将图像转换为一种“变种词向量”然后进行处理。多个Patch之间本身具有空间联系，类似于一种“空间语义”，从而获得了比较好的处理效果[^3]。 MLP Head是位于模型顶部的全连接前馈神经网络模块，用于将提取的图像特征表示转换为最终的分类结果或其他预测任务输出，通常跟在Transformer Encoder的输出之后，作为整个模型的最后一层[^4]。 ### 代码示例 以下是一个简单的ViT模型的代码示例（使用PyTorch）： ```python import torch import torch.nn as nn class ViT(nn.Module): def __init__(self, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim): super().__init__() num_patches = (image_size // patch_size) ** 2 patch_dim = 3 * patch_size ** 2 self.patch_embedding = nn.Linear(patch_dim, dim) self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim)) self.cls_token = nn.Parameter(torch.randn(1, 1, dim)) self.transformer = nn.TransformerEncoder( nn.TransformerEncoderLayer(d_model=dim, nhead=heads, dim_feedforward=mlp_dim), num_layers=depth ) self.mlp_head = nn.Sequential( nn.LayerNorm(dim), nn.Linear(dim, num_classes) ) def forward(self, x): # 将图像分割成patch并进行嵌入 patches = x.unfold(2, patch_size, patch_size).unfold(3, patch_size, patch_size) patches = patches.permute(0, 2, 3, 1, 4, 5).contiguous().view(x.shape[0], -1, 3 * patch_size ** 2) x = self.patch_embedding(patches) # 添加cls token cls_tokens = self.cls_token.expand(x.shape[0], -1, -1) x = torch.cat((cls_tokens, x), dim=1) # 添加位置编码 x += self.pos_embedding # 通过Transformer编码器 x = self.transformer(x) # 取cls token的输出 cls_output = x[:, 0] # 通过MLP Head进行分类 output = self.mlp_head(cls_output) return output # 示例使用 image_size = 224 patch_size = 16 num_classes = 10 dim = 768 depth = 12 heads = 12 mlp_dim = 3072 model = ViT(image_size, patch_size, num_classes, dim, depth, heads, mlp_dim) input_tensor = torch.randn(1, 3, image_size, image_size) output = model(input_tensor) print(output.shape) ```