简介面向深度学习课程设计与期末大作业场景基于Vision TransformerViT实现花卉图像分类的完整Python项目适合需要快速掌握ViT实际训练流程的初学者也可作为高分课程设计的参考。代码采用模块化组织共7个文件以6个Python脚本为主体另含1个keep占位文件压缩包整体仅10KB结构轻量但功能闭环。各脚本职责清晰数据加载与预处理负责图片读取和增强模型模块封装ViT主干结构工具函数提供通用辅助训练脚本管理损失计算与训练循环分类预测脚本用于推理参数配置统一集中便于调节超参数和路径。代码注释详尽关键逻辑均有说明初学者也能对照理解ViT在图像分类中的落地方案。项目下载后即可运行同时预留二次开发空间可替换数据集、调整模型配置进一步扩展。目前已有383人学习适合作为期末大作业、课程设计或图像分类入门实践的起步模板。1. 基于ViT来进行图像分类花卉识别拆解这份97分Python大作业代码这份python基于ViT来进行图像分类花卉识别的大作业代码核心是用Vision Transformer把花卉图片按品种分类整个项目由Modules目录下的config.py、DataLoader.py、model.py、train.py、DoClassify.py、utils.py几个脚本串起来解压就能跑。它解决的是期末大作业里最实在的问题结构完整、注释清楚、训练与推理链路闭环所以能在课程设计里拿到97分这档高分。适合三类人——正为Transformer课程设计发愁的本科生、已经把CNN玩熟想换ViT做图像分类的新手以及想用现成代码二开成自己数据集项目的同学。接下来从数据输入、模型搭建、训练调参一路拆到单图推理和避坑。2. 数据侧怎么做config.py和DataLoader.py把图片变成模型输入的完整链路数据链路主要涉及两个文件config.py把路径、尺寸、训练超参一次性定清楚DataLoader.py把磁盘图片读进来做随机增强、归一化最后产出模型需要的 [batch, 3, 224, 224] 张量。别小看这个环节——ViT对输入尺寸和归一化很敏感见过太多项目在模型上没翻车反而在数据加载阶段卡了一下午。2.1 config.py把可变参数集中成一个类很多同学喜欢在每个文件里直接写死数据路径train.py一套、DoClassify.py又复制一套。等数据集挪了位置换台电脑跑改起来四散各处。集中写到Config类里后面所有模块引用同一个类改一处、处处生效这也是课程设计评分里很加分的习惯。# Modules/config.py import torch class Config: # 数据与输出目录 data_root data/flower_photos # 数据集根目录 save_dir checkpoints # 模型保存目录 # 输入与分类头 image_size 224 # ViT B/16 的标准输入尺寸 num_classes 5 # 花卉类别数和数据集子文件夹数严格一致 # 训练超参数 epochs 50 batch_size 32 lr 3e-4 weight_decay 1e-4 warmup_epochs 5 # ViT 结构参数 patch_size 16 embed_dim 768 num_heads 12 depth 12 dropout 0.1 device cuda if torch.cuda.is_available() else cpu拿到手必须改的地方有三个data_root换成你的数据集实际路径num_classes要和数据集子文件夹数量一致否则分类头尺寸对不上batch_size和lr按显卡显存和loss表现调。image_size224是ViT-B/16的标准配置切成长宽各14个区域得到196个patch序列再加上一个cls_token共197个token。参数建议值控制什么调参方向lr3e-4每步参数更新幅度loss震荡就调小收敛过慢就调大batch_size32单次喂入样本数与显存占用显存不够降到16或8warmup_epochs5训练初期学习率爬升的时长数据多可减到2数据少可增到8patch_size16每个patch的边长显存敏感时可试32精度会有损失提示data_root用相对路径时以你敲python train.py的目录为准。建议数据集和代码放在同一个根目录下避免到处找不到路径。2.2 数据集目录结构ImageFolder的硬性要求DataLoader.py里最核心的类是torchvision.datasets.ImageFolder工作方式很直接给它一个根目录它自动把根目录下每个子文件夹当作一个类别并按字母顺序编号。所以数据集必须规规矩矩按下面这种结构放data/ └── flower_photos/ ├── daisy/ # 编号 0 ├── dandelion/ # 编号 1 ├── rose/ # 编号 2 ├── sunflower/ # 编号 3 └── tulip/ # 编号 4然后DataLoader.py里做的就是读取目录、套transform、包成DataLoader# Modules/DataLoader.py import torch from torch.utils.data import DataLoader, random_split from torchvision import datasets, transforms from config import Config def build_transform(modetrain): if mode train: return transforms.Compose([ transforms.Resize(256), # 先放大一点再随机裁剪 transforms.RandomCrop(Config.image_size), # 裁剪到 224x224 transforms.RandomHorizontalFlip(p0.5), # 随机翻转等价于增加样本 transforms.ToTensor(), # [0,255] 转 [0,1] 张量 transforms.Normalize( # 按 ImageNet 统计值归一化 mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225] ) ]) else: return transforms.Compose([ transforms.Resize((Config.image_size, Config.image_size)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) def build_dataloader(): train_ds datasets.ImageFolder( rootConfig.data_root /train, transformbuild_transform(train) ) val_ds datasets.ImageFolder( rootConfig.data_root /val, transformbuild_transform(val) ) train_loader DataLoader( train_ds, batch_sizeConfig.batch_size, shuffleTrue, num_workers2, pin_memoryTrue ) val_loader DataLoader( val_ds, batch_sizeConfig.batch_size, shuffleFalse, num_workers2, pin_memoryTrue ) return train_loader, val_loader, train_ds.classes注意这里的三点设计。第一训练集和验证集用两套transform验证集不做随机翻转和裁剪否则每个epoch对同一张图片的评估结果会上下浮动没法准确比较模型好坏。第二验证集shuffleFalse保证每批数据和标签顺序固定方便后面统计逐类准确率。第三函数最后把train_ds.classes返回出去这是ImageFolder排好序的类别名列表后面训练完保存下来就是推理时的标签依据。这里用的Resize(256)RandomCrop(224)相当于一种数据增强先从256缩放图中随机裁一块224区域每次训练看到同一张图的不同局部这比直接Resize(224)多了一些随机性。熟手可以把RandomRotation(15)也加进去但旋转角度别太大花卉图片旋转太多会引入无关背景干扰。2.3 随机拆分还是按目录拆分验证集比例怎么定如果你的数据集暂时没有分好train/val也可以用random_split按比例拆。两种方式各有利弊课程设计里要根据你的实验结果期待来选方式适用场景优点缺点random_split数据集只有一个根目录时快速出基线一行代码、比例随意调相似图片可能同时进训练和验证分数偏高、不够可信train/val目录分开比赛、论文、要写实验报告的作业互不干扰结果经得起推敲需要自己按类别搬运图片我一般建议如果这份作业要写实验报告、答辩时老师会追问数据划分就用目录分开的方式。随机拆分的结果多少带点乐观偏差老师经验丰富翻车的概率不小。验证集比例常规是8:2如果你的花卉数据集每类只有几十张建议用到9:1尽量多留训练样本给ViT。3. 模型与训练侧model.py里ViT网络与train.py训练循环的落地细节数据链路通了之后重头戏是model.py和train.py。model.py把 [batch, 3, 224, 224] 的图片张量映射成 [batch, num_classes] 的分类得分train.py通过反向传播把这套参数调到一个能用的状态。目前ViT的主流技术路线有两种直接用timm或torchvision现成模型改了分类头就跑以及自己实现Patch Embedding、Transformer Block把网络结构白纸黑字写在model.py里。想把分数做高、答辩讲得清至少需要自实现一版。3.1 自实现ViTPatch Embedding、Transformer Block与分类头ViT和CNN最大的差异在第一层就体现出来CNN用卷积核滑动扫描局部区域ViT则把图片切成固定大小的patch每个patch展平后线性投影成token向量再加位置编码和cls_token。下面这个Patch Embedding模块就是这个过程# Modules/model.py —— Patch Embedding 部分 import torch import torch.nn as nn class PatchEmbedding(nn.Module): def __init__(self, patch_size16, in_channels3, embed_dim768): super().__init__() self.patch_size patch_size self.proj nn.Conv2d(in_channels, embed_dim, kernel_sizepatch_size, stridepatch_size) self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed nn.Parameter(torch.zeros(1, 197, embed_dim)) def forward(self, x): B x.shape[0] # [B, 3, 224, 224] - [B, 768, 14, 14] - [B, 196, 768] x self.proj(x).flatten(2).transpose(1, 2) # 拼接类别 token [B, 197, 768] x torch.cat([self.cls_token.expand(B, -1, -1), x], dim1) # 加位置编码 x x self.pos_embed return x这里用一个卷积层实现patch切分是个很省事的技巧kernel_size和stride都等于patch_size时卷积天然就把图片切成非重叠patch且完成了线性投影。196个patch加上1个cls_token正好对应197的位置编码长度。cls_token是Transformer里用来汇总整张图信息的可学习向量全序列的attention结果最终都会传导到它身上。接下来是Transformer BlockViT的核心计算单元重复堆叠12层# Modules/model.py —— Transformer Block 部分 class TransformerBlock(nn.Module): def __init__(self, dim768, num_heads12, mlp_ratio4, dropout0.1): super().__init__() self.norm1 nn.LayerNorm(dim) self.attn nn.MultiheadAttention(dim, num_heads, dropoutdropout, batch_firstTrue) self.norm2 nn.LayerNorm(dim) self.mlp nn.Sequential( nn.Linear(dim, dim * mlp_ratio), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim * mlp_ratio, dim), ) def forward(self, x): # 残差 多头自注意力 x x self.attn(self.norm1(x))[0] # 残差 前馈网络 x x self.mlp(self.norm2(x)) return xself.attn(self.norm1(x))[0]取的是nn.MultiheadAttention返回的第一个值也就是加权求和后的特征。这里用batch_firstTrue可以让输入的张量形状保持 [batch, seq_len, embed_dim]和上面Patch Embedding的输出直接接上。LayerNorm放在attention之前是Pre-LN结构相比Post-LN训练更平稳ViT基本都采用这种写法。最后把这几层拼成一个完整的分类模型# Modules/model.py —— 完整模型 class ViTClassifier(nn.Module): def __init__(self, img_size224, patch_size16, in_channels3, embed_dim768, depth12, num_heads12, num_classes5): super().__init__() self.patch_embedding PatchEmbedding(patch_size, in_channels, embed_dim) self.blocks nn.Sequential( *[TransformerBlock(embed_dim, num_heads) for _ in range(depth)] ) self.norm nn.LayerNorm(embed_dim) self.head nn.Linear(embed_dim, num_classes) def forward(self, x): x self.patch_embedding(x) x self.blocks(x) x self.norm(x) # 只取 cls_token 位置的输出做分类 x self.head(x[:, 0]) return x为什么x[:, 0]取的是cls_token因为patch位置经过12层attention后已经把全局上下文融合进了cls_token它相当于整个序列的语义摘要。输出的维度是 [batch, num_classes]配合CrossEntropyLoss正好一个batch过一遍做反向传播。注意这个模型的depth、num_heads、embed_dim、patch_size都和config.py里的数值对应。如果改了config的depth这层nn.Sequential的重复次数也会变请保持一致。3.2 train.py训练循环里值得抄的几个细节train.py的逻辑和CNN训练代码很相似区别主要在优化器选择和保存策略上。ViT的标准优化器是AdamW权重衰减对Transformer的收敛有明显帮助学习率调度这里用Cosine Annealing让loss在训练后期慢慢降下去不震荡。# Modules/train.py import json import torch import torch.nn as nn from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR from tqdm import tqdm from config import Config from DataLoader import build_dataloader from model import ViTClassifier from utils import set_seed, plot_history set_seed(42) train_loader, val_loader, classes build_dataloader() model ViTClassifier(num_classesConfig.num_classes).to(Config.device) criterion nn.CrossEntropyLoss() optimizer AdamW(model.parameters(), lrConfig.lr, weight_decayConfig.weight_decay) scheduler CosineAnnealingLR(optimizer, T_maxConfig.epochs) # 保存类别映射推理时要用 with open(classes.json, w, encodingutf-8) as f: json.dump(classes, f) best_acc 0.0 for epoch in range(Config.epochs): model.train() total_loss, correct, total 0, 0, 0 loop tqdm(train_loader, descfEpoch {epoch1}/{Config.epochs}) for images, labels in loop: images, labels images.to(Config.device), labels.to(Config.device) output model(images) loss criterion(output, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() correct (output.argmax(1) labels).sum().item() total labels.size(0) loop.set_postfix(lossloss.item()) # 验证集评估 model.eval() val_correct, val_total 0, 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(Config.device), labels.to(Config.device) output model(images) val_correct (output.argmax(1) labels).sum().item() val_total labels.size(0) train_acc correct / total val_acc val_correct / val_total print(fEpoch {epoch1}: train_acc{train_acc:.4f}, fval_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), f{Config.save_dir}/best_model.pth) # 这里顺便把最后一条曲线数据存下来后续绘图用 plot_history(train_loss_list, val_loss_list, train_acc_list, val_acc_list) scheduler.step()上面代码里output.argmax(1)返回每个样本得分最高的类别索引与标签比较后得到正确数。数据加载阶段shuffleTrue训练时每轮的batch顺序都在变减轻过拟合。model.eval()和torch.no_grad()要成对出现前者切换BN和Dropout模式后者关掉梯度计算验证集跑起来能省不少显存和耗时。每次验证acc刷新就torch.save一次state_dict这是一个能救命的习惯。我一般不只保存best_model.pth还会在最后一轮额外保存last_model.pth。因为最佳模型未必出现在最后一个epoch如果最后几轮过拟合了没有best保存你只能从头再跑有best模型在手相当于给自己留了后悔药。3.3 utils.py固定随机种子和训练曲线答辩时拿得出证据utils.py在这个项目里承担两个职责固定随机种子和绘制训练曲线。固定种子看起来不起眼但在课程设计中非常关键——同一个项目今天跑85分、明天跑81分你写实验报告都不知道该写哪个数。固定种子之后结果可复现老师复现你的代码时也能拿到和你报告一致的数据。# Modules/utils.py import random import matplotlib.pyplot as plt import torch def set_seed(seed42): random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def plot_history(train_loss, val_loss, train_accNone, val_accNone, save_pathtrain_curve.png): plt.figure(figsize(12, 4)) plt.subplot(1, 2, 1) plt.plot(train_loss, labeltrain loss) plt.plot(val_loss, labelval loss) plt.xlabel(epoch) plt.ylabel(loss) plt.legend() if train_acc is not None: plt.subplot(1, 2, 2) plt.plot(train_acc, labeltrain acc) plt.plot(val_acc, labelval acc) plt.xlabel(epoch) plt.ylabel(acc) plt.legend() plt.savefig(save_path, dpi150)训练结束后把train_curve.png贴进实验报告比任何文字描述都直观。如果你的数据量不大自实现ViT从零训练容易偏低可以换成timm里现成的预训练模型把权重加载进来做finetune一行代码搞定model timm.create_model(vit_base_patch16_224, pretrainedTrue, num_classesConfig.num_classes)。这套做法适合时间紧、只想把分数刷高的场景自实现版本更适合答辩时讲清原理两个方向不冲突。4. 避坑排查跑通ViT花卉识别最容易翻车的五个场景虽然压缩包解压就能跑但换到你自己电脑上大概率会在下面几个地方卡住。按我踩过坑的频率从高到低排个序。4.1 环境依赖timm与torchvision版本冲突现象运行train.py直接报ModuleNotFoundError: No module named timm或者AttributeError: VisionTransformer object has no attribute forward_features这种诡异错误。原因timm这个库更新很频繁不同版本的API有改动另外机器上可能同时存在多个Python环境pip装到了另一个环境里。解决先确认当前环境是哪个基础环境然后安装固定版本pip install timm0.9.12 python -c import torch, torchvision; print(torch.__version__, torchvision.__version__)只要torch是2.x版本timm 0.9.12配合torchvision 0.15以上基本不会有接口问题。如果还报错把命令行里报错堆栈最后两行贴到搜索引擎看这招比乱猜版本快得多。4.2 数据为空ImageFolder找不到图片现象RuntimeError: Found 0 files in ...或者训练一开始loss就是NaN且accuracy为0。原因数据集目录层级放错了。ImageFolder只认“根目录/类别子目录/图片.jpg”这种两层结构如果把图片直接平铺在根目录、或者套了三层文件夹它都找不到样本。还有可能是路径里有中文Windows下中文路径经常引发编码问题。解决重新整理目录严格按第2.2节的结构摆放。路径中不要出现中文字符包括用户名目录和数据集文件夹名图片后缀保持jpg、jpeg、png大小写都行但别混用一些奇怪格式。另外Windows下如果多进程加载卡死把DataLoader的num_workers改成0再试确认能跑通再加到2。4.3 loss不降小数据集从零训练很痛苦现象train_loss前10个epoch几乎不动或者一直卡在0.69附近val_acc和瞎猜差不多。原因自实现ViT从头训练需要大量数据花卉数据集每类只有几百张时随机初始化的Transformer很难收敛另一个常见原因是学习率不对ViT用惯了CNN那套lr1e-3容易震荡而lr1e-4又可能学不动。解决先用一个“过拟合单batch”的测试来判断模型链路通不通把batch_size改成4固定一个batch迭代50步看loss能不能降到接近0。能降说明反向传播链路没问题问题出在数据量或超参降不动说明模型结构、梯度传播有bug。提示小数据集上最有效的解法不是疯狂调参而是换成预训练权重做finetune。数据量有限时让模型在ImageNet上学到的低级特征迁移到花卉上效果比从零开始训练高一个档次。4.4 撞显存ViT-B/16的batch设置现象CUDA out of memory训练程序中断显存小的卡直接崩掉。原因ViT-B/16在224x224输入下非常吃显存每张图约等于196个token做12层全局attentionbatch_size32在8G显存的卡上基本跑不动。解决把config里的batch_size降到8或16同时把num_workers降到0或2。还不行就换更小的输入尺寸比如image_size160对应patch数量从196降到100显存压力小一半精度损失在可接受范围内。最后的手段是用梯度累积每4个小batch后再更新一次梯度代码改动多在train循环里。# 梯度累积示例accum_steps4 optimizer.zero_grad() loss.backward() if (step 1) % accum_steps 0: optimizer.step()注意这个写法里optimizer.zero_grad()的位置要放在累积循环外部否则每步都清零累积就失效了正确做法是最外层zero_grad每隔accum_steps步step一次。4.5 推理错乱类别索引映射不一致现象训练时val_acc到了90%以上用DoClassify.py对单张图片分类结果完全不对而且每次都稳定地错成同一个类别。原因这是典型的classes映射错位。ImageFolder按字母顺序编号是daisy0, dandelion1, rose2, sunflower3, tulip4但训练保存的是模型权重推理脚本如果自己另写了一套编号或者加载了别人项目里顺序不同的classes列表预测的argmax索引对不上真实验标签。解决训练结束时把类别列表存成json推理时直接从json读。train.py里已经有一行json.dump(classes, f)DoClassify.py开头读取它不要自己手工维护类别映射。这是最容易忽略却最致命的一个坑强烈建议复盘一遍。5. 用DoClassify.py做单图推理验证Top-K输出与类别映射自检训练出来的模型最终要用在单张图片上DoClassify.py做的就是这件事加载权重、读图、预处理、前向推理、输出概率最高的前几个类别。下面是完整单图推理脚本的核心逻辑# Modules/DoClassify.py import json import torch from PIL import Image from torchvision import transforms from config import Config from model import ViTClassifier def classify(image_path, checkpointcheckpoints/best_model.pth): # 读取训练时保存的类别列表杜绝映射错位 with open(classes.json, r, encodingutf-8) as f: classes json.load(f) model ViTClassifier(num_classesConfig.num_classes) model.load_state_dict(torch.load(checkpoint, map_locationcpu)) model.eval() transform transforms.Compose([ transforms.Resize((Config.image_size, Config.image_size)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) img Image.open(image_path).convert(RGB) x transform(img).unsqueeze(0) with torch.no_grad(): logits model(x) probs torch.softmax(logits, dim1).squeeze(0) topk torch.topk(probs, k3) for score, idx in zip(topk.values, topk.indices): print(f{classes[idx]:15}: {score.item():.4f})推理时的预处理必须和验证集完全一致同样224尺寸、同样的mean和std归一化否则模型看到的输入分布和训练时不一样准确率会掉一截。map_locationcpu是保险写法即使训练在GPU上完成没有显卡的机器也能加载权重。验证模型是否真学会了我常用两个手段。一是把验证集完整过一遍统计每个类别的单独acc和混淆矩阵哪个类容易混就去看它和哪个类形态接近二是随手挑几张训练集之外的网络图片用DoClassify.py跑一次不仅看Top-1答案还要看Top-3的置信度分布——如果Top-1和Top-2都徘徊在30%左右说明模型对这个样本心里没底。从我自己的血泪经验看这类项目最容易翻车的从来不是网络结构而是刚开始说的数据路径和classes映射。从那以后我每次训练结束都要打开classes.json核对一遍推理前先打印类别表确认和目标数据集文件夹顺序完全一致才敢把结果写进报告。希望帮到你这份ViT花卉识别代码本身很完整按上面这条链路走一遍你能跑通、能讲清、还能顺手改成自己的数据集。本文还有配套的精品资源点击获取