Vision Transformer实战从零搭建PyTorch图像分类模型当卷积神经网络CNN长期统治计算机视觉领域时Transformer架构的横空出世彻底改变了游戏规则。2020年Google Research提出的Vision TransformerViT首次证明纯Transformer架构在图像分类任务上可以超越CNN。本文将带您从零开始用PyTorch实现一个完整的ViT模型深入解析每个关键组件的实现细节。1. ViT核心原理与架构设计ViT的核心思想是将图像处理为序列数据。与传统CNN逐层提取局部特征不同ViT将图像分割为固定大小的图块patches将这些图块线性投影后加上位置编码送入标准的Transformer编码器进行处理。关键创新点图像分块处理将224×224图像分割为16×16的图块共196个每个图块展平为768维向量可学习分类标记在序列开头添加特殊[CLS]标记其最终状态用于分类位置编码由于Transformer本身不包含位置信息必须显式添加位置编码实验表明当训练数据足够大时如JFT-300MViT的性能显著优于同规模的CNN模型尤其在捕捉长距离依赖关系方面表现突出2. 环境准备与数据预处理在开始编码前我们需要配置开发环境并准备数据集pip install torch1.12.0 torchvision0.13.0 pip install numpy matplotlib tqdm使用CIFAR-10数据集进行演示但需要注意原始ViT设计输入尺寸为224×224from torchvision import transforms, datasets # 数据增强与归一化 train_transform transforms.Compose([ transforms.Resize(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # 加载CIFAR-10数据集 train_set datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtrain_transform) train_loader torch.utils.data.DataLoader(train_set, batch_size32, shuffleTrue, num_workers2)3. ViT核心组件实现3.1 图块嵌入层将图像分割为图块并线性投影class PatchEmbedding(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, embed_dim768): super().__init__() self.img_size img_size self.patch_size patch_size self.n_patches (img_size // patch_size) ** 2 self.proj nn.Conv2d(in_chans, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): # x: [B, C, H, W] x self.proj(x) # [B, E, H/P, W/P] x x.flatten(2) # [B, E, N] x x.transpose(1, 2) # [B, N, E] return x3.2 位置编码与分类标记class ViTEmbeddings(nn.Module): def __init__(self, config): super().__init__() self.patch_embeddings PatchEmbedding( img_sizeconfig.img_size, patch_sizeconfig.patch_size, embed_dimconfig.hidden_size ) # 可学习的分类标记 self.cls_token nn.Parameter(torch.zeros(1, 1, config.hidden_size)) # 位置编码 self.position_embeddings nn.Parameter( torch.zeros(1, self.patch_embeddings.n_patches 1, config.hidden_size) ) self.dropout nn.Dropout(config.dropout_rate) def forward(self, x): batch_size x.shape[0] # 生成图块嵌入 patch_embeds self.patch_embeddings(x) # [B, N, E] # 扩展分类标记到batch维度 cls_tokens self.cls_token.expand(batch_size, -1, -1) # [B, 1, E] # 拼接分类标记 embeddings torch.cat((cls_tokens, patch_embeds), dim1) # [B, N1, E] # 添加位置编码 embeddings self.position_embeddings return self.dropout(embeddings)3.3 多头注意力机制实现class MultiHeadAttention(nn.Module): def __init__(self, config): super().__init__() self.num_heads config.num_heads self.head_dim config.hidden_size // config.num_heads self.query nn.Linear(config.hidden_size, config.hidden_size) self.key nn.Linear(config.hidden_size, config.hidden_size) self.value nn.Linear(config.hidden_size, config.hidden_size) self.out nn.Linear(config.hidden_size, config.hidden_size) self.dropout nn.Dropout(config.attention_dropout_rate) def forward(self, x): batch_size, seq_len, embed_dim x.shape # 线性投影得到Q, K, V q self.query(x) # [B, N, E] k self.key(x) # [B, N, E] v self.value(x) # [B, N, E] # 重塑为多头形式 q q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) v v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # 计算注意力分数 scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) attn_weights torch.softmax(scores, dim-1) attn_weights self.dropout(attn_weights) # 应用注意力权重 context torch.matmul(attn_weights, v) context context.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim) # 最终线性层 output self.out(context) return output4. 完整ViT模型集成将所有组件组合成完整模型class ViT(nn.Module): def __init__(self, config): super().__init__() self.config config # 嵌入层 self.embeddings ViTEmbeddings(config) # Transformer编码器 encoder_layer nn.TransformerEncoderLayer( d_modelconfig.hidden_size, nheadconfig.num_heads, dim_feedforwardconfig.mlp_dim, dropoutconfig.dropout_rate ) self.encoder nn.TransformerEncoder(encoder_layer, num_layersconfig.num_layers) # 分类头 self.classifier nn.Linear(config.hidden_size, config.num_classes) def forward(self, x): # 嵌入层 embeddings self.embeddings(x) # Transformer编码 encoded self.encoder(embeddings) # 使用分类标记进行分类 cls_token encoded[:, 0, :] logits self.classifier(cls_token) return logits5. 模型训练与优化配置训练参数并实现训练循环# 模型配置 class ViTConfig: img_size 224 patch_size 16 hidden_size 768 num_heads 12 mlp_dim 3072 num_layers 12 dropout_rate 0.1 attention_dropout_rate 0.0 num_classes 10 # 初始化模型 config ViTConfig() model ViT(config).to(device) # 定义损失函数和优化器 criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr3e-5, weight_decay0.01) # 训练循环 for epoch in range(10): model.train() running_loss 0.0 for i, (inputs, labels) in enumerate(train_loader): inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() if i % 100 99: print(fEpoch {epoch1}, Batch {i1}, Loss: {running_loss/100:.4f}) running_loss 0.06. 关键问题与解决方案问题1位置编码如何影响模型性能ViT使用可学习的位置编码而非固定的正弦编码。实验表明位置编码类型Top-1准确率可学习编码78.5%正弦编码77.9%无位置编码65.3%问题2如何处理不同尺寸的输入图像可以通过调整图块大小或使用金字塔结构# 动态调整图块大小 def adjust_patch_size(img_size, target_patches196): patch_size int(math.sqrt(img_size[0]*img_size[1]/target_patches)) return patch_size问题3如何加速ViT训练使用混合精度训练采用梯度检查点技术使用更大的batch size配合学习率warmup# 混合精度训练示例 scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()7. 模型评估与可视化实现注意力可视化可以帮助理解模型工作原理def visualize_attention(image, model, layer_idx0, head_idx0): # 注册钩子获取注意力权重 attention_weights [] def hook(module, input, output): attention_weights.append(output[1]) handle model.encoder.layers[layer_idx].self_attn.register_forward_hook(hook) # 前向传播 with torch.no_grad(): _ model(image.unsqueeze(0).to(device)) handle.remove() # 可视化特定头的注意力 attn attention_weights[0][0, head_idx, 0, 1:].reshape(14, 14) plt.imshow(attn.cpu().numpy(), cmaphot) plt.colorbar() plt.show()8. 进阶优化技巧技巧1知识蒸馏使用大型ViT模型如ViT-Large作为教师模型蒸馏到小型ViT# 定义蒸馏损失 def distillation_loss(student_logits, teacher_logits, temp2.0): soft_teacher F.softmax(teacher_logits/temp, dim-1) soft_student F.log_softmax(student_logits/temp, dim-1) return F.kl_div(soft_student, soft_teacher, reductionbatchmean) * (temp**2)技巧2混合架构结合CNN和Transformer的优势class HybridViT(nn.Module): def __init__(self): super().__init__() # 使用CNN提取低级特征 self.cnn_backbone torchvision.models.resnet18(pretrainedTrue) self.cnn_backbone nn.Sequential(*list(self.cnn_backbone.children())[:-2]) # ViT处理高级特征 self.vit ViT(config) def forward(self, x): cnn_features self.cnn_backbone(x) b, c, h, w cnn_features.shape cnn_features cnn_features.view(b, c, h*w).transpose(1, 2) return self.vit(cnn_features)技巧3数据高效训练使用CutMix数据增强应用MixUp正则化采用RandAugment自动增强策略# CutMix实现示例 def cutmix_data(x, y, alpha1.0): lam np.random.beta(alpha, alpha) rand_index torch.randperm(x.size(0)) target_a y target_b y[rand_index] bbx1, bby1, bbx2, bby2 rand_bbox(x.size(), lam) x[:, :, bbx1:bbx2, bby1:bby2] x[rand_index, :, bbx1:bbx2, bby1:bby2] lam 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (x.size()[-1] * x.size()[-2])) return x, target_a, target_b, lam
Vision Transformer实战：从零开始搭建ViT图像分类模型（PyTorch版）
Vision Transformer实战从零搭建PyTorch图像分类模型当卷积神经网络CNN长期统治计算机视觉领域时Transformer架构的横空出世彻底改变了游戏规则。2020年Google Research提出的Vision TransformerViT首次证明纯Transformer架构在图像分类任务上可以超越CNN。本文将带您从零开始用PyTorch实现一个完整的ViT模型深入解析每个关键组件的实现细节。1. ViT核心原理与架构设计ViT的核心思想是将图像处理为序列数据。与传统CNN逐层提取局部特征不同ViT将图像分割为固定大小的图块patches将这些图块线性投影后加上位置编码送入标准的Transformer编码器进行处理。关键创新点图像分块处理将224×224图像分割为16×16的图块共196个每个图块展平为768维向量可学习分类标记在序列开头添加特殊[CLS]标记其最终状态用于分类位置编码由于Transformer本身不包含位置信息必须显式添加位置编码实验表明当训练数据足够大时如JFT-300MViT的性能显著优于同规模的CNN模型尤其在捕捉长距离依赖关系方面表现突出2. 环境准备与数据预处理在开始编码前我们需要配置开发环境并准备数据集pip install torch1.12.0 torchvision0.13.0 pip install numpy matplotlib tqdm使用CIFAR-10数据集进行演示但需要注意原始ViT设计输入尺寸为224×224from torchvision import transforms, datasets # 数据增强与归一化 train_transform transforms.Compose([ transforms.Resize(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # 加载CIFAR-10数据集 train_set datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtrain_transform) train_loader torch.utils.data.DataLoader(train_set, batch_size32, shuffleTrue, num_workers2)3. ViT核心组件实现3.1 图块嵌入层将图像分割为图块并线性投影class PatchEmbedding(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, embed_dim768): super().__init__() self.img_size img_size self.patch_size patch_size self.n_patches (img_size // patch_size) ** 2 self.proj nn.Conv2d(in_chans, embed_dim, kernel_sizepatch_size, stridepatch_size) def forward(self, x): # x: [B, C, H, W] x self.proj(x) # [B, E, H/P, W/P] x x.flatten(2) # [B, E, N] x x.transpose(1, 2) # [B, N, E] return x3.2 位置编码与分类标记class ViTEmbeddings(nn.Module): def __init__(self, config): super().__init__() self.patch_embeddings PatchEmbedding( img_sizeconfig.img_size, patch_sizeconfig.patch_size, embed_dimconfig.hidden_size ) # 可学习的分类标记 self.cls_token nn.Parameter(torch.zeros(1, 1, config.hidden_size)) # 位置编码 self.position_embeddings nn.Parameter( torch.zeros(1, self.patch_embeddings.n_patches 1, config.hidden_size) ) self.dropout nn.Dropout(config.dropout_rate) def forward(self, x): batch_size x.shape[0] # 生成图块嵌入 patch_embeds self.patch_embeddings(x) # [B, N, E] # 扩展分类标记到batch维度 cls_tokens self.cls_token.expand(batch_size, -1, -1) # [B, 1, E] # 拼接分类标记 embeddings torch.cat((cls_tokens, patch_embeds), dim1) # [B, N1, E] # 添加位置编码 embeddings self.position_embeddings return self.dropout(embeddings)3.3 多头注意力机制实现class MultiHeadAttention(nn.Module): def __init__(self, config): super().__init__() self.num_heads config.num_heads self.head_dim config.hidden_size // config.num_heads self.query nn.Linear(config.hidden_size, config.hidden_size) self.key nn.Linear(config.hidden_size, config.hidden_size) self.value nn.Linear(config.hidden_size, config.hidden_size) self.out nn.Linear(config.hidden_size, config.hidden_size) self.dropout nn.Dropout(config.attention_dropout_rate) def forward(self, x): batch_size, seq_len, embed_dim x.shape # 线性投影得到Q, K, V q self.query(x) # [B, N, E] k self.key(x) # [B, N, E] v self.value(x) # [B, N, E] # 重塑为多头形式 q q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) k k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) v v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) # 计算注意力分数 scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) attn_weights torch.softmax(scores, dim-1) attn_weights self.dropout(attn_weights) # 应用注意力权重 context torch.matmul(attn_weights, v) context context.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim) # 最终线性层 output self.out(context) return output4. 完整ViT模型集成将所有组件组合成完整模型class ViT(nn.Module): def __init__(self, config): super().__init__() self.config config # 嵌入层 self.embeddings ViTEmbeddings(config) # Transformer编码器 encoder_layer nn.TransformerEncoderLayer( d_modelconfig.hidden_size, nheadconfig.num_heads, dim_feedforwardconfig.mlp_dim, dropoutconfig.dropout_rate ) self.encoder nn.TransformerEncoder(encoder_layer, num_layersconfig.num_layers) # 分类头 self.classifier nn.Linear(config.hidden_size, config.num_classes) def forward(self, x): # 嵌入层 embeddings self.embeddings(x) # Transformer编码 encoded self.encoder(embeddings) # 使用分类标记进行分类 cls_token encoded[:, 0, :] logits self.classifier(cls_token) return logits5. 模型训练与优化配置训练参数并实现训练循环# 模型配置 class ViTConfig: img_size 224 patch_size 16 hidden_size 768 num_heads 12 mlp_dim 3072 num_layers 12 dropout_rate 0.1 attention_dropout_rate 0.0 num_classes 10 # 初始化模型 config ViTConfig() model ViT(config).to(device) # 定义损失函数和优化器 criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr3e-5, weight_decay0.01) # 训练循环 for epoch in range(10): model.train() running_loss 0.0 for i, (inputs, labels) in enumerate(train_loader): inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() if i % 100 99: print(fEpoch {epoch1}, Batch {i1}, Loss: {running_loss/100:.4f}) running_loss 0.06. 关键问题与解决方案问题1位置编码如何影响模型性能ViT使用可学习的位置编码而非固定的正弦编码。实验表明位置编码类型Top-1准确率可学习编码78.5%正弦编码77.9%无位置编码65.3%问题2如何处理不同尺寸的输入图像可以通过调整图块大小或使用金字塔结构# 动态调整图块大小 def adjust_patch_size(img_size, target_patches196): patch_size int(math.sqrt(img_size[0]*img_size[1]/target_patches)) return patch_size问题3如何加速ViT训练使用混合精度训练采用梯度检查点技术使用更大的batch size配合学习率warmup# 混合精度训练示例 scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): outputs model(inputs) loss criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()7. 模型评估与可视化实现注意力可视化可以帮助理解模型工作原理def visualize_attention(image, model, layer_idx0, head_idx0): # 注册钩子获取注意力权重 attention_weights [] def hook(module, input, output): attention_weights.append(output[1]) handle model.encoder.layers[layer_idx].self_attn.register_forward_hook(hook) # 前向传播 with torch.no_grad(): _ model(image.unsqueeze(0).to(device)) handle.remove() # 可视化特定头的注意力 attn attention_weights[0][0, head_idx, 0, 1:].reshape(14, 14) plt.imshow(attn.cpu().numpy(), cmaphot) plt.colorbar() plt.show()8. 进阶优化技巧技巧1知识蒸馏使用大型ViT模型如ViT-Large作为教师模型蒸馏到小型ViT# 定义蒸馏损失 def distillation_loss(student_logits, teacher_logits, temp2.0): soft_teacher F.softmax(teacher_logits/temp, dim-1) soft_student F.log_softmax(student_logits/temp, dim-1) return F.kl_div(soft_student, soft_teacher, reductionbatchmean) * (temp**2)技巧2混合架构结合CNN和Transformer的优势class HybridViT(nn.Module): def __init__(self): super().__init__() # 使用CNN提取低级特征 self.cnn_backbone torchvision.models.resnet18(pretrainedTrue) self.cnn_backbone nn.Sequential(*list(self.cnn_backbone.children())[:-2]) # ViT处理高级特征 self.vit ViT(config) def forward(self, x): cnn_features self.cnn_backbone(x) b, c, h, w cnn_features.shape cnn_features cnn_features.view(b, c, h*w).transpose(1, 2) return self.vit(cnn_features)技巧3数据高效训练使用CutMix数据增强应用MixUp正则化采用RandAugment自动增强策略# CutMix实现示例 def cutmix_data(x, y, alpha1.0): lam np.random.beta(alpha, alpha) rand_index torch.randperm(x.size(0)) target_a y target_b y[rand_index] bbx1, bby1, bbx2, bby2 rand_bbox(x.size(), lam) x[:, :, bbx1:bbx2, bby1:bby2] x[rand_index, :, bbx1:bbx2, bby1:bby2] lam 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (x.size()[-1] * x.size()[-2])) return x, target_a, target_b, lam
