Vision Transformer实战教程：从零搭建ViT图像分类模型

我一个做了一段时间图像分类项目的朋友最近问我：现在大家都在聊 Vision Transformer，到底值不值得从 CNN 切过来？我的回答是，如果你手上正好有图像分类、目标检测或语义分割的需求，ViT 现在已经不是论文里的概念，而是完全可以直接跑在自己数据集上的成熟方案。这次的实战教程就是要带你走一遍完整的流程——从环境准备、数据处理、模型搭建到训练推理，全程用代码说话。

ViT（Vision Transformer）的核心思路一句话就能讲明白：把一张图片切成一个一个的小 patch，然后把每个 patch 当成 NLP 里的一个 token 送进 Transformer。听起来简单，但这里面涉及 patch embedding 的维度换算、位置编码的加法和注意力机制的计算流程，每一个环节在写代码时都有门道。我会先用一套易理解的代码把整个模型搭起来，再用 CIFAR-10 这类公开数据集做训练示例，最后告诉你如何替换成自己的数据集。

这套内容适合两类读者：一是对 Transformer 架构有基础了解、想在视觉任务上做尝试的开发者；二是被各种 ViT 变体绕晕、想直接上手跑通一个最小实现的新手。不管你是想给自己的分类任务换一个更强的骨干网络，还是想理解注意力机制在视觉里怎么发挥作用，这篇教程都会给你一条可参考的路线。

1. 上手之前：先搞懂 ViT 凭什么能打

1.1 为什么突然都在折腾 ViT

大概从 ViT 被提出以来，视觉领域的很多任务就开始多了一个选择方向。CNN 靠卷积核提取局部特征，一层层堆叠扩大感受野；而 ViT 走的是另一条路——它通过自注意力机制直接建模图像中任意两个位置之间的关系。也就是说，一张图中距离很远的两个像素，CNN 需要很多层才能建立联系，而 ViT 在一层之内就能完成全局交互。

这种思路带来的直接优势有两个。第一，在数据量足够大的情况下，ViT 的上限往往比同量级的 CNN 更高，因为注意力机制没有卷积那种局部归纳偏置的限制，能学到更灵活的特征表达。第二，ViT 的架构和 NLP 里的 Transformer 高度统一，意味着多模态模型可以共用一套骨干网络设计，工程上能省不少事。

但要注意一个前提：数据量足够大。ViT 不像 CNN 那样自带“附近像素更相关”的先验知识，所以小数据集上直接训练 ViT 很容易欠拟合。这也是为什么后来出现了 DeiT、Swin Transformer 等一系列改进方案，它们主要就是在解决“如何让 ViT 在小数据上也能训好”的问题。我们这次实战不会刻意避坑，而是把这个问题摊开来，让你知道什么时候该用 ViT、什么时候该用 ConvNeXt 或者 ResNet。

1.2 ViT 的核心流程拆解

把 ViT 的 forward 过程拆开来看，其实就是一个“图→序列→Transformer→分类”的流水线：

• 输入一张 H×W×C 的图片，先缩放到固定尺寸，比如 224×224。
• 按照设定的 patch size（比如 16×16）把图片切成 N 个小块，N = 224/16 × 224/16 = 196 个 patch。
• 每个 patch 展平成一个向量，再用一个线性层映射到 D 维空间，这就是 patch embedding，这一步等价于把每个 patch 变成 d_model 维的 token。
• 加上一个可学习的类别 token（class token），这个 token 的作用是吸收整张图的全局信息，最终用于分类。
• 加上位置编码（positional embedding），让模型知道每个 token 在原始图片中的相对位置。
• 送入 L 层 Transformer Encoder，每一层由 LayerNorm + Multi-Head Self-Attention + MLP 组成。
• 取 class token 对应的输出向量，过一层分类头，得到最终的类别概率。

这个流程里最容易写错的就是维度变换。一个 224×224×3 的图，切成 196 个 16×16×3 的 patch，展平后是 196×768，经过线性层变成 196×D，假设 D=768，那么和 class token 拼接后就是 197×768。这个 197 就是 Transformer 的序列长度，768 就是 embedding 维度。后面所有 attention 计算都要围绕这个形状来设计，稍不留神就 shape 不匹配报错。

2. 环境准备与数据集处理

2.1 跑通 ViT 需要哪些依赖

我推荐直接用 PyTorch 生态，因为它的动态图和社区生态对 ViT 这类模型非常友好。以下是我实际测试下来的依赖组合：

torch>=2.0 torchvision>=0.15 timm>=0.9 numpy>=1.24 matplotlib>=3.7 tqdm>=4.65

torch 和 torchvision 是训练和数据集加载的地基，timm 库很有用，它提供了很多现成的 Vision Transformer 变体、以及各种数据增强策略。你要是想从零实现一个 ViT 来加深理解，那么 timm 不是必须的；但如果想快速对比 ViT-Base/DeiT-Small/Swin-T 这些模型的性能，timm 能帮你节省大量时间。

安装方式很简单：

pip install torch torchvision timm numpy matplotlib tqdm

如果有 GPU，建议提前确认 CUDA 版本并安装对应的 torch 版本。用 CPU 训练 ViT 会非常煎熬，一个 epoch 可能要跑十几分钟甚至更久，所以有条件的话至少搞一张 6GB 以上显存的卡。

2.2 数据集格式与准备

为了方便复现，这里先用 CIFAR-10 作为示例数据集。CIFAR-10 共 10 类、6 万张 32×32 的彩色图片，直接用 torchvision 就能下载加载。但你切换到自己的数据集时，要遵循一个通用的目录组织方式：

dataset/ ├── train/ │ ├── class_0/ │ ├── class_1/ │ └── ... ├── val/ │ ├── class_0/ │ ├── class_1/ │ └── ...

PyTorch 的ImageFolder会自动按子目录名识别类别，前提是 train 和 val 下的子目录名称要一一对应。我自己在项目中换过好几个数据集，只要按这个格式整理，代码里只需要改数据路径和一个num_classes参数，其他地方都不用动。

加载数据的核心代码：

from torchvision import datasets, transforms from torch.utils.data import DataLoader transform_train = transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(), transforms.RandomCrop(224, padding=4), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) transform_val = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) train_dataset = datasets.ImageFolder(root='dataset/train', transform=transform_train) val_dataset = datasets.ImageFolder(root='dataset/val', transform=transform_val) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=4, pin_memory=True)

这里有一点值得注意：我用了Resize((224, 224))，但在很多任务里直接拉伸会破坏长宽比，导致物体形变。如果你的数据集中物体比例相对固定，直接用 Resize 问题不大；但如果物体形状差异很大，建议改成Resize((256, 256))后CenterCrop(224)，让模型看到的内容更自然。

2.3 数据增强与预处理

ViT 对数据增强的依赖比 CNN 更明显。因为 ViT 缺乏卷积那种平移等变性先验，如果你不做足够的增强，它非常容易在训练集上死记硬背，验证集上掉点。除了翻转、裁剪这些基础操作，我强烈建议加入以下几种增强：

• RandomResizedCrop：随机裁剪后缩放，模拟物体尺度变化。
• RandAugment：自动组合平移、旋转、颜色抖动等操作。
• MixUp / CutMix：在图像层面混合样本，让模型学习更平滑的决策边界。
• Label Smoothing：防止模型对训练标签过度自信。

timm 里提供了RandAugment和Mixup的实现，直接调就行：

from timm.data import rand_augment_transform from timm.data.mixup import Mixup transform_timm = rand_augment_transform('rand-m9-mstd0.5', {}) mixup_fn = Mixup(mixup_alpha=0.8, cutmix_alpha=1.0, label_smoothing=0.1, num_classes=10)

我个人的体会是，CIFAR-10 这种小数据集上，RandAugment 加上 Mixup 能把准确率稳定提升 2 到 3 个百分点。数据增强库的版本差异会影响结果，建议锁定稳定版本。

3. 模型搭建与代码解析

3.1 从零搭建 ViT 主模型

这里我会手写一个完整的 ViT 实现，而不是直接调 timm 里的现成模型。这样做的原因是：只有当你真正逐行写过 patch embedding、positional embedding 和 attention 的实现，你才能在遇到维度错误时快速定位问题。理解比调用更重要。

import torch import torch.nn as nn import math class PatchEmbed(nn.Module): def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768): super().__init__() self.num_patches = (img_size // patch_size) ** 2 self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size=patch_size, stride=patch_size) def forward(self, x): # x: [B, 3, 224, 224] -> [B, 768, 14, 14] -> [B, 196, 768] x = self.proj(x) x = x.flatten(2) x = x.transpose(1, 2) return x

这段代码里，nn.Conv2d的卷积核大小和 stride 都等于 patch_size，相当于把每个 patch 一次性投影成 embed_dim 维向量。用卷积实现 patch embedding 在计算上比手动切片更高效，这也是 PyTorch 实践中的常用技巧。

然后是实现 Multi-Head Self-Attention 和 Transformer Encoder Block：

class Attention(nn.Module): def __init__(self, dim, num_heads=8, qkv_bias=True): super().__init__() self.num_heads = num_heads self.head_dim = dim // num_heads self.scale = self.head_dim ** -0.5 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.proj = nn.Linear(dim, dim) def forward(self, x): B, N, C = x.shape qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] attn = (q @ k.transpose(-2, -1)) * self.scale attn = attn.softmax(dim=-1) x = (attn @ v).transpose(1, 2).reshape(B, N, C) x = self.proj(x) return x

注意力计算的流程是：q、k、v 都是从同一个输入通过线性层得到的，然后计算 q 和 k 的点积，除以 head_dim 的平方根做缩放，softmax 得到权重，再对 v 做加权求和。缩放的原因很实际：如果不除以根号 head_dim，当维度增大时点积结果会很大，softmax 会进入饱和区，梯度变得很小，训练就不稳定。

Transformer Encoder Block：

class Mlp(nn.Module): def __init__(self, in_features, hidden_features=None, out_features=None, drop=0.1): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features * 4 self.fc1 = nn.Linear(in_features, hidden_features) self.act = nn.GELU() 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 class Block(nn.Module): def __init__(self, dim, num_heads, mlp_ratio=4.0, drop=0.1): super().__init__() self.norm1 = nn.LayerNorm(dim) self.attn = Attention(dim, num_heads) self.norm2 = nn.LayerNorm(dim) self.mlp = Mlp(dim, int(dim * mlp_ratio), drop=drop) def forward(self, x): x = x + self.attn(self.norm1(x)) x = x + self.mlp(self.norm2(x)) return x

这里用的是 Pre-LN 结构，即先做 LayerNorm 再进 Attention/MLP，残差连接保留原始信息。相比 Post-LN，Pre-LN 在训练早期更稳定，是 ViT 能成功训练的关键设计之一。

最终的完整 ViT：

class ViT(nn.Module): def __init__(self, img_size=224, patch_size=16, in_channels=3, num_classes=10, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4.0, drop=0.1): super().__init__() self.patch_embed = PatchEmbed(img_size, patch_size, in_channels, 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(p=drop) self.blocks = nn.Sequential(*[ Block(embed_dim, num_heads, mlp_ratio, drop) for _ in range(depth) ]) self.norm = nn.LayerNorm(embed_dim) self.head = nn.Linear(embed_dim, num_classes) # 初始化 nn.init.trunc_normal_(self.pos_embed, std=0.02) nn.init.trunc_normal_(self.cls_token, std=0.02) self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): nn.init.trunc_normal_(m.weight, std=0.02) if m.bias is not None: nn.init.zeros_(m.bias) elif isinstance(m, nn.LayerNorm): nn.init.zeros_(m.bias) nn.init.ones_(m.weight) def forward(self, x): B = x.shape[0] x = self.patch_embed(x) cls_tokens = self.cls_token.expand(B, -1, -1) x = torch.cat((cls_tokens, x), dim=1) x = x + self.pos_embed x = self.pos_drop(x) x = self.blocks(x) x = self.norm(x) x = x[:, 0] x = self.head(x) return x

这段代码里有一个细节值得解释：self.cls_token.expand(B, -1, -1)用的是 expand 而不是 repeat。expand 只改变视图不复制内存，效率更高；但 expand 后这个张量不是连续的，某些操作可能会隐式地复制，实际上这里因为只用于 cat 和加法，性能没有影响。如果你之后再对这个张量做view操作，就要小心先contiguous()。

3.2 Patch Embedding 到底做了什么

如果从数学角度理解，Patch Embedding 其实就是一个形状为(patch_size * patch_size * in_channels, embed_dim)的线性层，作用是把每个 patch 展平后的向量映射到高维空间。用卷积实现时，输入输出形状的变化更直观：

假设输入是[B, 3, 224, 224]，patch_size=16，embed_dim=768。卷积层输出是[B, 768, 14, 14]，这 14×14=196 个位置就对应 196 个 patch。flatten(2)变成[B, 768, 196]，transpose(1, 2)变成[B, 196, 768]，每个 patch 就是一个长度为 768 的 token。

为什么需要 class token？这是 ViT 设计里的一个巧妙之处。Transformer 本身处理的是序列，序列里每个 token 都携带局部信息，但没有一个显式的“全局汇总”位置。class token 作为额外的一个 token，通过在自注意力层中与所有 patch token 交互，可以聚合全局信息，最终在输出层只取这个 token 的向量做分类。取了 class token 后，相当于让模型自己决定“哪些全局信息对分类最有用”。

3.3 位置编码与 Attention 的核心

ViT 没有循环结构和卷积结构，所以模型根本不知道 token 的先后顺序。位置编码就是给每个 token 注入位置信息的手段。ViT 原文里用的是可学习的 1D 位置编码，它是一组形状为(num_patches + 1, embed_dim)的参数，训练前随机初始化，训练中随模型一起更新。

使用可学习位置编码而不是正弦位置编码，原因在于视觉任务的 patch 排列天然是二维的，但 ViT 直接拉平成一维序列；可学习编码可以让模型自己决定如何编码这种空间关系，灵活性更高。虽然有人提过 2D 位置编码能小幅度提升性能，但实践中 1D 可学习编码已经够用，而且实现更简单。

Attention 的核心就是刚才代码里的那几行：

attn = (q @ k.transpose(-2, -1)) * scale attn = attn.softmax(dim=-1) x = (attn @ v)

如果你把 q 和 k 的点积理解成“相似度打分”，那么这个分数决定了每个 token 在聚合信息时该关注其他 token 多少。softmax 把分数变成和为 1 的权重分布，最后加权求和 v，得到当前 token 的新表示。多头就是把 embedding 维度切成 num_heads 份，每个头独立做上述计算，最后拼接起来。多头的意义在于让模型能同时关注不同类型的语义关系，比如一个头关注局部纹理，另一个头关注全局形状。

4. 训练循环与关键细节

4.1 损失函数和优化器的选择逻辑

分类任务最常用的损失函数是 CrossEntropyLoss，PyTorch 里直接调用即可。但如果你用了 MixUp 或 Label Smoothing，就不能直接用它，需要稍微改造：

class SoftTargetCrossEntropy(nn.Module): def __init__(self): super().__init__() def forward(self, x, target): # target 是软标签，形状与 x 相同 loss = torch.sum(-target * F.log_softmax(x, dim=-1), dim=-1) return loss.mean()

关于优化器，ViT 训练通常有两个主流选择：一是 AdamW，收敛快，泛化好；二是 SGD with momentum，调好之后效果也不差，但需要更多 epoch。我个人在 ViT 上更推荐 AdamW，配合余弦退火学习率调度器，稳定性和最终精度都很不错。原因在于 ViT 的优化空间更加崎岖，Adam 的逐参数自适应学习率能更快穿越平坦区域。

有个关键参数要特别注意：weight decay。ViT 对 weight decay 比较敏感，太小容易过拟合，太大则可能欠拟合。CIFAR-10、100 epoch 的训练任务，weight decay 取 5e-2 到 1e-1 之间是一个经验区间。如果让你的 ViT 在小数据集上学习，建议从 5e-2 开始调。

4.2 超参数设置

ViT 家族里没有一个通用的最优超参数组合，但在中小型数据集上，一套比较稳的基准配置是：

embed_dim 与 num_heads 的关系需要注意：embed_dim 必须能被 num_heads 整除。一般每个头的维度 head_dim 取 64 左右是一种常见设置，embed_dim=384、num_heads=6 时 head_dim=64，刚好对齐。

4.3 训练中的监控与保存

训练循环里除了跑 forward 和 backward，我强烈建议每 5 个 epoch 做一次验证，并且记录以下指标：

• 训练集 loss 和准确率
• 验证集 loss 和准确率
• 当前 learning rate
• GPU 显存占用
• 单个 epoch 耗时

验证逻辑：

def evaluate(model, val_loader, criterion, device): model.eval() total_loss = 0.0 correct = 0 total = 0 with torch.no_grad(): for images, labels in val_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) loss = criterion(outputs, labels) total_loss += loss.item() * images.size(0) _, predicted = outputs.max(1) total += labels.size(0) correct += predicted.eq(labels).sum().item() avg_loss = total_loss / total acc = 100.0 * correct / total return avg_loss, acc

保存模型时，不要只保存 state_dict，最好把 optimizer state 也存下来，方便中断后继续训练：

checkpoint = { 'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict(), 'epoch': epoch, 'best_acc': best_acc, } torch.save(checkpoint, 'checkpoint.pth')

训练主循环里有个容易忽略的细节：在训练模式下一定要调用model.train()，在验证模式下调用model.eval()。看似简单，但Dropout和LayerNorm（实际是BatchNorm类）在两种模式下的行为完全不同。ViT 里主要用的是 LayerNorm，它不受model.eval()影响，但 Dropout 会。如果你漏了model.train()，Dropout 不生效，训练精度会虚高，推理时反而掉点严重。

5. 踩坑实录与问题排查

5.1 显存爆掉怎么办

ViT 的显存占用本来就不低，尤其你把batch_size设为 128、分辨率拉满时，OOM 是常有的事。这里我按优先级给你几个实操方案：

• 减小 batch size，这是最直接的办法，但千万别忘了同步调整学习率。经验做法是把学习率乘以 batch size 的变化比例，比如 batch size 从 64 减到 32，学习率大致减半。
• 打开混合精度训练，用torch.cuda.amp可以显著降低显存占用，同时训练速度可能提升 1.5 倍。
• 用梯度累积模拟大 batch，比如每隔 4 个 step 做一次 backward，效果接近直接使用 4 倍 batch size。
• 检查是否有不必要的张量被保留在了计算图中。比如验证时用with torch.no_grad()包裹，能省下大量显存。

混合精度的关键代码：

scaler = torch.cuda.amp.GradScaler() for images, labels in train_loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() with torch.cuda.amp.autocast(): outputs = model(images) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

要注意：用 AMP 时，如果 loss 出现 NaN，通常在网络前向或损失函数构造中已经存在问题，AMP 不会自动修复。另外，不同 GPU 对 AMP 的支持程度不同，老显卡可能不支持某些半精度算子，遇到报错换回 FP32 即可。

5.2 训练 loss 不下降怎么排查

如果你的 ViT 训练 loss 一直不下降，这里有一个可执行的排查顺序：

• 检查数据增强是不是过度了。比如 MixUp 的 alpha 设得太大，会过度混合样本，模型很难拟合。先关掉所有高级增强，用最朴素的 Resize + Normalize 跑 20 个 epoch，如果 loss 能下降，说明增强强度需要降低。
• 检查学习率是否过大或过小。过大导致 loss 抖动不收敛，过小则下降缓慢得像蜗牛。用学习率探测法在一个小批量上跑 100 step，观察 loss 曲线，找到最合适的数量级。
• 检查标签是否有问题。用一小部分数据手动过一遍模型，查看输出的 logits 分布是否合理。比如类别数是否等于你的标签范围，标签有没有从 1 开始而不是从 0，这类低级错误排查起来很费时间。
• 检查位置编码和 cls token 的初始化。位置编码随机初始化是正常的，但如果位置编码范围太大或太小，会让注意力分布异常。确认trunc_normal_的 std 是 0.02 而不是 0.2。

5.3 常见问题速查表

5.4 加快训练速度的实用手段

ViT 训练耗时是出了名的长，这里分享几个我实测有效的提速方案。

第一，用 F.cross_entropy 集成 label smoothing，避免自己实现一遍各种细节导致版本差异：

loss = nn.CrossEntropyLoss(label_smoothing=0.1)

第二，把 dataloader 的 num_workers 调到合适的值。一般来说 4~8 个 num_workers 对单机训练已经够用，再多不一定更快，反而可能拖慢主进程。可以做一个简单的 benchmark，对比不同 num_workers 下每个 epoch 的耗时。

第三，使用torch.backends.cudnn.benchmark = True，这个开关让 cuDNN 针对固定输入尺寸自动选择最快的卷积算法，对 CNN 和 ViT 都有帮助。

第四，如果你的 GPU 支持，直接上 Flash Attention。Flash Attention 通过分块计算把注意力矩阵的中间结果不落显存，显存占用和计算速度都大幅优化。在 PyTorch 中可以通过torch.nn.functional.scaled_dot_product_attention直接使用，无需额外安装依赖。

6. 我的一些实操体会

折腾了这么久 ViT，我自己最大的感受是：不要神话它。在 CIFAR-10 这种 32×32 的小数据上，从零训练 ViT 其实很难打赢同量级的 ResNet；但在 ImageNet 这种量级或者迁移学习场景下，ViT 的优势就体现出来了。所以当你准备用自己的数据集做分类时，先用小规模实验判断一下：数据量是否有十万级以上？如果不够，建议先跑通流程，再用迁移学习加载预训练权重——这比从零训练高效得多。

后续如果想往更深走，可以尝试给 ViT 加入更高效的位置编码（如 2D 可学习位置编码），或者把 patch embedding 换成重叠 patch 的变体，甚至可以试着把 ViT 接到检测框架里做端到端的目标检测。总之，先把基础版本跑通，再逐步往上叠改进点，才算把 ViT 真正吃透。

最后分享一个小技巧：把模型和数据集理顺之后，日志系统值得多花一点功夫。我习惯把每个 epoch 的 loss、acc、lr 都记录到 CSV 文件里，训练结束后直接画曲线。它能帮你快速判断是过拟合还是欠拟合，也能让你在调整超参数时更有方向。训练自己的数据集，从来都是个工程问题，不只是模型结构的堆叠。