简介面向Transformer与计算机视觉初学者这份实践项目以经典的“猫狗大战”分类任务为载体完整演示如何用ViT模型搭建图像分类流程。项目不局限于猫狗场景只需替换数据集并调整类别数参数即可迁移至其他图像分类任务具备良好泛化能力适合希望快速上手ViT的深度学习入门者。压缩包内共2000个文件包含1998张jpg格式的猫狗图像数据用于训练与验证另有2个python脚本分别承担模型定义、训练与推理等核心逻辑代码结构清晰、注释简洁便于逐行理解。整个资源约218.41MB已吸引1469人学习使用。通过该实践项目学习者可以直观认识Transformer注意力机制在计算机视觉中的工作方式掌握ViT的原理与结构并学会在真实项目中组织数据、编写训练流程和调整参数是兼顾理论与实战的入门级参考资料。1. 一段代码让猫狗分类告别卷积ViT 不是玩具Vision TransformerViT这几年从 NLP 跨界到计算机视觉最常被拿来练手的就是图像分类任务而“猫狗大战”这种经典二分类数据集正好是验证注意力机制能不能替代卷积的试金石。这个实践项目把一张 224×224 的猫狗图片切成 16×16 的小 patch每个 patch 映射成一个向量再扔进 Transformer Encoder 里做全局建模最后用 CLS token 输出“猫”或“狗”的概率。我最初也怀疑纯注意力模型在小数据集上会不会直接过拟合但实测下来配合数据增强和 warmup在没有预训练权重的情况下ViT-Tiny 也能在猫狗数据集上训到 95% 左右的准确率。如果你想知道 Transformer 在视觉任务上到底怎么落地、动手改模块时该碰哪些代码这个项目是很好的起点。2. 技术选型与工程结构为什么是 ViT目录怎么摆2.1 技术选型注意力机制取代卷积归纳偏置CNN 做分类靠的是局部感受野和权值共享先看纹理再看形状这是一种很强的归纳偏置——假设相邻像素相关性最高。ViT 的思路反过来不假设局部性把图像当成一串有序的 token让全局注意力机制自己决定该关注哪些区域。猫和狗的差异不止在耳朵和胡子这种局部特征整体轮廓和身体姿态也很关键所以这个任务很适合拿来对比两类模型的差异。模型规模上我选的是 ViT-Tiny 和 ViT-Small 两个配置而不是 ViT-Base。原因很实际猫狗数据集一般就一两万张图ViT-Base 有 8600 万参数没有大规模预训练权重撑着很容易过拟合到训练集上训练速度也慢。Tiny 版参数只有 500 万左右单张 2080 Ti 级别的显卡就能跑迭代一轮很快方便反复改超参。Transformer 的 self-attention 复杂度是 O(n²)这是选 patch size 时要考虑的。224×224 的图切成 16×16 patch得到 14×14196 个 tokenn196 带来的计算量完全可接受。如果切成 8×8 patchtoken 数量变成 784注意力矩阵的尺寸会膨胀 16 倍小数据集上没必要这么干。2.2 工程结构一个训练脚本能从零跑到出结果整个项目我按“配置、数据、模型、训练、推理”五个方向拆分避免把所有逻辑塞进一个 train.py 里。目录结构长这样vit-cat-dog/ ├── config.py # 超参数集中管理改参数不用翻代码 ├── train.py # 训练入口数据加载、模型构建、训练循环 ├── infer.py # 推理脚本加载权重对单张图片预测 ├── src/ │ ├── dataset.py # 数据增强、Dataset 构建、归一化统计 │ ├── model.py # ViT 模型定义patch embed、attention、encoder │ └── utils.py # 日志、checkpoint 保存、准确率统计 ├── data/ │ ├── train/ │ │ ├── cat/ # 猫的训练图片 │ │ └── dog/ # 狗的训练图片 │ └── val/ │ ├── cat/ # 猫的验证图片 │ └── dog/ # 狗的验证图片 └── weights/ # 训练产出的 checkpoint 目录data 目录直接用 ImageFolder 的规范组织torchvision 的datasets.ImageFolder能自动按子目录名生成类别标签。config.py里集中放 patch size、depth、heads、dropout、学习率这些参数我每次调参只改这个文件不碰主逻辑。核心训练流程是读取配置 → 构建 train/val 两个 Dataset → 初始化 ViT 模型 → 用 AdamW 优化器跑训练循环 → 每个 epoch 在验证集上算准确率 → 保存最好的 checkpoint。下文每一层都会拆开讲。模块职责关键参数data/原始图片按类别存放训练集约 1.8 万张验证集约 2000 张src/model.pyViT 模型定义patch_size16, dim192, depth12, heads3src/dataset.py增强和归一化RandomResizedCrop(224), normalize(0.5, 0.5)train.py训练编排epochs50, batch64, lr3e-4infer.py单图推理加载 best.pth, softmax 输出概率3. 数据准备猫狗图片如何清洗成 ViT 要的 token 序列3.1 归一化统计别随手抄 ImageNet 的 mean 和 std很多人在自己的数据集上直接套 ImageNet 的归一化参数这能用但不一定最优。猫狗图片的色彩分布和 ImageNet 差挺多图库来源不同、白平衡不同mean/std 都有偏移。这个项目里我建议先跑一段脚本统计训练集的全局 mean 和 std然后在 dataset.py 里固定写死验证集也沿用同一份统计不做单独计算。# src/dataset.py import torch from torch.utils.data import Dataset from torchvision import datasets, transforms from PIL import Image def get_transform(is_train: bool, mean, std): if is_train: return transforms.Compose([ transforms.RandomResizedCrop(size224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize(meanmean, stdstd), ]) else: return transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(meanmean, stdstd), ]) def build_dataset(data_root: str, mean, std): train_ds datasets.ImageFolder( rootf{data_root}/train, transformget_transform(is_trainTrue, meanmean, stdstd) ) val_ds datasets.ImageFolder( rootf{data_root}/val, transformget_transform(is_trainFalse, meanmean, stdstd) ) return train_ds, val_ds参数说明scale(0.7, 1.0)表示裁剪区域占原图面积的 70% 到 100%让模型看到不同尺度的猫狗RandomHorizontalFlip是猫狗分类最安全的增强因为左右翻转不改变语义——猫不管面朝哪边都是猫。ColorJitter我控制得比较轻三个通道的抖动都压在 0.2防止颜色变化太大导致模型把颜色当决定性特征。ViT 对位置信息很敏感如果像训 CIFAR 那样用 RandomRotation(30) 这种强旋转猫狗的头部朝向会被打乱而 ViT 的位置编码会在训练中记住“常见位置关系”旋转后的样本反而成了噪声。所以这里我没加旋转。3.2 Debug 工具训练前先打印一版样本数据部分最大的坑是图能读进来但增强后图片已经不像猫狗了。我的习惯是在跑训练之前单独写一段可视化脚本把 train_loader 的 batch 拿出来反归一化存成 4×4 的网格图看一遍。def inspect_batch(loader, mean, std): images, labels next(iter(loader)) images images.permute(0, 2, 3, 1).numpy() images images * std mean # 反归一化还原到 [0,1] images (images * 255).astype(uint8) for i in range(min(4, len(images))): img Image.fromarray(images[i]) img.save(fdebug_sample_{i}_label_{labels[i].item()}.jpg)跑完这个脚本重点看三件事裁剪区域是否保留了猫狗的主体颜色有没有失真label 和图片内容是否对得上。很多情况是数据目录分类放反了或者增强参数太激进导致图片被裁成一片空白这些问题在训练前发现只要几分钟训练到一半才发现就是几个小时白跑。4. 核心代码拆解从 Patch Embedding 到分类头的完整实现4.1 Patch Embedding用卷积实现 patch 切分ViT 的第一步是把图片切成 patch 并映射成向量。最直接的实现是用F.unfold把图片展开成 patch再经过一个线性层映射。但工程上更高效的做法是用一个Conv2d搞定卷积核大小和步长都等于 patch_size输出通道就是 embedding 维度。import torch import torch.nn as nn class PatchEmbed(nn.Module): def __init__(self, in_chans3, img_size224, patch_size16, embed_dim192): 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): # 输入 x: [B, 3, 224, 224] x self.proj(x) # [B, 192, 14, 14] x x.flatten(2) # [B, 192, 196] x x.transpose(1, 2) # [B, 196, 192] return x逻辑说明proj是一个步长等于卷积核大小的卷积224×224 输入经过 16×16、stride16 的卷积后特征图变成 14×14。flatten(2)把 14×14 的空间维度压平得到 196 个 patch 的序列最后 transpose 成[batch, num_patches, embed_dim]这是 Transformer Encoder 标准的输入形状。这里有个值得注意的设计用卷积实现 patch embedding 比 unfold Linear 少一次内存拷贝在 batch 较大时能明显提升数据加载效率。embed_dim 是每个 patch 映射到的向量长度这个值决定了后续注意力计算的宽度太小会损失表达力太大在小数据集上容易过拟合。4.2 多头自注意力qkv 合一矩阵省一次显存申请多头注意力是 ViT 的核心。实现上我习惯把 query、key、value 三个投影合并成一个线性层输出维度是 3 倍 embed_dim然后切成三段避免写三次nn.Linear。class Attention(nn.Module): def __init__(self, dim, num_heads3, qkv_biasTrue, attn_drop0.0): 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, biasqkv_bias) self.attn_drop nn.Dropout(attn_drop) 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) qkv qkv.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) attn self.attn_drop(attn) x (attn v).transpose(1, 2).reshape(B, N, C) x self.proj(x) return x参数和逻辑说明head_dim是每个头负责的向量维度dim192、num_heads3 时每个头是 64 维。scale head_dim ** -0.5是 scaled dot-product attention 里的缩放因子防止向量内积过大把 softmax 推入饱和区。reshape加permute把 qkv 拆成三份并调整维度顺序为[3, B, num_heads, N, head_dim]。注意力分数矩阵的形状是[B, num_heads, 196, 196]196 是 patch 数量这个矩阵可以可视化出来看模型关注了哪些区域后面二次开发章节会用到。proj是注意力输出后的线性投影相当于把多头结果融合回 dim 维度。4.3 Encoder BlockPre-Norm 结构 MLP 残差Transformer Encoder 的每个 Block 由 Attention 和 MLP 组成ViT 用的是 Pre-Norm 结构也就是先做 LayerNorm 再进 Attention。相比 Post-NormPre-Norm 在训练初期更稳定梯度传播更顺畅这也是 ViT 能堆叠深度的原因之一。class Mlp(nn.Module): def __init__(self, dim, mlp_ratio4.0, dropout0.1): super().__init__() hidden_dim int(dim * mlp_ratio) self.fc1 nn.Linear(dim, hidden_dim) self.act nn.GELU() self.fc2 nn.Linear(hidden_dim, dim) self.drop nn.Dropout(dropout) 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 EncoderBlock(nn.Module): def __init__(self, dim, num_heads, mlp_ratio4.0, dropout0.1): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn Attention(dim, num_headsnum_heads, attn_dropdropout) self.norm2 nn.LayerNorm(dim) self.mlp Mlp(dim, mlp_ratio, dropoutdropout) def forward(self, x): x x self.attn(self.norm1(x)) x x self.mlp(self.norm2(x)) return xmlp_ratio4表示隐藏层维度是 embedding 维度的 4 倍这是 Transformer 的通用设置从 NLP 迁移过来的经验值。激活函数用的是 GELU它在负区间保留了非零梯度比 ReLU 平滑训练曲线更稳定。残差连接x attention(...)是必须的没有它深层模型很难收敛这也是为什么 Transformer 能堆到 12 层甚至 24 层。4.4 完整前向传播CLS token 的加入和位置编码ViT 的分类头不是对 196 个 token 做全局平均池化而是模仿 BERT 的做法在序列开头拼接一个特殊的 CLS token。这个 token 不承载任何具体 patch 的信息它只在训练中通过注意力机制聚合全局信息最终用于分类。class VisionTransformer(nn.Module): def __init__(self, img_size224, patch_size16, in_chans3, num_classes2, embed_dim192, depth12, num_heads3, mlp_ratio4.0, dropout0.1): super().__init__() self.patch_embed PatchEmbed(in_chans, img_size, patch_size, 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.randn(1, num_patches 1, embed_dim)) self.pos_drop nn.Dropout(dropout) self.blocks nn.Sequential(*[ EncoderBlock(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) def forward(self, x): B x.shape[0] x self.patch_embed(x) # [B, 196, 192] cls_tokens self.cls_token.expand(B, -1, -1) x torch.cat((cls_tokens, x), dim1) # [B, 197, 192] x x self.pos_embed x self.pos_drop(x) x self.blocks(x) # [B, 197, 192] x self.norm(x) cls_final x[:, 0] return self.head(cls_final)cls_token初始化为零向量这是刻意的。位置编码用randn随机初始化如果也用零前几轮训练所有 token 没有区分度梯度更新会变慢。expand把 shape 为[1, 1, 192]的 CLS token 复制到整个 batch。x self.pos_embed是 broadcast 加法pos_embed 的 shape 是[1, 197, 192]自动广播到 batch 维度。最终输出取x[:, 0]也就是每个样本的 CLS token 经过 12 层编码后的状态再过一个线性层映射到 2 个类别。整个 forward 从输入到输出一个 batch 的完整链路是卷积切 patch → 拼接 CLS token → 加位置编码 → 12 层 AttentionMLP → LayerNorm → 线性分类头。ViT-Small 的完整配置我用了一张表记录方便对照调试参数名值说明patch_size16每边切成 16 份共 196 个 patchembed_dim192patch 映射后的向量长度depth12Encoder Block 堆叠层数num_heads3注意力头数192/364 维每头mlp_ratio4.0MLP 隐藏层是 embed_dim 的 4 倍dropout0.1全连接层和注意力层的 dropoutnum_classes2猫狗二分类5. 训练与排查损失不降、过拟合、显存爆炸的处置记录5.1 损失曲线直接上天warmup 和初始学习率现象第一个 epoch 的 loss 从 0.7 左右瞬间跳到 3.5后面几个 epoch 才慢慢降下来前两轮基本在瞎猜。原因Transformer 对学习率极其敏感ViT 在视觉任务上没有卷积那样的平滑归纳偏置初始学习率太大AdamW 的动量会把参数推出合理范围。我一开始直接用了 3e-4 的固定学习率没加 warmup效果就是这样。解决改成前 5 个 epoch 做线性 warmup学习率从 1e-6 逐步升到 3e-4后面再用 cosine annealing 缓慢衰减。实现上用一个 step 计数器动态返回当前学习率即可。# train.py 中的学习率调度逻辑 def lr_scheduler(step, warmup_steps, total_steps, base_lr): if step warmup_steps: return base_lr * (step / warmup_steps) progress (step - warmup_steps) / (total_steps - warmup_steps) return base_lr * 0.5 * (1.0 torch.cos(torch.tensor(progress * 3.1415926))) # 每个 step 调用optimizer.param_groups[0][lr] lr_scheduler(...)参数说明warmup_steps在猫狗分类数据集上我设成总步数的 10%也就是说如果训练 50 个 epoch、batch 为 64、约 280 个 batch 每 epoch总步数是 14000warmup 就是 1400 步左右。base_lr用 3e-4AdamW 的 beta 和 weight decay 保持 PyTorch 默认值。5.2 验证集卡在 90% 上不去过拟合与标签噪声现象训练集准确率已经到 99%验证集卡在 88%~90% 不再上升损失在验证集上还出现了轻微上翘。原因两件事叠在一起。一是 ViT 参数多小数据集上容易把训练集的细节背下来二是猫狗数据集里有少量标签错误的图比如贵宾犬被剪成猫的形状这些噪声样本在训练后期成了“硬记忆”的对象。解决先加大数据增强把 RandomResizedCrop 的 scale 下限调到 0.5让模型看到更多局部细节再把模型自身的 dropout 从 0.1 提到 0.3尤其是 MLP 里的 dropout。标签噪声靠 label smoothing把标签从硬编码的 0/1 改成 0.95/0.05给模型留一点对错误标签的容忍度。self.criterion nn.CrossEntropyLoss(label_smoothing0.1) # 优化器设置bias 和 norm 层不做 weight decay param_groups [ {params: [p for n, p in model.named_parameters() if bias in n or norm in n], weight_decay: 0.0}, {params: [p for n, p in model.named_parameters() if bias not in n and norm not in n], weight_decay: 0.05}, ] optimizer torch.optim.AdamW(param_groups, lr3e-4)说明label_smoothing0.1意味着真实标签是 cat 时目标分布变成 cat0.95、dog0.05模型不会被强迫 output 无限接近 1对误标样本的梯度冲击会小很多。weight decay 只作用于权重矩阵不作用于 bias 和 LayerNorm 的参数这是 AdamW 优化 ViT 的通用做法能避免把归一化层的 scale 也压小。5.3 单卡训练太慢显存不够不是只能换显卡现象batch 设 128 直接 OOM设 64 又觉得 GPU 利用率不够训练一个 epoch 要 8 分钟50 个 epoch 要等大半天。原因ViT 的自注意力矩阵是[batch, heads, 197, 197]batch 翻倍时这一项显存也翻倍加上中间 12 层的激活值都保存在显存里做反向传播显存占用比同规模 CNN 高不少。解决两个手段组合。一是混合精度训练用torch.cuda.amp把前向和反向计算切到 fp16显存占用直接减半而且现在主流显卡的 tensor core 对 fp16 更快二是梯度累积batch 设 32每 2 个 batch 再更新一次参数模拟 batch64 的效果。scaler torch.cuda.amp.GradScaler() for step, (images, labels) in enumerate(train_loader): images, labels images.cuda(), labels.cuda() with torch.cuda.amp.autocast(): outputs model(images) loss criterion(outputs, labels) scaler.scale(loss).backward() if (step 1) % accumulation_steps 0: scaler.step(optimizer) scaler.update() optimizer.zero_grad()混合精度训练有两个容易翻车的细节loss 必须用GradScaler.scale(loss).backward()如果手动.backward()再用 scaler.step梯度会精度不足每次scaler.step之后要调用scaler.update()否则动态 loss scale 不会自动调整。梯度累积时注意optimizer.zero_grad()的位置它应该每个累积周期执行一次而不是每个 batch 都执行。5.4 一张图预测时结果反复横跳现象训练时验证集准确率还行但单张图片预测时换个预处理尺寸结果就从猫变成狗同一张图 reszie 成不同大小结果不稳定。原因验证集用的是 Resize(256)CenterCrop(224)而单图推理时可能直接 Resize 成 224×224破坏了长宽比。猫狗图片很多是竖构图或横构图强行拉伸到正方形会导致特征畸变。解决infer.py 里和验证集保持完全相同的预处理流水线Resize(256) 后再 CenterCrop(224)。另外训练时的 RandomResizedCrop 已经让模型见过不同尺度的目标单图推理时 follow 验证集的均匀裁剪是最稳妥的。6. 二次开发换数据集、可视化注意力与部署加速6.1 换一个数据集要改什么attention 可视化怎么做ViT 换数据集改动比 CNN 少主要动三处num_classes改成新任务的类别数数据目录替换成新数据集最后重新统计 mean/std。但有一个坑必须注意——位置编码的尺寸依赖 patch 数。如果新数据集不是 224×224比如换成了 384×384位置编码的 shape 会不匹配。此时要么把输入统一 resize 到 224×224要么对位置编码做双线性插值后者实现稍微复杂我一般优先选择前者。注意力可视化是理解 ViT 有没有学到东西的直接手段。把第 4 章 Attention 里的attn权重抽出来对多头取平均就能得到每个 patch 对其他 patch 的注意力分布def visualize_attention(model, img_tensor): model.eval() with torch.no_grad(): x model.patch_embed(img_tensor) cls_token model.cls_token.expand(x.shape[0], -1, -1) x torch.cat((cls_token, x), dim1) x x model.pos_embed for i, blk in enumerate(model.blocks): # 取最后一个 block 的注意力层手动前向一次 if i len(model.blocks) - 1: attn blk.attn(blk.norm1(x)) attn_map attn.mean(dim1) # [B, 197, 197] cls_attn attn_map[0, 0, 1:].reshape(14, 14) cls_attn torch.nn.functional.interpolate( cls_attn.unsqueeze(0).unsqueeze(0), size(224, 224), modebilinear ) # 将 cls_attn 归一化到 [0,1] 后与原始图片叠加显示这段代码里attn.mean(dim1)是把 3 个头的注意力矩阵做平均cls_attn[0, 0, 1:]取的是 CLS token 对所有 patch 的注意力也就是模型分类时“看”的位置。插值回 224×224 后叠加在原图上如果猫脸区域的注意力值明显高于背景说明 ViT 学到了真实的判别特征而不是在靠背景纹理硬猜。6.2 部署加速导出 ONNX 和固定输入尺寸训练完要落地到服务端一般不会直接用 PyTorch 模型跑推理因为 Python 环境和 GIL 会限制并发。这个项目导出 ONNX 的改动很小用torch.onnx.export即可注意把输入尺寸固定ViT 的序列长度可变但位置编码是固定的动态尺寸对 Transformer 的分类模型收益不大固定成 224×224 最省心。导出后可以用 ONNX Runtime 的 Python 接口做一次精度对齐比对 PyTorch 和 ONNX 的 softmax 输出差值在 1e-4 以内就说明导出没出问题。如果后续要跑 TensorRTONNX 是现成的中间格式配合 fp16 精度单张图推理时间从几毫秒级别优化到亚毫秒级这对线上服务是有意义的优化。最近一次换新数据集时我图省事直接沿用了旧任务的 mean/std结果 loss 开头两轮完全不降日志打印出来的 batch 均值明显偏离 0 和 1。从那以后我每次换数据集都会先单独统计 one epoch 的均值方差再打印一版增强后的样本图确认细节还在最后才开训练——跑通之后才意识到ViT 对数据分布比 CNN 更敏感这步检查省不得。希望这个项目的拆解过程能帮你在自己的数据上少踩几个类似的坑。本文还有配套的精品资源点击获取