
这次我们直接手撕 Transformer。不是调transformers库也不是套别人封装好的vit_pytorch而是从 Patch Embedding 到多头自注意力再到完整的 Forward 前向传播用 PyTorch 把 Vision TransformerViT的核心链路逐行写出来。如果你一直觉得 Transformer 的代码难啃其实核心就三件事第一怎么把图片变成模型能处理的序列第二怎么在序列上做自注意力计算第三怎么把编码后的特征变成最终的分类输出。ViT 的代码恰好把这三件事完整串起来而且结构比 GPT、BERT 这类超大模型简单得多是最适合用来“手撕”的版本。这篇文章就围绕这三件事展开读完你就能自己改 ViT 源码做图像分类、特征提取和批量推理。本文适用读者很明确已经会用 PyTorch 搭建 CNN 分类模型但对 Transformer 内部实现处于“看得懂论文、写不出代码”阶段的同学以及需要在图像分类、图像检索、特征提取任务里自己调整 ViT 结构的工程同学。文中的所有代码都不依赖第三方 Transformer 库只依赖 PyTorch 和 torchvisionCPU 环境就能完成前向验证训练环节也可以用少量数据跑通。下面直接开始。1. 核心知识速览先把这篇文章会手写的代码能力列成一张表看完你就知道整个实现链路覆盖到哪里。能力项说明输入表示图像张量 (B, C, H, W) 转换为 Patch 序列 (B, N, D)Patch Embedding用 Conv2d 实现图像切块 线性映射位置编码可学习的 Positional Embedding给序列补充位置信息自注意力QKV 映射 缩放点积注意力 多头拼接 输出映射Encoder BlockLayerNorm Attention MLP 残差连接分类头cls_token 提取特征 Linear 输出 logits完整 Forward端到端前向传播与每一层形状推导训练验证MNIST 图像分类最小实验批量推理文件夹批量预测 Flask API 封装示例整个 ViT 的前向链路可以概括为图片 - Patch 序列 - 加位置编码 - 过 N 个 Transformer Encoder Block - 取 cls_token - 分类头。下面按这个链路逐个模块实现。2. 为什么是 Transformer从序列建模到视觉任务Transformer 最早是为 NLP 设计的核心机制是自注意力。自注意力的核心计算是对于序列中的每一个 token都去计算它和其他所有 token 的相关性再根据相关性加权聚合信息。这样做的好处是任意两个 token 之间都有一条直接的信息通路距离不再是问题。CNN 的感受野靠堆叠卷积层逐步扩大RNN 靠逐步传递隐藏状态而 Transformer 一步就能让全局信息互相可见。Vision TransformerViT的核心思想是把这个机制搬到图像上。图像本身不是序列所以第一步必须做转换把一张 H×W 的图片切成若干个小 Patch每个 Patch 拉平后做线性映射变成一个向量。这样一张图片就变成了一个向量序列可以直接送进 Transformer Encoder。论文里的典型设置是 224×224 的图片切成 16×16 的 Patch得到 14×14196 个 Patch每个 Patch 展平后映射成 768 维向量。相比 CNNViT 的全局建模能力更强尤其是在大数据集上ViT 可以超越同参数量级的 CNN。但它的缺点也很明显自注意力的计算复杂度是序列长度的平方图片分辨率越高Patch 数量越多计算量增长越快。这也是后面 Swin Transformer、DeiT 等改进方向要解决的问题。理解了 ViT 的 Forward再去看这些变体就轻松多了。3. 环境准备与前置条件这篇文章的代码只需要一个深度学习基础环境门槛很低。# Python 3.8 以上推荐 3.10 python -m pip install torch torchvision # 验证安装 python -c import torch; print(torch.__version__)操作系统Windows / Linux / macOS 都可以代码不涉及平台相关操作。硬件CPU 即可完成前向验证和小规模训练。有 NVIDIA GPU 会快很多但不是必须。内存建议 8GB 以上训练 MNIST 时数据集会加载到内存。磁盘PyTorch 安装约 2GB 左右MNIST 数据集几十 MB。GPU 要求没有强制要求。如果你有 CUDA 环境训练代码会自动调用cuda没有就回退到cpu。import torch device torch.device(cuda if torch.cuda.is_available() else cpu) print(Using device:, device)如果这一步打印的是cpu完全不影响后面的代码运行。下面所有代码我都按 CPU/GPU 均可运行的方式写。4. 图像预处理与 Patch 切块原理动手写代码之前必须先搞清楚一个核心问题图片是怎么变成序列的。一张 RGB 图片在 PyTorch 里的张量形状是 (B, C, H, W)B 是 batch sizeC 是通道数H 和 W 是宽高。ViT 的做法是先把图片切成一个个不重叠的小方块每个小方块称为一个 Patch。以 224×224 的输入、16×16 的 Patch 为例每张图片被切成 (224/16) × (224/16) 14 × 14 196 个 Patch。每个 Patch 的大小是 3×16×16展平后是 768 维向量。最终得到序列形状 (B, 196, 768)序列长度是 196每个 token 的维度是 768。在 PyTorch 里最优雅的实现方式不是用循环切图而是用一个 stride 等于 kernel_size 的 Conv2d 一步完成。卷积核大小设为 patch_size步长也设为 patch_size每个卷积核的输出通道对应一个 embedding 维度。这样输出的空间尺寸自然就是 H/patch × W/patch然后把空间维度展平就得到了 Patch 序列。形状变化如下表输入形状操作输出形状(B, 3, 224, 224)Conv2d 16×16, stride16(B, 768, 14, 14)(B, 768, 14, 14)flatten(2)(B, 768, 196)(B, 768, 196)transpose(1, 2)(B, 196, 768)这就是 Patch Embedding 的全部秘密。下面直接写代码。5. 核心模块代码实现5.1 Patch Embedding 代码实现import torch import torch.nn as nn class PatchEmbedding(nn.Module): def __init__(self, img_size224, patch_size16, in_channels3, embed_dim768): super().__init__() self.img_size img_size self.patch_size patch_size self.num_patches (img_size // patch_size) ** 2 self.embed_dim embed_dim # 用卷积实现切块 线性映射 self.proj nn.Conv2d( in_channels, embed_dim, kernel_sizepatch_size, stridepatch_size, ) def forward(self, x): # x: (B, C, H, W) B, C, H, W x.shape assert H self.img_size and W self.img_size, \ f输入尺寸应为 {self.img_size}x{self.img_size}当前为 {H}x{W} x self.proj(x) # (B, embed_dim, H/patch, W/patch) x x.flatten(2) # (B, embed_dim, num_patches) x x.transpose(1, 2) # (B, num_patches, embed_dim) return xflatten(2)会把 (B, D, H/patch, W/patch) 的后两维展平得到 (B, D, num_patches)。transpose(1, 2)把维度顺序调整为 (B, num_patches, D)这样每个 Patch 就是一个 token可以直接输入 Transformer。5.2 Positional Encoding 代码实现Patch 序列本身没有顺序信息。Transformer 的自注意力对序列顺序不敏感把第 1 个 Patch 和第 196 个 Patch 交换位置注意力计算结果是一样的所以必须显式加入位置编码。ViT 使用的是可学习的位置编码每个位置对应一个可训练向量初始化为小的随机值。class PositionalEncoding(nn.Module): def __init__(self, num_patches, embed_dim, dropout0.0): super().__init__() # 多一个位置给 cls_token self.pos_embed nn.Parameter(torch.zeros(1, num_patches 1, embed_dim)) nn.init.trunc_normal_(self.pos_embed, std0.02) self.dropout nn.Dropout(dropout) def forward(self, x): # x: (B, N, D) return self.dropout(x self.pos_embed)注意num_patches 1ViT 在序列开头额外加了一个 cls_token所以位置编码也要多一个。后面实现完整模型时会看到 cls_token 的用法。5.3 Multi-Head Self-Attention 代码实现自注意力是整个 Transformer 的灵魂。多头注意力做的事情是把输入序列通过三个线性映射分别变成 Q、K、V然后分成多个头分别计算注意力最后拼接并做输出映射。class MultiHeadSelfAttention(nn.Module): def __init__(self, embed_dim, num_heads, dropout0.0): super().__init__() assert embed_dim % num_heads 0, embed_dim 必须能被 num_heads 整除 self.num_heads num_heads self.head_dim embed_dim // num_heads self.qkv nn.Linear(embed_dim, embed_dim * 3) self.proj nn.Linear(embed_dim, embed_dim) self.dropout nn.Dropout(dropout) def forward(self, x): # x: (B, N, D) B, N, D x.shape qkv self.qkv(x) # (B, N, 3D) qkv qkv.reshape(B, N, 3, self.num_heads, self.head_dim) 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.head_dim ** -0.5) attn attn.softmax(dim-1) attn self.dropout(attn) x attn v # (B, num_heads, N, head_dim) x x.transpose(1, 2).contiguous().reshape(B, N, D) x self.proj(x) # (B, N, D) return x几个细节值得注意缩放因子head_dim ** -0.5是为了防止点积结果过大导致 softmax 梯度消失。Q 和 K 的维度是 head_dim点积结果的方差会随 head_dim 增大而增大除以根号 head_dim 可以稳定训练。permute之后的 q、k、v 形状是 (B, num_heads, N, head_dim)注意力矩阵形状是 (B, num_heads, N, N)其中 N 是序列长度。transpose之后必须加.contiguous()再reshape否则会报运行时错误。5.4 Transformer Encoder Block 代码实现一个标准的 Transformer Encoder Block 由两个子层组成多头注意力子层和 MLP 子层每个子层前面有 LayerNorm后面有残差连接。ViT 默认使用 Pre-LayerNorm 结构即先归一化再送入子层。class TransformerEncoderBlock(nn.Module): def __init__(self, embed_dim, num_heads, mlp_ratio4.0, dropout0.0): super().__init__() self.norm1 nn.LayerNorm(embed_dim) self.attn MultiHeadSelfAttention(embed_dim, num_heads, dropout) self.norm2 nn.LayerNorm(embed_dim) hidden_dim int(embed_dim * mlp_ratio) self.mlp nn.Sequential( nn.Linear(embed_dim, hidden_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(hidden_dim, embed_dim), nn.Dropout(dropout), ) def forward(self, x): x x self.attn(self.norm1(x)) x x self.mlp(self.norm2(x)) return xMLP 的隐藏层维度通常是 embed_dim 的 4 倍激活函数使用 GELU。残差连接的作用是让梯度可以跨层传播避免深层网络退化。整个 Block 的输入输出形状都是 (B, N, D)所以可以叠加任意多层。5.5 Vision Transformer 完整代码实现把前面几个模块拼起来加上 cls_token 和分类头就是完整的 ViT 前向链路。class VisionTransformer(nn.Module): def __init__(self, img_size224, patch_size16, in_channels3, num_classes1000, embed_dim768, depth12, num_heads12, mlp_ratio4.0, dropout0.0): super().__init__() self.patch_embed PatchEmbedding(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(dropout) self.blocks nn.ModuleList([ TransformerEncoderBlock(embed_dim, num_heads, mlp_ratio, dropout) 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, module): if isinstance(module, nn.Linear): nn.init.trunc_normal_(module.weight, std0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.LayerNorm): nn.init.zeros_(module.bias) nn.init.ones_(module.weight) def forward(self, x): # x: (B, C, H, W) B x.size(0) x self.patch_embed(x) # (B, N, D) cls_token self.cls_token.expand(B, -1, -1) # (B, 1, D) x torch.cat([cls_token, x], dim1) # (B, N1, D) x x self.pos_embed x self.pos_drop(x) for block in self.blocks: x block(x) # (B, N1, D) x self.norm(x) cls_out x[:, 0] # 取 cls_token 位置的特征 logits self.head(cls_out) # (B, num_classes) return logitscls_token 是 ViT 里一个非常巧妙的设计。它一开始是随机初始化的向量在训练过程中通过注意力机制不断聚合整张图片的全局信息。最终分类时只取 cls_token 对应位置的输出不需要对全部 Patch 的输出做平均池化这样就能让模型自己决定哪些 Patch 的信息对分类更重要。6. 完整模型汇总与 Forward 验证上面 5 个模块合在一起就是一份不依赖任何第三方 Transformer 库的 ViT 实现。下面用一个随机张量验证 Forward 是否跑通。if __name__ __main__: # 构造一个小型 ViT 配置 model VisionTransformer( img_size224, patch_size16, in_channels3, num_classes10, embed_dim192, depth4, num_heads6, mlp_ratio4.0, dropout0.1, ) x torch.randn(2, 3, 224, 224) logits model(x) print(输入形状:, x.shape) print(输出形状:, logits.shape) # 预期输出: 输出形状 torch.Size([2, 10])运行这段代码如果能看到输出形状是torch.Size([2, 10])说明整个前向链路已经打通。在动手训练之前建议先逐层打印中间张量形状确认每个模块的输出符合预期def inspect_forward(model, x): print(输入:, x.shape) x model.patch_embed(x) print(Patch Embedding 后:, x.shape) B x.size(0) cls_token model.cls_token.expand(B, -1, -1) x torch.cat([cls_token, x], dim1) print(拼接 cls_token 后:, x.shape) x x model.pos_embed print(加位置编码后:, x.shape) for i, block in enumerate(model.blocks): x block(x) print(f第 {i1} 个 Encoder Block 后:, x.shape) x model.norm(x) cls_out x[:, 0] print(cls_token 特征:, cls_out.shape) logits model.head(cls_out) print(分类输出:, logits.shape)这种逐层打印的方式排查维度错误效率最高。如果某一步形状和预期不一致问题基本就出在那一步的 reshape、transpose 或 padding 逻辑上。7. 图像分类训练实战Forward 跑通之后用 MNIST 做一个小实验验证模型真的能学到东西。MNIST 是灰度图所以in_channels设为 1num_classes设为 10。import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms transform_train transforms.Compose([ transforms.Resize((64, 64)), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)), ]) transform_test transforms.Compose([ transforms.Resize((64, 64)), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)), ]) train_dataset datasets.MNIST(root./data, trainTrue, downloadTrue, transformtransform_train) test_dataset datasets.MNIST(root./data, trainFalse, downloadTrue, transformtransform_test) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue, num_workers2) test_loader DataLoader(test_dataset, batch_size64, shuffleFalse, num_workers2) device torch.device(cuda if torch.cuda.is_available() else cpu) model VisionTransformer( img_size64, patch_size8, in_channels1, num_classes10, embed_dim128, depth4, num_heads4, mlp_ratio4.0, dropout0.1, ).to(device) criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr1e-4) def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss 0.0 correct 0 total 0 for images, labels in loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() total_loss loss.item() correct (outputs.argmax(dim1) labels).sum().item() total labels.size(0) return total_loss / len(loader), correct / total def evaluate(model, loader, device): model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in loader: images, labels images.to(device), labels.to(device) outputs model(images) correct (outputs.argmax(dim1) labels).sum().item() total labels.size(0) return correct / total for epoch in range(3): train_loss, train_acc train_one_epoch(model, train_loader, criterion, optimizer, device) test_acc evaluate(model, test_loader, device) print(fEpoch {epoch 1}: loss{train_loss:.4f}, train_acc{train_acc:.4f}, test_acc{test_acc:.4f})这里有两个建议。第一如果 CPU 训练建议先用torch.utils.data.Subset取 5000 张训练数据跑 1 个 epoch确认链路无误后再全量训练避免一个下午等来一个维度报错。第二Resize((64, 64))配合patch_size8得到 8×864 个 Patch计算量比 224×224 小很多适合在没有 GPU 的环境里验证。从经验上看这种小型 ViT 在 MNIST 上 3 个 epoch 就能看到明显下降的 loss 和上升的准确率不需要达到 SOTA目标只是证明 Forward、反向传播和梯度更新这条链路是通的。如果 loss 完全不下降优先检查学习率、数据归一化和模型初始化。8. 批量推理与 Flask API 封装模型训练完后实际工程里通常有两种使用方式一种是批量处理本地文件夹里的图片另一种是把模型封装成 HTTP API 服务供其他系统调用。两种都写出来。8.1 文件夹批量推理import torch from pathlib import Path from PIL import Image from torchvision import transforms def build_transform(img_size64): return transforms.Compose([ transforms.Resize((img_size, img_size)), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)), ]) def predict_folder(model, image_dir, img_size64, batch_size32, devicecpu): transform build_transform(img_size) image_files sorted(Path(image_dir).glob(*.jpg)) sorted(Path(image_dir).glob(*.png)) results [] for i in range(0, len(image_files), batch_size): batch_files image_files[i:i batch_size] batch_tensors [] for f in batch_files: img Image.open(f).convert(L) # MNIST 训练时用灰度图 batch_tensors.append(transform(img)) batch torch.stack(batch_tensors).to(device) with torch.no_grad(): logits model(batch) preds logits.argmax(dim-1).tolist() for f, pred in zip(batch_files, preds): results.append((str(f), pred)) print(f{f.name}: {pred}) return results批量推理时需要注意几个细节图片读取后一定要convert成和训练时一致的通道数归一化参数必须和训练保持一致model.eval()和torch.no_grad()缺一不可否则 dropout 和 BN 统计量会导致推理结果不稳定。8.2 Flask API 封装示例python -m pip install flaskimport io from flask import Flask, request, jsonify from PIL import Image from torchvision import transforms app Flask(__name__) device torch.device(cuda if torch.cuda.is_available() else cpu) model VisionTransformer( img_size64, patch_size8, in_channels1, num_classes10, embed_dim128, depth4, num_heads4, mlp_ratio4.0, dropout0.1, ).to(device) model.load_state_dict(torch.load(vit_mnist.pth, map_locationdevice)) model.eval() transform transforms.Compose([ transforms.Resize((64, 64)), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)), ]) app.route(/predict, methods[POST]) def predict(): if image not in request.files: return jsonify({error: 缺少 image 文件}), 400 file request.files[image] image Image.open(io.BytesIO(file.read())).convert(L) tensor transform(image).unsqueeze(0).to(device) with torch.no_grad(): logits model(tensor) pred logits.argmax(dim-1).item() prob torch.softmax(logits, dim-1).max().item() return jsonify({prediction: pred, confidence: round(prob, 4)}) if __name__ __main__: app.run(host127.0.0.1, port5000)启动服务后用 curl 测试curl -X POST -F imagetest.png http://127.0.0.1:5000/predict返回结果示例{ prediction: 7, confidence: 0.9821 }接口服务上线前至少要验证三件事输入图片大小不符合预期时是否返回明确错误、并发请求下显存或内存是否稳定、模型推理耗时是否满足业务要求。9. 资源占用与性能观察ViT 的计算开销主要来自两个部分Patch Embedding 的卷积计算和自注意力矩阵计算。自注意力的复杂度是 O(N²·D)N 是 Patch 数量D 是 embedding 维度。图片分辨率翻倍Patch 数量变成 4 倍自注意力计算量变成 16 倍。这是 ViT 在高分辨率场景下最需要关注的问题。几个常见模型配置的规模对比标准公开配置模型embed_dimdepthnum_heads参数量级ViT-Tiny192123约 5MViT-Small384126约 22MViT-Base7681212约 86M实际训练时的显存占用取决于四个主要因素batch size、patch 数量、embedding 维度和模型深度。显存不够时优先降低 batch size如果 batch size 已经降到 1 还不够就降低输入分辨率或使用 gradient checkpointing。# 统计参数量 def count_parameters(model): return sum(p.numel() for p in model.parameters() if p.requires_grad) model VisionTransformer(img_size224, patch_size16, num_classes1000) print(f可训练参数量: {count_parameters(model) / 1e6:.2f}M)在 CPU 上做前向验证时建议用 64×64 输入和 depth4 的小配置单次前向快在 GPU 上训练时才使用 224×224 和 depth12 的较大配置。反过来如果你在 CPU 上跑 224×224 的 ViT-Base一次前向可能需要几十秒这属于正常现象不是代码写错了。10. 常见问题与排查方法问题现象可能原因排查方式解决方案输入尺寸报错图片 H/W 与 img_size 不一致打印输入张量形状预处理中 Resize 到模型指定尺寸维度不匹配transpose 后未 contiguous查看报错堆栈reshape 前加 .contiguous()embed_dim 无法整除 num_heads配置错误检查模型参数修改 embed_dim 或 num_heads显存不足batch size 太大观察显存占用降低 batch size 或分辨率CUDA 不可用驱动/版本问题torch.cuda.is_available()用 CPU 跑或升级驱动训练 loss 不下降学习率过高/过低打印 loss 值调整学习率到 1e-4 附近输出全是同一个类别分类头未收敛检查验证集准确率降低学习率、增加训练轮次API 请求 400没有上传 image 字段检查请求格式用 curl -F 上传文件推理结果和训练差异大未调用 model.eval()检查代码推理时加 model.eval() 和 no_grad批量任务中途卡住数据读取异常查看日志对单张图片单独测试定位最常用的调试手段有两个一是逐层打印张量形状对照第 6 节的inspect_forward二是先用 batch size 为 1 的随机张量跑通 Forward再换真实数据。这两个方法能解决 90% 以上的维度问题。11. 最佳实践与使用建议先跑小模型再上大模型。第一次实现时用img_size64, patch_size8, embed_dim128, depth4这组配置Forward 通了再逐步放大。固定随机种子。ViT 训练对初始化敏感复现实验前先固定torch.manual_seed和np.random.seed。数据归一化必须一致。训练和推理的 Normalize 参数不一致是部署时最常见的隐性错误。优先使用预训练权重。对于真实业务场景不建议从零训练 ViT直接加载在 ImageNet 上预训练的权重做微调效果和收敛速度都会好很多。模型文件与代码分目录管理。权重文件、数据集、输出结果不要混在一个目录训练和推理脚本也要分开。输入素材和输出结果要留日志。批量推理时记录每张图片的文件名、预测结果和置信度方便事后回溯。接口服务要限制访问范围。部署 API 时绑定127.0.0.1或内网 IP不要直接暴露公网如果需要对外提供服务必须加鉴权和限流。承担数据合规责任。使用公开数据集要遵守数据集 license处理真实业务图片时要确认有合法授权涉及人脸、车牌、医疗影像等敏感信息时更要谨慎。12. 总结与下一步这篇文章从 Patch Embedding 开始到多头自注意力、Encoder Block、完整 Forward再到 MNIST 训练验证、批量推理和 Flask API 封装把 Vision Transformer 核心代码完整手写了一遍。整条链路不依赖任何第三方 Transformer 库代码可以直接复制到自己的项目里做二次开发。最值得保留的部分是第 5 节的模块实现和第 6 节的逐层形状验证。前者让你具备修改 ViT 源码的能力后者帮你快速定位维度问题。最容易踩的坑集中在 transpose 后的 reshape、通道数不一致、归一化参数不统一这三处建议收藏备用。下一步可以往两个方向扩展一是换用 Swin Transformer用窗口注意力降低高分辨率下的计算量二是研究 DeiT通过蒸馏方式让小模型在中小数据集上也能训出不错的效果。理解了本文的 Forward 链路之后再去看这些变体的源码你基本只需要关注它们改了哪个模块其他部分都是你熟悉的结构。