★ FEATURED ARTICLE 简介面向Transformer初学者的视觉TransformerViT实战项目以经典的猫狗数据集完成图像分类全流程。项目覆盖数据读取、模型搭建、训练与评估等关键环节代码结构清晰、注释到位适合零基础或刚接触注意力机制的读者。更重要的是该方案不局限于猫狗分类只需更换数据集并调整类别数目参数即可迁移至花卉、车型等任意图像分类任务展现良好泛化能力。资源共有2000个文件包含1998张JPG图片与2个Python脚本压缩包体积约218.41MB图片数据支持直接用于训练验证脚本则实现了完整的ViT分类流程方便直接运行或二次开发。目前已有1469人浏览学习热度较高。通过本项目读者既能深入理解ViT原理与结构又能获得一套可复用的图像分类模板将其灵活应用到自己的研究或课程设计中为后续探索Transformer在视觉领域的应用打下扎实基础。1. Vision Transformer(ViT)实践猫狗分类入门Transformer视觉任务值得做吗很多从卷积网络切到Transformer的人第一反应都觉得把一张猫图切成小方块再当句子喂给模型很反直觉。但Vision Transformer(ViT)恰恰证明这条路在图像分类上走得通猫狗分类这种二分类任务是验证它有没有真正学到全局特征的极佳试金石。我之前接手一个图像分类Demo时就用这个任务把ViT全流程跑了一遍切Patch、位置编码、微调、看损失、排查过拟合。下面就把这条链路拆开说清每个参数为什么那样设以及几个容易翻车的地方。适合想从CNN转ViT的开发者也适合刚拿到视觉数据但拿不准要不要上Transformer的从业者。2. 把ViT拆开看关键变量Patch、注意力机制与猫狗分类的选型理由ViT和CNN最大的不同是一上来就把图像拆成一组有序向量而不是在局部窗口里做卷积。这个结构决定了它对全局关系的建模天然更强但代价是需要更多数据或更强的预训练权重否则很容易过拟合。猫狗分类恰好适合用来观察这种差异猫脸、狗脸、耳朵、背景之间的空间关系比单纯纹理细节更值得模型去建模。2.1 从Patch到Token图像是怎么变成序列的ViT把输入切成固定大小的Patch。以patch_size16、输入224×224为例图像被切成14×14196个Patch。每个Patch展平成16×16×3768维向量经过一层线性投影变成768维的embedding再额外拼接一个CLS token序列长度变成197最后加上一个可学习的位置编码位置编码的维度和embedding一样是768。这样一张图就变成一条“token序列”交给后面的Transformer Encoder处理。CLS token很重要。它不像CNN那样直接输出特征图ViT需要这个固定位置的token来汇总全局信息。原理和BERT里的CLS类似最后一层多头自注意力会把所有Patch的信息聚合到CLS上再通过分类头映射到猫/狗两个类别。这里有一个可视化好处attention map可以直接看出模型在“盯”图像哪个位置这个能力在CNN里往往需要额外工具才能做到。ViT里几个关键参数的默认值参数常见默认值在猫狗分类里的影响patch_size16越小序列越长局部细节保留越好但显存和耗时上升embed_dim768每个token的特征维度越大表达力越强也越吃显存depth / layers12Encoder层数影响全局交互能力num_heads12多头注意力个数多个头会关注不同区域关系mlp_ratio4FFN中间层维度是embed_dim的4倍patch_size尽量不要随便改。如果你把patch改成8序列长度会变成原来的4倍注意力计算量和显存会大幅上升。在猫狗分类这种主体占画面大部分的任务里patch_size16是平衡点如果目标是远处的小猫或小狗可以换用patch_size8的变体但要先确认显存够用。位置编码是可学习的它的作用是让模型知道第三个patch大概对应原图哪一行哪一列。ViT抛弃了卷积的滑动窗口相当于放弃“相邻像素有关系”的归纳偏置位置编码就是找回这部分空间信息的手段。这也解释了为什么ViT对位置编码插值、图像分辨率变化会比较敏感后面会说这个坑。2.2 猫狗分类选ViT的三个前提数据量、算力与可解释性很多人觉得ViT效果一定比CNN好但实际要从三方面看一是数据量。从零训练ViT在猫狗数据上会非常吃力至少需要百万级图片量级才能体现优势。现在能快速跑通基本都靠ImageNet预训练权重在几万张猫狗图上做微调效果可以做到和CNN相当甚至更好。如果手里只有两三千张图还是建议先用CNN跑一个baseline再用ViT对比否则ViT很容易记住背景而不是猫狗本身。二是算力。ViT base模型参数量大概在86M左右比常见CNN骨干大不少前向和反向都更吃显存。8GB显存开大batch很容易OOM需要梯度累积或者激活检查点来救。如果机器实在跑不动可以换vit_small_patch16_224甚至vit_tiny_patch16_224效果下降幅度在二分类任务里可接受。三是调试门槛。代码本身并不比CNN难一行create_model就能加载模型门槛主要在训练参数上。ViT的loss曲线对学习率和warmup极敏感调不好表现为loss震荡、val acc卡在50%。这一点在避坑章节会反复提到。维度CNNViT局部纹理建模强一般全局关系、长距离依赖弱强小数据从零训练相对容易很难预训练微调稳定效果好但对lr敏感注意力可视化需要额外工具直接可看所以我的判断是如果你是第一次在视觉任务上实验Transformer猫狗分类这种二分类是试错成本最低的起点但前提是使用预训练权重。选型时看数据量、显存、可调参数这三点缺一个都会在训练阶段多吃苦头。3. 从数据准备到训练闭环猫狗分类用ViT微调的完整流程这一章直接给可复现的流程。数据准备、数据增强、训练循环、微调策略按顺序走踩坑点会标注在对应位置。3.1 数据集组织与标签解析先花十分钟检查目录结构拿到猫狗分类数据第一件事不是写模型而是看清目录结构。常见的是train目录下再分cat和dog两个子目录val目录同理这种结构可以直接用现成的ImageFolder。但有些来源的数据是单目录平铺文件名里带Cat和Dog路径大小写还混着这种情况需要自己写Dataset否则标签很容易错乱。data/ ├── train/ │ ├── cat/ │ │ ├── cat.1.jpg │ │ └── cat.2.jpg │ └── dog/ │ ├── dog.1.jpg │ └── dog.2.jpg └── val/ ├── cat/ └── dog/如果目录不规范我习惯写一个轻量Dataset而不是先花时间重排所有文件from torch.utils.data import Dataset from PIL import Image class CatDogDataset(Dataset): def __init__(self, root, transformNone): # root 指向 train/ 或 val/ self.paths sorted([str(p) for p in root.glob(*/*.jpg)]) self.transform transform def __len__(self): return len(self.paths) def __getitem__(self, idx): path self.paths[idx] # 统一把反斜杠转成斜杠再用 /cat/ 判断更保险 norm_path path.replace(\\, /) label 0 if /cat/ in norm_path else 1 img Image.open(path).convert(RGB) if self.transform: img self.transform(img) return img, label这段代码有几个值得注意的点。用root.glob(/.jpg)而不是glob.glob(str(root / / .jpg))是为了跨平台兼容Windows和Linux下路径分隔符不同。判断标签时用/cat/而不是cat in path是因为文件名可能叫dogcat_5.jpg用子目录判断更稳。Image.open之后必须convert(RGB)猫狗数据集里偶尔混入灰度图不转RGB会直接导致后续通道数不匹配。3.2 数据增强和Normalize训练和验证必须两套transformViT微调的数据增强直接影响收敛效果。训练和推理要用两套完全不同的transform这是老生常谈但每次还是有人把RandomResizedCrop用到测试集上。from torchvision import transforms train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.6, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])RandomResizedCrop的scale参数默认是(0.08, 1.0)这个默认值在ImageNet分类问题不大但猫狗分类里主体通常占据画面大部分裁剪比例太低会频繁裁出大片背景毛色等于给模型增加无效噪声。我把scale下限提到0.6收敛明显更稳。ColorJitter三个通道的值不要拉得太猛0.2对猫狗毛色来说已经足够太强反而让模型学到颜色不变性之外的无用信息。提示这里的Normalize均值和标准差对应预训练权重的统计值不要自己改成(0.5, 0.5, 0.5)。改了之后模型看到的数据分布和预训练时不一致效果会掉得莫名其妙。3.3 训练闭环从create_model到OneCycleLR模型我用timm加载读者自己复现时注意版本差异timm.create_model会下载预训练权重第一次运行需要联网如果是离线环境需要提前下载权重并手动指定。核心训练代码如下import torch import torch.nn as nn import timm model timm.create_model(vit_base_patch16_224, pretrainedTrue, num_classes2) device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) criterion nn.CrossEntropyLoss(label_smoothing0.05) optimizer torch.optim.AdamW(model.parameters(), lr1e-5, weight_decay0.05) scheduler torch.optim.lr_scheduler.OneCycleLR( optimizer, max_lr3e-5, epochsepochs, steps_per_epochlen(train_loader), pct_start0.15 ) for epoch in range(epochs): model.train() total_loss 0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() with torch.cuda.amp.autocast(): logits model(images) loss criterion(logits, labels) loss.backward() optimizer.step() scheduler.step() total_loss loss.item() avg_loss total_loss / len(train_loader) val_acc evaluate(model, val_loader, device) print(fepoch {epoch1}/{epochs} loss {avg_loss:.4f} val_acc {val_acc:.4f})pretrainedTrue是猫狗分类能快速收敛的关键直接随机初始化ViT base在这个数据量上很容易过拟合。CrossEntropyLoss里的label_smoothing0.05能起到轻微正则作用ViT参数量大这一点平滑对验证集有帮助。AdamW的lr我建议设1e-5~3e-5之间超过1e-4大概率会把预训练特征冲坏。OneCycleLR的max_lr我设成3e-5optimizer初始lr是1e-5曲线会先warmup到max_lr再开始衰减pct_start0.15表示前15%的step用于warmup。这里特别容易踩的坑是scheduler.step()放在每个batch后不要放在epoch循环里否则warmup节奏完全是乱的。3.4 微调策略与位置编码哪些该训练、哪些不该动ViT微调策略可以分两档。数据量中等、显存够时直接全量微调效果最好所有参数都参与更新。数据量很小或明显过拟合时冻结Patch Embedding和前几层Encoder只训练分类头和最后几层能有效降低模型容量。# 冻结前6层Encoder只训练最后6层和分类头 for name, param in model.named_parameters(): if blocks.0. in name or blocks.1. in name or blocks.2. in name: param.requires_grad False常见做法是先冻结跑10个epoch观察验证集准确率停滞后再解冻最后几层继续微调。这样比一开始就全量训练更稳尤其适合训练样本几千张的场景。位置编码和图像分辨率的关系容易被忽略。vit_base_patch16_224的pos_embed只在224×224下有效如果强行把输入改成384×384序列长度变了位置编码对不上。timm在较新版本里会对pos_embed做插值但不同版本实现细节有差异。如果实在想用384分辨率优先选官方支持384输入的预训练变体比如vit_base_patch16_384不要自己手动改pos_embed然后去猜有没有问题。4. ViT猫狗分类常见问题与排查别让五个细节吃掉准确率这一章写的是我实际跑猫狗分类时反复踩过的坑。每一条都按现象、原因、解决展开你可以直接对照自己的训练日志排查。4.1 loss震荡、val acc卡在50%先查学习率与warmup现象训练2到3个epoch后loss在0.65到0.70之间来回跳验证集准确率始终在50%附近模型等于什么都没学到。原因学习率太大或者完全没有warmup。ViT的qkv投影在预训练后分布比较紧微调一开始用大学习率会把整个特征空间顶出原有区域。很多人直接把CNN训练用的1e-3搬过来loss自然震荡。解决把AdamW的lr降到2e-5左右并在一开始的5%到15%步数内做warmup。用OneCycleLR时lr从1e-5升到3e-5再衰减基本能解决这个问题。如果你用的是固定lr也建议手动把前几百步的lr乘一个小的warmup系数。4.2 训练集准确率很高、验证集上不去增强太弱与过拟合现象训练后期train acc接近99%val acc却只有70%上下而且两者差距越拉越大。原因ViT base有约86M参数猫狗数据量有限时模型会把训练集里的背景、水印、拍摄风格一起记住。只用Resize加ToTensor的配置训练ViT过拟合几乎是必然的。解决数据增强至少包含RandomResizedCrop加RandomHorizontalFlip再加一点ColorJitter。把scale下限从0.08提升到0.6避免模型靠背景毛色分类。CrossEntropyLoss加label_smoothing0.05ViT的Dropout从默认0.1提升到0.2。最重要的还是加载预训练权重而不是从头训练。如果数据集实在太小回到3.4的冻结策略。4.3 显存OOM别急着调小batch试试梯度累积现象8GB显存跑ViT basebatch_size设32直接爆显存改成8之后勉强能跑但验证集准确率掉了两三个点。原因AdamW的动量参数本身就占显存ViT的attention中间变量又大batch太小会让每个batch的梯度估计噪声变大表现为训练震荡。解决用梯度累积把有效batch补回来。显存装下batch_size8就固定8把累积步数设成4等效batch_size32accumulation_steps 4 optimizer.zero_grad() for step, (images, labels) in enumerate(train_loader): images, labels images.to(device), labels.to(device) with torch.cuda.amp.autocast(): logits model(images) loss criterion(logits, labels) / accumulation_steps loss.backward() if (step 1) % accumulation_steps 0: optimizer.step() optimizer.zero_grad()loss除以accumulation_steps是关键这样累积后的梯度均值等价于大batch。如果显存还是不够再考虑torch.utils.checkpoint激活检查点但训练时间会变长也可以直接换vit_small或vit_tiny。这里有个习惯问题如果用了AMP的GradScalerscaler.step和scaler.update也要放在累积完成的位置否则loss又会出现NaN。4.4 AMP混合精度loss变NaN问题不一定在数据现象训练前几轮一切正常到第10个epoch左右loss突然变成NaN之后整个训练报废。原因fp16动态范围有限ViT的LayerNorm和自注意力里大数值很容易溢出另一部分情况是用了torch.cuda.amp但没有正确处理梯度缩放喂给优化器的梯度已经是inf。解决先确认AMP写法是不是完整的三件套。标准写法是scaler torch.cuda.amp.GradScaler() with torch.cuda.amp.autocast(): logits model(images) loss criterion(logits, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()如果已经这么写还是NaN优先尝试退回到fp32。ViT微调在猫狗分类上不用AMP也能接受速度慢一点但稳定性高很多。如果硬件支持bfloat16可以用torch.autocast(device_typecuda, dtypetorch.bfloat16)动态范围更大基本不会溢出。4.5 验证集正常但真实图片翻车推理transform和训练时必须分开现象训练集和验证集指标都不错但随便拿一张手机拍的真实猫图进去预测结果错了而且同一张图跑两次结果还不一样。原因最常见的是把训练时的RandomResizedCrop、RandomHorizontalFlip也用在推理上。推理时随机裁剪会让模型看到不同区域输出必然抖动。解决推理单独定义eval_transform只做Resize(256)、CenterCrop(224)、Normalize。输入图片如果是PNG带alpha通道或者灰度图务必先convert(RGB)。另外一个容易忽略的细节如果训练时RandomResizedCrop的scale用了(0.6, 1.0)而推理用Resize(256)CenterCrop(224)两者看到的物体占比其实不太一致可能掉1到2个点。可以尝试把推理Resize从256改到280再做CenterCrop有时候能弥补这个gap。从那时起我把eval_transform单独封装成一个函数训练、验证、推理各用各的这类问题再没出现过。5. 猫狗分类之外单张推理、注意力可视化与迁移到自己的数据跑通训练只是第一步真正落到使用场景还需要把模型拿出来做推理、确认模型学到的东西、以及换到自己的数据上。5.1 用训练好的模型做单张推理细节决定稳定性推理代码不长但有几个常见错误值得注意。加载权重时如果之前用DataParallel训练保存的state_dict会带module.前缀直接load会报key不匹配。state torch.load(catdog_vit.pth, map_locationdevice) new_state {k.replace(module., ): v for k, v in state.items()} model.load_state_dict(new_state) model.eval() image val_transform(Image.open(test.jpg).convert(RGB)).unsqueeze(0).to(device) with torch.no_grad(): logits model(image) prob torch.softmax(logits, dim-1)[0] pred int(logits.argmax(dim-1)[0])model.eval()会关闭Dropout和Attention Dropout这个状态在推理时不能省。with torch.no_grad()可以减少显存和内存占用同时避免累积计算图。输出的prob是类别概率取argmax得到0或1分别对应猫和狗。5.2 Attention Map可视化验证模型看的是猫还是背景ViT自带注意力机制可以直接把最后层的attention权重拿出来看模型到底“盯”在哪里。常见做法是取最后一个Encoder里CLS token对各个Patch的注意力然后还原成14×14的网格再上采样到224×224。这个操作在不同版本的模型接口里略有差异可以给一个整体流程参考# attn_cls 形状为 [heads, tokens, tokens] # 取CLS对所有Patch的注意力去掉CLS自身 attn_map attn_cls.mean(dim0)[0, 1:].reshape(14, 14) attn_map F.interpolate(attn_map[None, None], size(224, 224), modebilinear)得到的热力图叠加到原图上以后如果高亮区域集中在猫耳朵或狗鼻子附近说明模型学到了有效的视觉特征如果高亮区域集中在背景墙面或地板上说明训练数据里存在很强的背景bias模型再调参也要先考虑数据问题。对猫狗分类来说这是我用来判断“有没有训练歪”的最快手段。5.3 换到自己的数据哪些参数必须动哪些可以不用动从猫狗二分类切到自己的项目第一反应是只把num_classes改掉这个方向对但不够。如果新数据是十类分类头需要重新随机初始化主干保持预训练权重分类头的lr可以给到主干的10倍让新分类头更快收敛。如果类别数量不均衡不能只看accuracy要看每个类的recall和precision必要时用WeightedRandomSampler或对样本少的类别过采样。如果任务变成了多标签“猫还是狗”的二分类逻辑就不适用了CrossEntropyLoss要换成多标签损失输出层激活函数也要从Softmax改成Sigmoid。换数据之后最好把5.2的Attention Map可视化流程再跑一遍。我自己每次换数据集都逼自己先看一眼模型是不是真的在盯目标对象再继续调参这个方法帮我避开了好几次假精度。从那以后我每次做迁移实验都会强制走这一遍希望帮到你。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助？
