Vision Transformer（ViT）原理与PyTorch从零复现：图像分类实战指南

如果你接触深度学习有一段时间大概率听过 Vision TransformerViT这个名字。2020 年 Google Brain 团队发表论文An Image is Worth 16x16 Words直接把 Transformer 从 NLP 搬到了图像分类任务上用“图像切块 标准 Encoder”的方式在 ImageNet 上打平甚至超过当时的 CNN 模型。从那以后CLIP、SAM、DINOv2 等一批视觉大模型底层骨干基本都是 ViT 或其变体。可以说搞懂 ViT是理解当前视觉大模型体系的底线要求。这次这篇文章不准备只讲概念。我们会把 ViT 的核心算法原理拆成五六个环节然后用 PyTorch 从零写一个完整可训练的 VisionTransformer 分类模型再放到 CIFAR-10 上实际跑一遍训练和推理。整个过程不需要任何预训练权重代码按顺序复制就能运行。适合正在入门深度学习、准备复现论文源码或者想了解图像 Transformer 内部结构的读者。1. 核心能力速览项目说明模型类型视觉 Transformer 图像分类模型论文来源2020 年 Google BrainAn Image is Worth 16x16 Words核心思想图像切成 Patch映射为 Token送入标准 Transformer Encoder最常用配置ViT-Base/16patch16embed_dim768depth12heads12参数量级Base/16 通常在 8600 万参数左右具体以实际代码统计为准依赖环境Python PyTorch torchvisionCPU 可跑通小模型批量支持天然支持 Batch 输入适合批量图像推理接口 API模型本身不提供可自行封装为 HTTP 服务本文实操从零复现、CIFAR-10 训练、单张图片推理、性能观察适用读者很简单刚学完 CNN、想理解 Transformer 怎么处理图像的论文里看到了“Patch Embedding”“CLS Token”但不清楚具体代码怎么写的准备把 ViT 接到自己项目里做分类或特征提取的。2. 为什么视觉任务会转向 TransformerCNN 统治了视觉任务很多年核心原因是它自带两个归纳偏置局部性和平移等变性。卷积核只看小邻域参数共享这让 CNN 在中小数据集上很省力。但反过来CNN 的感受野是逐层扩大的想捕捉真正长距离的像素关系往往需要很深的网络。ViT 直接放弃了这两个偏置。它把图像当成一组 Token让自注意力机制在整张图范围内计算任意两个 Patch 之间的相关性。第一层 Encoder 就能看到全图这种全局建模能力在语义分割、目标检测、多模态对齐等复杂任务上优势明显。从工程角度看Transformer 架构还有一个实用价值通用性。NLP 用 Transformer视觉用 Transformer音频也可以切片成 Token 用 Transformer。统一结构意味着可以共享推理框架、分布式训练方案也方便做多模态联合训练。这也是现在 AI 大模型动辄“视觉语言一体”的底子。更具体地说ViT 在超大数据集上出现了类似语言模型的尺度效应模型越大数据越多效果越稳定这个特性让它在现代大模型训练流程里占据了绝对主力位置。理解 ViT不是为了替代 CNN而是为了理解大模型视觉模块的通用语言。3. ViT 算法原理拆解ViT 的完整前向流程可以概括为五步图像切块、线性映射得到 Token、加入位置编码、经过多层 Transformer Encoder、用分类 Token 输出结果。3.1 图像切块与 Patch Embedding假设输入是一张 224×224×3 的 RGB 图像设置 patch_size16。那么 ViT 会把图像切成 (224/16)^2 196 个 Patch每个 Patch 是 16×16×3 的小块。在实现层面这一步通常用卷积来完成。一个 kernel_size16、stride16 的二维卷积输入通道 3输出通道 embed_dim卷积一次就能把整张图切成 Patch 并完成线性映射。196 个 Patch 每个变成 768 维的向量整个图像就变成了一个形状为 [B, 196, 768] 的 Token 序列。这里就是论文标题所说“An Image is Worth 16x16 Words”一张图等于 196 个单词每个单词长度是 768。3.2 为什么需要位置编码自注意力机制本身是排列不变的。对注意力计算来说把第 1 个 Patch 和第 196 个 Patch 交换位置结果完全一样。但图像显然有空间结构左上角和右下角的语义完全不同。如果不加任何位置信息模型只会得到一堆无序 Patch 的集合。ViT 的解决方案是给每个 Token 加上一个可学习的位置向量。最终会有一个形状为 [1, 197, 768] 的位置编码矩阵其中 197 表示 196 个 Patch Token 加 1 个分类 Token。模型会在训练过程中自己学会这些位置向量应该长什么样。这种可学习位置编码的好处是实现简单缺点是输入分辨率变化时需要做插值。后面工程部分会提到这个问题。3.3 Transformer Encoder注意力、MLP 与残差Encoder Block 是 ViT 的核心计算单元。每个 Block 由两部分组成多头自注意力模块Multi-Head Self-AttentionMSA多层感知机模块MLP每个模块之前都做 LayerNorm模块之后都接残差连接。用伪代码表示x x MSA(LayerNorm(x)) x x MLP(LayerNorm(x))这就是 pre-norm 结构。它在训练稳定性上比后置 LayerNorm 更好是目前最主流的写法。自注意力的核心公式是Attention(Q, K, V) Softmax(QK^T / sqrt(d_k)) V每个 Token 生成 Query、Key、Value 三个向量。Query 和所有 Key 做点积得到注意力分数除以 sqrt(d_k) 是为了防止点积过大导致 Softmax 梯度消失。Softmax 之后得到 0 到 1 的权重再用这个权重加权求和 Value。多头注意力就是把维度切成 num_heads 份分别做注意力再拼接让模型在不同子空间关注不同模式。MLP 通常是两层线性层加 GELU 激活中间隐藏层维度是 embed_dim 的 4 倍即 MLP ratio4。这层是全图共享的负责对每个 Token 做非线性变换。整个 Encoder 堆叠 depth 层一般 Base 模型是 12 层Large 是 24 层。3.4 分类 Token 与分类头ViT 处理分类任务时会在 Patch Token 前面额外拼接一个可学习的分类 Token也就是 CLS Token。它的初始状态是随机向量和所有 Patch Token 一起经过 Encoder 后最后一层 CLS Token 的输出可以看成整张图像的全局表征。为什么用 CLS Token 而不是把所有 Patch 求平均核心原因是 CLS Token 给了模型一个显式的“汇总全图信息”的 Agent。在训练过程中注意力机制会学会让 CLS Token 主动去聚合所有 Patch 的信息。bert 里也是同样的思路。最后的分类头是一个简单的线性层输入维度是 embed_dim输出维度是类别数。3.5 归纳偏置ViT 和 CNN 的本质区别CNN 把“图像是局部连续的一块块像素”这个先验直接写在网络结构里所以小数据也能学得不错。ViT 几乎没有这种先验它只用一个 Patch 切块操作后面的注意力机制完全靠数据学。好处是上限更高坏处是数据量不够时很难训练。这解释了为什么当初 ViT 在 ImageNet-21k 上预训练后表现好但在纯 ImageNet-1k 上直接训练一开始不如同规模 CNN。理解这一点对后续做小数据集实验很有帮助。4. 源码复现从零实现 ViT下面直接进入代码。我们会按 Patch Embedding、Attention、Encoder Block 的顺序组装出一个完整的 VisionTransformer。代码以可读性优先参考了常见开源实现的结构但不依赖任何第三方大库只有 PyTorch。4.1 Patch Embedding 实现import torch import torch.nn as nn class PatchEmbed(nn.Module): 图像切块 线性映射 输入 x: [B, C, H, W] 输出 x: [B, num_patches, embed_dim] def __init__(self, img_size224, patch_size16, in_chans3, embed_dim768): super().__init__() self.img_size img_size self.patch_size patch_size self.num_patches (img_size // patch_size) ** 2 self.proj nn.Conv2d( in_chans, embed_dim, kernel_sizepatch_size, stridepatch_size ) def forward(self, x): # [B, embed_dim, H/patch, W/patch] x self.proj(x) # 展平 H/W 两个维度 x x.flatten(2) # [B, embed_dim, num_patches] # 调整为 Transformer 需要的 [B, num_patches, embed_dim] x x.transpose(1, 2) return x这里没有对 Patch 做像素拼接再线性映射而是直接用卷积实现结果完全等价但计算效率更高。4.2 多头注意力实现class Attention(nn.Module): def __init__(self, dim, num_heads8, qkv_biasTrue, attn_drop0.0): super().__init__() self.num_heads num_heads self.scale (dim // num_heads) ** -0.5 self.qkv nn.Linear(dim, dim * 3, biasqkv_bias) self.attn_drop nn.Dropout(attn_drop) self.proj nn.Linear(dim, dim) def forward(self, x): # x: [B, N, C] B, N, C x.shape # 一次线性层生成 Q、K、V qkv self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads) qkv qkv.permute(2, 0, 3, 1, 4) # [3, B, num_heads, N, head_dim] q, k, v qkv[0], qkv[1], qkv[2] attn (q k.transpose(-2, -1)) * self.scale attn attn.softmax(dim-1) attn self.attn_drop(attn) x (attn v).transpose(1, 2).reshape(B, N, C) x self.proj(x) return x这里的关键操作是一次线性层生成三份向量然后 reshape 成多头结构。注意q k.transpose(-2, -1)得到的注意力矩阵形状是 [B, num_heads, N, N]N 是 Token 数。4.3 MLP 与 Encoder Blockclass Mlp(nn.Module): def __init__(self, in_features, hidden_features, drop0.0): super().__init__() self.fc1 nn.Linear(in_features, hidden_features) self.act nn.GELU() self.fc2 nn.Linear(hidden_features, in_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_ratio4.0, drop0.0, attn_drop0.0): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn Attention(dim, num_headsnum_heads, attn_dropattn_drop) self.norm2 nn.LayerNorm(dim) self.mlp Mlp(dim, int(dim * mlp_ratio), dropdrop) def forward(self, x): x x self.attn(self.norm1(x)) x x self.mlp(self.norm2(x)) return x4.4 完整 VisionTransformer 模型class VisionTransformer(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, num_classes1000, embed_dim768, depth12, num_heads12, mlp_ratio4.0, drop0.1): super().__init__() self.patch_embed PatchEmbed( img_sizeimg_size, patch_sizepatch_size, in_chansin_chans, embed_dimembed_dim ) num_patches self.patch_embed.num_patches # 分类 Token模仿 BERT 的 [CLS] self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) # 位置编码1 个 CLS Token num_patches 个 Patch Token self.pos_embed nn.Parameter(torch.zeros(1, num_patches 1, embed_dim)) self.pos_drop nn.Dropout(pdrop) self.blocks nn.Sequential(*[ Block(dimembed_dim, num_headsnum_heads, mlp_ratiomlp_ratio, dropdrop) 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, 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) 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 Token cls_token self.cls_token.expand(B, -1, -1) x torch.cat([cls_token, x], dim1) # 加位置编码 x x self.pos_embed x self.pos_drop(x) # 过 Encoder x self.blocks(x) x self.norm(x) # 取 CLS Token 的分类表征 cls_out x[:, 0] return self.head(cls_out)到这里一个标准的 ViT 模型就完整了。可以试一下model VisionTransformer( img_size224, patch_size16, in_chans3, num_classes1000, embed_dim768, depth12, num_heads12 ) x torch.randn(1, 3, 224, 224) out model(x) print(out.shape) # [1, 1000]如果能正常输出 shape说明前向流程已经跑通。5. 环境准备与数据说明复现 ViT 不需要特别复杂的依赖。推荐环境如下依赖项建议版本Python3.9 或更高PyTorch2.x1.8 以上基本可用torchvision与 PyTorch 匹配的版本CUDA可选有 NVIDIA 显卡可加速CPU 也能跑通小模型安装环境pip install torch torchvision如果没有 GPU也能用 CPU 跑下面的实验只要把 batch_size 调小、模型调小即可。测试数据集我们使用 CIFAR-10torchvision 自带下载不用手动准备数据。由于 ViT 天然缺少 CNN 的归纳偏置直接在 CIFAR-10 这类小数据集上从零训练效果不会很好这是正常现象。这里的目的是验证模型能训练、能收敛、能推理。真正要用 ViT 做分类建议加载 ImageNet 预训练权重或者先在大数据集上做自监督预训练。6. 在 CIFAR-10 上训练小型 ViT为了在普通电脑上快速跑通我们使用一个小型配置图像 32×32patch4embed_dim192depth6num_heads6。这个模型远小于 ViT-Base但对流程验证完全够用。import torch import torch.optim as optim import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, transforms def build_dataloaders(batch_size64): transform transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ]) train_ds datasets.CIFAR10(root./data, trainTrue, downloadTrue, transformtransform) test_ds datasets.CIFAR10(root./data, trainFalse, downloadTrue, transformtransform) train_loader DataLoader(train_ds, batch_sizebatch_size, shuffleTrue, num_workers2) test_loader DataLoader(test_ds, batch_sizebatch_size, shuffleFalse, num_workers2) return train_loader, test_loader def create_model(): return VisionTransformer( img_size32, patch_size4, in_chans3, num_classes10, embed_dim192, depth6, num_heads6, mlp_ratio4.0, drop0.1 )训练循环使用 AdamW 优化器和余弦退火学习率这在 Transformer 类模型上是很常见的组合。def train(model, train_loader, epochs10, lr1e-3): device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lrlr, weight_decay5e-2) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxepochs) for epoch in range(epochs): model.train() total_loss 0.0 correct 0 total 0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() * images.size(0) _, preds outputs.max(1) correct (preds labels).sum().item() total labels.size(0) scheduler.step() acc correct / total print(fEpoch {epoch 1}/{epochs} | fLoss: {total_loss / total:.4f} | Acc: {acc:.4f})这里加载模型前需要确保 VisionTransformer 类已经在同一命名空间中定义。把前一步的模型类代码放在同一文件即可。运行训练if __name__ __main__: train_loader, test_loader build_dataloaders(batch_size64) model create_model() train(model, train_loader, epochs10)如果电脑是 Windows建议把训练代码包在if __name__ __main__:里避免 DataLoader 多进程问题。训练过程的期望结果前几轮 loss 明显下降准确率从 10% 逐步提高。由于数据量不大10 轮后达到 50% 以上的测试准确率都是可能的但如果差距大也不影响流程验证。重点是模型正常收敛没有出现梯度爆炸或维度错误。7. 推理测试与效果验证训练完成后的验证分成两层测试集整体准确率和单张图片推理。测试集验证def evaluate(model, test_loader): device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in test_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, preds outputs.max(1) correct (preds labels).sum().item() total labels.size(0) print(fTest Accuracy: {correct / total:.4f}) return correct / total单张图片推理from PIL import Image classes [airplane, automobile, bird, cat, deer, dog, frog, horse, ship, truck] transform transforms.Compose([ transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)), ]) image Image.open(test_cat.jpg).convert(RGB) tensor transform(image).unsqueeze(0) device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) tensor tensor.to(device) with torch.no_grad(): logits model(tensor) pred logits.argmax(dim1).item() print(Predicted:, classes[pred])判断成功的标准很简单模型能输出一个类别索引且索引落在合法类别范围内测试集准确率明显高于随机水平CIFAR-10 随机水平是 10%。如果单张图片推理结果不对先看图片内容和标签是否匹配再检查数据预处理是不是和训练时一致。Normalize 的均值和标准差必须完全一致。8. 资源占用与性能观察方法ViT 的资源占用和输入分辨率、Patch Size、Embedding 维度、层数、Batch Size 都有关系具体占多少显存与硬件相关不经过实测不能写死。但有几个通用观察方法。8.1 查看参数量def count_parameters(model): return sum(p.numel() for p in model.parameters() if p.requires_grad) print(参数量:, count_parameters(model))本文的小型 ViT 参数量在百万级别完整 ViT-Base/16 通常超过 8000 万参数。8.2 查看激活显存占用PyTorch 可以直接查询当前显存峰值torch.cuda.reset_peak_memory_stats() with torch.no_grad(): output model(tensor) peak_mem torch.cuda.max_memory_allocated() / 1024**2 print(f显存峰值: {peak_mem:.2f} MB)推理时不开梯度显存占用主要来自模型参数和中间激活值。8.3 分辨率对计算量的影响ViT 的注意层计算复杂度约为 O(N^2)其中 N 是 Token 数而 Token 数等于 (H/patch) × (W/patch)。分辨率从 224×224 提高到 448×448Token 数变成原来的 4 倍注意力计算量大约变成 16 倍。这是 ViT 在高分辨率下吃显存的核心原因。降低占用可以从这几个方向入手方法效果减小 Batch Size直接降低激活显存增大 Patch SizeToken 数减少例如 patch 16 改 32Token 数变成 1/4使用混合精度FP16 推理和训练可以显著降低显存梯度检查点训练时用计算换显存多维注意力优化使用 Swin 等局部注意力变体8.4 计算 FLOPs如果需要定量比较模型计算量可以使用 thop 或 fvcorepip install thopfrom thop import profile flops, params profile(model, inputs(torch.randn(1, 3, 224, 224),)) print(fFLOPs: {flops / 1e9:.2f} G, Params: {params / 1e6:.2f} M)需要注意这里的 FLOPs 只是前向传播的粗略估算实际数值会因实现方式不同有差异。9. 常见问题与排查方法问题现象可能原因排查方式解决方案图像尺寸和 Patch 不匹配图像尺寸不能被 patch_size 整除检查输入 shape 和 patch_size调整图像尺寸或使用 padding训练 loss 不下降学习率不合适、数据没归一化、数据集太小打印 loss检查输入数据范围调低学习率或增加 warmup加载预训练权重显存不足 OOMBatch 太大、分辨率太高、模型太大查看报错位置和显存占用降低 batch、提高 patch_size、开混合精度测试时分类结果全集中在一个类别模型欠拟合或类别不均衡看训练准确率是否偏低加大训练轮数检查数据增强位置编码维度对不上换输入分辨率后 Token 数变了打印 pos_embed shape对 pos_embed 做插值或固定分辨率CPU 训练极慢模型和 batch 对 CPU 不友好观察每轮耗时调小模型、调小 batch、使用 GPU下载 CIFAR-10 失败网络问题重试或手动下载后放入 ./data使用镜像源或离线数据包如果训练过程中 loss 突然变成 NaN优先检查学习率是否过大以及数据中是否存在异常值。ViT 对初始化也敏感本文代码里使用了常见的 trunc_normal 初始化正常情况不会出现数值不稳定。10. 最佳实践与使用建议第一第一次实验不要直接复现 ViT-Base。先在 32×32 或 64×64 的小图、小模型上跑通全流程确认训练、保存、加载、推理、指标统计都正常再逐步放大。这样可以省掉大量排错时间。第二训练和推理的数据预处理必须保持一致。包括 Resize 尺寸、Normalize 的均值和标准差、颜色通道顺序。很多效果问题其实不是模型问题而是数据管线不一致。第三模型输出管理和日志要做在前面。建议把训练日志、模型权重、推理结果分目录保存方便追踪实验。每次实验前固定随机种子保证可复现。第四涉及开源预训练权重时要检查权重来源的许可证和使用条款。商业项目里使用 ViT 权重不能只看能下载就随便用。同理如果准备把人脸、车牌等敏感图像喂进模型一定要确认数据来源合法、使用范围合规。第五ViT 的批量推理很简单只要用 DataLoader 设置大 batch 即可。但如果要对外提供 API 服务建议用 FastAPI 封装并且限制请求大小和并发量。# 通用 API 服务封装模板需按实际项目调整 from fastapi import FastAPI from pydantic import BaseModel import io import torch from PIL import Image class PredictRequest(BaseModel): image_base64: str app FastAPI() app.post(/predict) def predict(req: PredictRequest): # 解析 base64 - PIL Image - tensor - model - label # 此处为模板实际逻辑根据模型输入输出调整 return {status: ok}第六模型导出。训练好的 ViT 可以用torch.onnx.export或 TorchScript 导出方便部署到服务端。导出时固定输入尺寸最省事如果必须支持动态分辨率需要额外处理位置编码的插值问题。11. 总结与下一步ViT 的核心并不神秘就是先把图像切成 Patch再让 Transformer Encoder 做全局注意力建模最后用 CLS Token 输出分类结果。源码复现的关键也就三处Patch Embedding 的实现、多头注意力、Encoder Block 的残差结构把这三部分串起来完整的 ViT 就写好了。这篇文章值得先收藏。第一次运行可以先拿最小的 CIFAR-10 配置测试确认环境没问题再逐步替换成自己的数据集和更大的配置。后续扩展可以直接往这几个方向走试试用 ViT 提取特征再接 FPN 做目标检测换成 Swin Transformer 理解局部注意力如何降低计算量或者加载预训练权重在自定义数据集上做微调。理解了这一套再去读 CLIP、SAM、DINOv2 的代码会发现底层都是熟悉的 Encoder 结构。