1. 认识Vision Transformer当Transformer遇上计算机视觉2017年Transformer横空出世彻底改变了自然语言处理领域的游戏规则。但谁曾想到这个原本为文本序列设计的模型竟能在计算机视觉领域掀起一场革命2020年Google Research团队发表的《An Image is Worth 16x16 Words》论文首次将纯Transformer架构成功应用于图像分类任务这就是我们今天要实现的Vision TransformerViT。ViT的核心思想非常直观把图像拆解成一个个小块patch把这些图像块视为视觉单词然后像处理文本序列一样用Transformer进行处理。这种看似简单的思路却带来了惊人的效果——当训练数据足够多时ViT在ImageNet等基准测试上超越了当时最先进的CNN模型。我在实际项目中测试发现ViT特别擅长捕捉图像中的长距离依赖关系比如识别一只猫的尾巴和耳朵之间的关联这种全局理解能力正是传统CNN的短板。要实现ViT我们需要解决几个关键问题如何将二维图像转换为适合Transformer处理的一维序列如何保留图像特有的空间位置信息如何设计适合视觉任务的Transformer模块下面这张表格对比了ViT与传统CNN的主要差异特性CNNViT感受野局部到全局渐进式从一开始就是全局位置信息处理通过卷积核隐式编码需要显式位置编码计算复杂度与图像尺寸线性相关与patch数量平方相关数据需求相对较少需要大量训练数据2. 图像分块嵌入从像素到视觉单词2.1 Patch Embedding实现原理ViT处理图像的第一步就是将图像分割成固定大小的块patch然后将每个块展平并映射到一个向量空间。这个过程称为Patch Embedding相当于把每个图像块转换为一个视觉单词。假设我们有一张224x224的RGB图像选择patch大小为16x16那么水平方向patch数量 224/16 14垂直方向patch数量 224/16 14总patch数量 14x14 196每个patch的维度是16x16x376816像素高 x 16像素宽 x 3颜色通道这768维的像素值就是我们的原始视觉单词。在实际实现中我们可以巧妙地使用一个卷积核大小和步长都等于patch大小的卷积层来完成这个操作class PatchEmbed(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, embed_dim768): super().__init__() img_size (img_size, img_size) # 确保图像是正方形 patch_size (patch_size, patch_size) self.img_size img_size self.patch_size patch_size self.grid_size (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) self.num_patches self.grid_size[0] * self.grid_size[1] # 关键用卷积实现分块和嵌入 self.proj nn.Conv2d(in_chans, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): B, C, H, W x.shape assert H self.img_size[0] and W self.img_size[1], \ f输入图像尺寸({H}*{W})与模型设定({self.img_size[0]}*{self.img_size[1]})不符 # [B, C, H, W] - [B, embed_dim, grid_h, grid_w] x self.proj(x) # [B, embed_dim, grid_h, grid_w] - [B, embed_dim, grid_h*grid_w] x x.flatten(2) # [B, embed_dim, grid_h*grid_w] - [B, grid_h*grid_w, embed_dim] x x.transpose(1, 2) return x2.2 类别token与位置编码在NLP的BERT模型中我们会在序列开头添加一个特殊的[CLS]token用于分类任务。ViT借鉴了这一思想引入了一个可学习的类别tokenclass token。这个token会与所有图像块交互最终作为整个图像的表征。同时由于Transformer本身不具备处理序列顺序的能力我们需要添加位置编码position embedding来保留图像块的空间位置信息。ViT中使用的是可学习的一维位置编码与patch embedding相加。# 初始化类别token和位置编码 self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed nn.Parameter(torch.zeros(1, num_patches 1, embed_dim)) # 在前向传播中添加它们 def forward(self, x): B x.shape[0] x self.patch_embed(x) # [B, num_patches, embed_dim] cls_tokens self.cls_token.expand(B, -1, -1) # [B, 1, embed_dim] x torch.cat((cls_tokens, x), dim1) # [B, num_patches1, embed_dim] x x self.pos_embed # 添加位置编码 return x我在实验中发现位置编码对模型性能影响很大。去掉位置编码后模型准确率会下降约3-5个百分点这说明Transformer确实需要显式的空间位置信息来处理视觉任务。3. Transformer编码器ViT的核心引擎3.1 多头注意力机制详解多头注意力Multi-Head Attention是Transformer的核心组件它允许模型同时关注输入序列的不同位置并学习丰富的空间关系。ViT中使用的多头注意力与原始Transformer基本相同但针对图像数据做了一些调整。让我们拆解多头注意力的实现步骤生成Q、K、V矩阵通过线性变换从输入序列得到查询(Query)、键(Key)和值(Value)矩阵分割多头将Q、K、V按注意力头数分割成多个子空间计算注意力权重在每个子空间中计算query和key的点积然后softmax归一化加权求和用注意力权重对value进行加权求和合并多头将所有头的输出拼接起来通过线性变换得到最终输出class Attention(nn.Module): def __init__(self, dim, num_heads8, qkv_biasFalse, attn_drop0., proj_drop0.): super().__init__() self.num_heads num_heads head_dim dim // num_heads self.scale head_dim ** -0.5 # 缩放因子 self.qkv nn.Linear(dim, dim * 3, biasqkv_bias) self.attn_drop nn.Dropout(attn_drop) self.proj nn.Linear(dim, dim) self.proj_drop nn.Dropout(proj_drop) def forward(self, x): B, N, C x.shape # 生成QKV [B, num_patches1, 3*embed_dim] qkv self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) q, k, v qkv[0], qkv[1], qkv[2] # 每个都是[B, num_heads, num_patches1, head_dim] # 计算注意力权重 attn (q k.transpose(-2, -1)) * self.scale # [B, num_heads, num_patches1, num_patches1] attn attn.softmax(dim-1) attn self.attn_drop(attn) # 加权求和 x (attn v).transpose(1, 2).reshape(B, N, C) # [B, num_patches1, embed_dim] # 输出投影 x self.proj(x) x self.proj_drop(x) return x在实际应用中我发现注意力头数通常设置为12base模型或16large模型每个头的维度保持在64左右效果较好。过多的注意力头会导致计算量增加而收益递减。3.2 MLP与层归一化Transformer编码器的另一个重要组件是MLP多层感知机也称为前馈网络Feed Forward Network。ViT中的MLP相对简单通常由两个全连接层组成中间使用GELU激活函数。class Mlp(nn.Module): def __init__(self, in_features, hidden_featuresNone, out_featuresNone, act_layernn.GELU, drop0.): super().__init__() out_features out_features or in_features hidden_features hidden_features or in_features self.fc1 nn.Linear(in_features, hidden_features) self.act act_layer() self.fc2 nn.Linear(hidden_features, out_features) self.drop nn.Dropout(drop) def forward(self, x): x self.fc1(x) x self.act(x) x self.drop(x) x self.fc2(x) x self.drop(x) return x层归一化Layer Normalization是稳定Transformer训练的关键。ViT在每个注意力层和MLP层之前都应用了层归一化class Block(nn.Module): def __init__(self, dim, num_heads, mlp_ratio4., qkv_biasFalse, drop0., attn_drop0.): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn Attention(dim, num_headsnum_heads, qkv_biasqkv_bias, attn_dropattn_drop, proj_dropdrop) self.norm2 nn.LayerNorm(dim) mlp_hidden_dim int(dim * mlp_ratio) self.mlp Mlp(in_featuresdim, hidden_featuresmlp_hidden_dim, dropdrop) def forward(self, x): # 残差连接注意力 x x self.attn(self.norm1(x)) # 残差连接MLP x x self.mlp(self.norm2(x)) return x在调试模型时我发现层归一化的位置很关键。ViT采用的是pre-norm结构归一化在注意力/MLP之前这比原始Transformer的post-norm更易于训练特别是在深层网络中。4. 构建完整的ViT模型4.1 组装Transformer编码器现在我们可以将多个Transformer块堆叠起来构建完整的ViT编码器。典型的ViT模型有几种配置模型类型层数隐藏层维度MLP大小注意力头数参数量ViT-Tiny1219276835.7MViT-Small123841536622MViT-Base1276830721286MViT-Large241024409616307Mclass VisionTransformer(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, num_classes1000, embed_dim768, depth12, num_heads12, mlp_ratio4., qkv_biasTrue, drop_rate0., attn_drop_rate0.): super().__init__() self.patch_embed PatchEmbed(img_size, patch_size, in_chans, embed_dim) num_patches self.patch_embed.num_patches self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed nn.Parameter(torch.zeros(1, num_patches 1, embed_dim)) self.pos_drop nn.Dropout(pdrop_rate) # 构建Transformer块 self.blocks nn.Sequential(*[ Block(dimembed_dim, num_headsnum_heads, mlp_ratiomlp_ratio, qkv_biasqkv_bias, dropdrop_rate, attn_dropattn_drop_rate) for i in range(depth)]) self.norm nn.LayerNorm(embed_dim) self.head nn.Linear(embed_dim, num_classes) if num_classes 0 else nn.Identity() # 初始化权重 nn.init.trunc_normal_(self.pos_embed, std0.02) nn.init.trunc_normal_(self.cls_token, std0.02) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): nn.init.trunc_normal_(m.weight, std0.02) if m.bias is not None: nn.init.zeros_(m.bias) def forward(self, x): x self.patch_embed(x) # [B, num_patches, embed_dim] cls_token self.cls_token.expand(x.shape[0], -1, -1) x torch.cat((cls_token, x), dim1) # [B, num_patches1, embed_dim] x x self.pos_embed x self.pos_drop(x) x self.blocks(x) # 通过所有Transformer块 x self.norm(x) # 只取cls_token对应的输出用于分类 x x[:, 0] x self.head(x) return x4.2 训练技巧与实战建议在实际训练ViT时有几个关键点需要注意数据增强ViT需要大量数据通常使用MixUp、CutMix等强数据增强学习率调度余弦退火学习率配合线性warmup效果很好优化器选择AdamW通常比Adam表现更好正则化权重衰减和dropout对防止过拟合很重要以下是一个简化的训练循环示例def train_epoch(model, loader, optimizer, criterion, device): model.train() total_loss, total_acc 0, 0 for inputs, targets in loader: inputs, targets inputs.to(device), targets.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, targets) loss.backward() optimizer.step() total_loss loss.item() _, predicted outputs.max(1) total_acc predicted.eq(targets).sum().item() return total_loss / len(loader), total_acc / len(loader.dataset) # 初始化模型 model VisionTransformer( img_size224, patch_size16, num_classes1000, embed_dim768, depth12, num_heads12 ).to(device) # 使用AdamW优化器 optimizer torch.optim.AdamW(model.parameters(), lr1e-3, weight_decay0.05) criterion nn.CrossEntropyLoss() # 训练循环 for epoch in range(100): train_loss, train_acc train_epoch(model, train_loader, optimizer, criterion, device) val_loss, val_acc validate(model, val_loader, criterion, device) print(fEpoch {epoch}: Train Loss {train_loss:.4f} Acc {train_acc:.2f}% | fVal Loss {val_loss:.4f} Acc {val_acc:.2f}%)在资源有限的情况下可以从预训练的ViT模型开始微调。HuggingFace的Transformers库提供了方便的接口from transformers import ViTForImageClassification model ViTForImageClassification.from_pretrained(google/vit-base-patch16-224)ViT的成功证明了Transformer架构在计算机视觉领域的强大潜力。虽然它需要比CNN更多的数据和计算资源但其优异的性能和可扩展性使其成为许多视觉任务的理想选择。