简介一套基于Transformer模型的木薯叶病虫害分类Python源码项目适合作为期末大作业、课程设计或深度学习入门实践参考。源码已在本地编译通过难度适中内容经助教审定可放心参考运行。压缩包共12个文件包含6个Python脚本、5个编译生成的pyc文件及1个说明文档整体仅11KB其中py文件按职责拆分为数据读取、Transformer模型构建、全局参数配置、GPU管理与训练入口等功能模块结构清晰便于逐模块理解与二次开发。已有199人浏览学习适合需要快速搭建图像分类基线或熟悉Vision Transformer实战流程的读者。除核心实现外还提供Markdown说明文档可辅助快速了解项目结构配置与复现要点是一份精炼且完整的高分项目范例。1. python 木薯叶病虫害分类源码一个能跑的 transformer 工程长什么样先说结论再拆包这个题目对应的不是玩具 demo而是一套从数据读取到模型训练、再到推理验证的完整 python 工程核心模型用 transformer 架构手写实现不是调库一行带过的那种。木薯叶病虫害分类任务的目标是区分 5 类叶片状态包括三种病毒病、一种螨害和健康叶片这类任务在农业视觉里很典型拿到课程设计或期末大作业里做题目比通用的猫狗分类要扎实得多。我拿到压缩包第一反应是看它是不是那种代码能跑、但换个环境就废的黑匣子。把 zip 解开之后六个 py 文件加一个 README结构很干净数据集类、模型类、全局参数、GPU 检查、训练入口和推理入口各司其职。对想复现 transformer 图像分类流程的人来说这份源码的价值在于把 ViT 的核心步骤原样摊开了适合当骨架改自己的任务也适合直接跑通后对着代码讲原理。2. 把压缩包拆开看文件职责、入口顺序与全局参数拿到源码包第一步不是跑 run.py而是先把这个工程的文件结构读明白。一个 transformer 分类项目只要超过 300 行代码就该拆文件而这个包正是这么拆的。先看文件清单再谈每个文件的边界和调用方式。2.1 六份 py 文件对应哪个环节这个工程的拆法很直白数据归数据模型归模型入口归入口。文件职责关键内容Global_Variable.py全局参数集中地数据路径、batch_size、epochs、学习率、图片尺寸、类别数Gpu.py设备判定优先 cuda不可用时回退 cpuCassavaDataset.py数据加载Dataset 子类负责读图、预处理、返回样本和标签Model.py模型定义transformer 分类网络结构含 patch 化与注意力层run.py训练入口数据装载、损失函数、优化器、训练循环main.py推理/验证入口加载权重对单张或一批图片做预测需要注意 run.py 和 main.py 是分开的这是这个工程做得比较讲究的地方。训练循环里带着梯度、优化器、学习率调度和强数据增强而推理只需要模型权重和预处理两者混在一个脚本里会让代码又长又容易误操作。拆开之后训练阶段跑 run.py验证或预测的时候走 main.py逻辑清楚答辩的时候也好讲。2.2 全局参数集中改Global_Variable.py 里值得关注的配置我习惯拿到工程先看全局参数文件因为所有能改的旋钮基本都在这一个文件里。常见做法是像下面这样集中定义而不是把参数散落到各个函数里DATA_ROOT ./data # 数据集根目录 TRAIN_DIR os.path.join(DATA_ROOT, train) VAL_DIR os.path.join(DATA_ROOT, val) NUM_CLASSES 5 # 五类CMD / CBSD / CGM / MCD / healthy IMG_SIZE 224 # 输入图片统一 resize 到 224x224 PATCH_SIZE 16 # patch 大小224/1614共 196 个 patch BATCH_SIZE 32 # 显存不够时最先降这个 EPOCHS 120 LR 1e-4 # transformer 建议比 CNN 更小的学习率 WEIGHT_DECAY 1e-4 WARMUP_EPOCHS 5 # 前 5 个 epoch 学习率线性上升 SEED 42这里几个参数值得单独说明。PATCH_SIZE 和 IMG_SIZE 的关系是硬约束图片边长必须是 patch 边长的整数倍。224 除以 16 得到 1414 乘 14 是 196 个 patch token再加一个 class token序列长度就是 197。如果把 IMG_SIZE 改成 128 之类位置编码的维度就要跟着重新算。真实复现时最容易翻车的点就在这里后面避坑章节还会再提。学习率我在跑 transformer 图像分类时一般不会开到 1e-3这是和 ResNet 这类 CNN 最大的习惯差异。ViT 的优化对学习率非常敏感1e-3 经常出现 loss 震荡甚至不降1e-4 到 5e-4 是相对稳的区间。这个工程把 LR 放到全局变量里改起来很方便训练跑挂的时候先回来调它。2.3 run.py 与 main.py训练入口和推理入口的分工run.py 里做的是一整套标准训练流程构建 dataset 和 dataloader实例化模型定义交叉熵损失和 AdamW 优化器然后按 epoch 循环。核心训练循环的大致逻辑是这样的# run.py 训练主循环的核心结构 model CassavaViT(num_classesNUM_CLASSES).to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lrLR, weight_decayWEIGHT_DECAY) # warmup cosine 衰减的调度器 scheduler get_cosine_schedule_with_warmup(optimizer, num_warmup_stepsWARMUP_EPOCHS * len(train_loader), num_training_stepsEPOCHS * len(train_loader)) for epoch in range(EPOCHS): model.train() for imgs, labels in train_loader: imgs, labels imgs.to(device), labels.to(device) pred model(imgs) loss criterion(pred, labels) optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step()训练循环本身并不复杂值得关注的是模型输出直接和交叉熵做比较这说明 Model.py 返回的是 5 类的 logits没有在内部额外套 softmax这是 PyTorch 的标准做法——CrossEntropyLoss 内部已经包含 softmax 过程模型输出 logits 即可。如果自己改代码时在模型里先 softmax 再进 loss数值上虽然不会报错但梯度路径绕了一圈没必要。main.py 是推理链路它只做三件事加载训练好的权重文件对输入图片做和训练时一致的预处理然后取 argmax 作为预测结果。典型写法是这样的# main.py 加载模型并推理 model.load_state_dict(torch.load(best_model.pt, map_locationdevice)) model.eval() with torch.no_grad(): logits model(img_tensor.unsqueeze(0)) pred logits.argmax(dim1).item()推理阶段必须用 torch.no_grad() 包住前向过程否则会为不需要的梯度分配显存batch 量大的时候白白吃掉几百 MB 显存。另外 model.eval() 不能省虽然 ViT 里没有 BN 层但 dropout 在 eval 模式下才会被正确关闭。3. CassavaDataset 与预处理从图片文件到能喂进 patch embedding 的 Tensor数据加载是整个工程里最容易被低估的部分。transformer 模型本身不吃图片它吃的是 patch 序列而 patch 是从像素张量切出来的。所以数据管线要做的事情是把磁盘上的 jpg 文件变成归一化后的 3×224×224 张量顺序、尺寸、数值范围都不能出错。3.1 Dataset 骨架三件套init/len/getitemCassavaDataset.py 实现的是 PyTorch Dataset 接口这个类在三处方法里把整个数据读取逻辑封装好# CassavaDataset.py 核心结构 class CassavaDataset(Dataset): def __init__(self, img_paths, labels, transformNone): self.img_paths img_paths self.labels labels self.transform transform def __len__(self): return len(self.img_paths) def __getitem__(self, idx): img Image.open(self.img_paths[idx]).convert(RGB) label self.labels[idx] if self.transform is not None: img self.transform(img) return img, labelgetitem里有两个细节值得注意。第一是 convert(RGB)有些叶片图片是灰度或带 alpha 通道的不统一转换会导致后面通道数对不上报错。第二是返回的是图像和标签的配对transform 只作用在图像上标签不做任何变换这一点在数据增强时很关键——随机翻转和裁剪只能动图不能动类别。3.2 标签组织文件夹即类别与 CSV 映射两种方式标签怎么和图片对应是这类源码包里最常见的分叉点。一种方式是文件夹即类别train 目录下每个子文件夹名就是类别名ImageFolder 可以直接读另一种方式是 CSV 映射图片路径和标签 id 分开存放Dataset 构造时传入两个列表。这个工程的数据集类是自定义实现的所以更贴合第二种方式把所有样本的路径读进一个列表标签读进另一个列表在getitem里按索引取。这种方式的好处是划分 train/val 时只需要切分两个列表不用复制图片文件也不受目录结构的限制。实际切分时我一般会固定随机种子避免每次运行生成不同的划分导致前后结果没法对比# 按比例切分并固定随机种子保证可复现 random.seed(Global_Variable.SEED) indices list(range(len(all_paths))) random.shuffle(indices) val_len int(len(indices) * 0.2) val_idx, train_idx indices[:val_len], indices[val_len:]3.3 训练/验证划分与增强边界木薯叶任务常用的公开数据集大致是 2 万张量级的图片、5 个类别每类样本数并不均衡某些病害类别天然比健康叶片少。这种不均衡下乱切 train/val 很容易让某个类别在验证集里只有几十张指标波动很大。我的习惯是按标签分层划分也就是先按类别分组再在每个类别内部按比例抽验证集保证每个类别在训练集和验证集里的占比接近。这一步用到 sklearn 的 StratifiedShuffleSplit 就能实现写起来很短但对后续评估的稳定性帮助很大。增强策略上训练集可以放开手脚随机裁剪、随机水平翻转、颜色抖动都是常见选项。验证集和测试集只能做 resize 和归一化任何随机性都会让评估结果忽高忽低。归一化参数几乎固定用 ImageNet 统计量 mean[0.485, 0.456, 0.406]、std[0.229, 0.224, 0.225]因为预训练权重都是在这个分布下学出来的自己算数据集均值反而容易把数值范围带偏。4. Transformer 分类模型的代码逻辑patch 化、自注意力与组装这一章是核心。模型部分从零手写了一个简化版 ViT输入图片先被切成 patch每个 patch 线性投影成一个 token加上位置编码后进入多层 transformer encoder最后用 class token 接分类头输出 5 类 logits。整个链路在 Model.py 里组装。4.1 PatchEmbed把 224×224 图像切成 196 个 tokentransformer 本身是序列模型它不知道图像是什么。把图像变成序列的标准做法是用一个步长等于卷积核大小的卷积把图片一次性切成不重叠的 patch并完成线性投影# PatchEmbed图像 - patch 序列 - token 序列 class PatchEmbed(nn.Module): def __init__(self, in_ch3, img_size224, patch_size16, embed_dim768): super().__init__() n_patch (img_size // patch_size) ** 2 # 14x14 196 self.patch_size patch_size # 卷积核和步长都等于 patch_size实现不重叠切块 self.proj nn.Conv2d(in_ch, embed_dim, kernel_sizepatch_size, stridepatch_size) # class token 和位置编码1 是给 class token 留的位置 self.cls_token nn.Parameter(torch.zeros(1, 1, embed_dim)) self.pos_embed nn.Parameter(torch.zeros(1, n_patch 1, embed_dim)) def forward(self, x): B x.shape[0] # batch size x self.proj(x) # (B, 768, 14, 14) x x.flatten(2).transpose(1, 2) # (B, 196, 768) cls self.cls_token.expand(B, -1, -1) # (B, 1, 768) x torch.cat([cls, x], dim1) # (B, 197, 768) x x self.pos_embed # 加上位置编码 return x这个模块是整个模型第一个关键点。proj 是参数化的训练过程中会学到把 patch 内容映射为 768 维向量的权重所以不是简单地把像素拉平而是投影。flatten(2) 把 H 和 W 两个维度合并成序列维度transpose 调整维序让序列维在前。cls_token 是 ViT 的经典设计它不来自任何 patch最终分类时取它在最后一层的输出作为整张图的特征相当于一个可学习的全局归纳向量。4.2 多头自注意力qkv 的拆分、缩放与 softmaxTransformer 之所以能替代卷积做图像分类核心就在多头自注意力。对序列里的每个 token它都能直接和序列里所有其他 token 计算关联不受卷积核大小的感受野限制。代码上实现注意力层是这个工程最值得逐行读的部分# Attention多头自注意力的核心计算 class Attention(nn.Module): def __init__(self, dim768, num_heads12): super().__init__() self.num_heads num_heads head_dim dim // num_heads # 每个头 64 维 self.scale head_dim ** -0.5 # sqrt(64) 的倒数缩放因子 self.qkv nn.Linear(dim, dim * 3) # 一次线性变换同时产生 qkv 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, C // self.num_heads) qkv qkv.permute(2, 0, 3, 1, 4) # (3, B, 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) out (attn v).transpose(1, 2).reshape(B, N, C) return self.proj(out)scale 这个细节很多人会忽略但它直接影响训练稳定性。如果不除以 sqrt(head_dim)点积结果的方差会随着 head_dim 增大而变大softmax 之后分布过于尖锐梯度容易消失。除以缩放因子是原始 transformer 论文的标准做法。代码里用三倍宽度的 Linear 一次算出 q、k、v比拆成三个 Linear 更省参数和计算量reshape 和 permute 负责把维度拆到多头让每个头在独立的子空间里计算注意力。4.3 Encoder 堆叠与分类头Model.py 里的最终输出有了 PatchEmbed 和 Attention 之后把两者组装成完整的 transformer encoder 是标准的残差 LayerNorm结构每个 encoder 层内部包含一个注意力子层和一个前馈网络子层各自带残差连接。这个工程的 Model.py 外层组装大致如下# Model.py完整 ViT 分类模型组装简化示意 class CassavaViT(nn.Module): def __init__(self, num_classes5, depth12): super().__init__() self.patch_embed PatchEmbed() self.blocks nn.ModuleList([EncoderLayer() for _ in range(depth)]) self.norm nn.LayerNorm(embed_dim) self.head nn.Linear(embed_dim, num_classes) # 分类头输出 5 类 logits def forward(self, x): x self.patch_embed(x) for blk in self.blocks: x blk(x) x self.norm(x) cls x[:, 0] # 取 class token 的特征 return self.head(cls) # (B, 5)EncoderLayer 里每层都先用 LayerNorm 再做注意力这是 Pre-LN 结构和原始 transformer 的 Post-LN 有区别。Pre-LN 在训练深模型时更稳定梯度更容易回传这在这种从零训练的课程项目里很实用。分类头只接在 class token 上而不是对 197 个 token 输出做全局池化这是 ViT 和普通序列模型分类的关键差异——cls token 在训练中会被引导成聚合了整张图信息的向量。5. 训练与复现避坑清单5 个跑通后最容易翻车的位置源码本地编译过能跑不等于换到你的机器上也能跑。我复现过不少 transformer 图像分类工程这五个坑是出现频率最高的而且顺序基本固定先是路径再是学习率然后是显存最后是环境和数据划分。5.1 数据路径分隔符与数据集找不到现象运行 run.py 报 FileNotFoundError 或者 dataloader 里 image not found但打开目录看文件明明在。原因Windows 和 Linux 的路径分隔符不一样代码里如果硬编码了 / 或 \换系统必然出问题。另外有些同学把数据集解压后目录多套了一层比如 data/cassava/train/train路径和全局参数里写的不一致。解决统一用 os.path.join 拼接路径不要让路径字符串里出现硬编码分隔符。代码开头加一行断言最快暴露问题assert os.path.exists(TRAIN_DIR), f训练目录不存在: {TRAIN_DIR} print(训练集样本数:, len(os.listdir(TRAIN_DIR)))5.2 学习率设成 1e-3 直接不收敛现象训练 30 个 epochloss 在 1.6 附近反复横跳不下降或者更严重loss 直接变成 nan。原因transformer 对学习率比 CNN 敏感得多。ResNet 用 1e-3 甚至 1e-2 能跑但 ViT 从零训练时 1e-3 经常震荡。如果用了 AdamW 且 weight_decay 偏大参数更新幅度会被放大loss 更容易崩。解决把学习率降到 1e-4并把 warmup 打开。学习率从 0 线性上升到目标值的过程相当于让模型先用小步幅走稳再做大步长搜索。我当时排查的时候把 loss 曲线打成日志发现崩的点基本都在第一个 epoch 的末尾换成 1e-4 之后曲线从第 5 个 epoch 就开始稳定下降。5.3 显存炸掉与 patch 序列长度现象batch_size 32 在 GPU 上跑报 CUDA out of memory但同一份代码用 ResNet 跑同样的 batch 没有压力。原因自注意力的显存消耗是序列长度的平方量级。197 个 token 两两计算注意力每个头都要维护 197×197 的注意力矩阵12 个头叠起来再乘上 batch 32显存占用远超同尺寸 CNN。解决先降 batch_size 到 8 或 16 确认能跑通再逐步往上加。如果非要用大 batch可以降低输入分辨率比如把 IMG_SIZE 改成 160patch 数从 196 降到 100序列长度几乎减半显存压力显著下降。代价是 160 不能被 16 整除的话要同步调整 patch_size否则 token 数对不上位置编码。5.4 Python 3.7 环境与pycache旧字节码现象代码能跑但启动时出现一堆 FutureWarning 或 SyntaxWarning甚至提示 python 版本编译的字节码过期。原因压缩包里带着pycache目录里面的 .pyc 文件是 Python 3.7 编译的。如果你的环境是 3.8 以上解释器会重新编译原文件旧的 pyc 虽然不影响正确性但会产生警告。反过来如果环境低于 3.7某些语法可能直接没法解析。解决最干净的做法是把pycache整个目录删掉再运行让解释器用自己的版本重新生成字节码。Windows 命令行下删除后再执行 python run.py警告会消失。依赖方面建议用 requirements 里的版本号创建独立环境不要用全局环境硬跑。5.5 验证集指标虚高的数据泄漏现象验证集 accuracy 能到 0.97看起来非常漂亮但把模型拿到新图片上预测效果明显不如指标显示的那么好。原因大概率是 train 和 val 划分时混入了同一来源的相似图片。木薯叶数据集里同一棵植株的不同叶片照片往往是在同一时间、同一背景下拍摄的如果按文件列表顺序直接切分训练集和验证集会存在大量背景几乎一样的样本。解决划分时按类别分层抽样并且尽量让同源的图片进同一个集合。严格的做法是按图片 id 分组划分或者在切分前对图片计算感知哈希去重。课程项目里至少做到固定随机种子 分层划分并且每次提交前打印每个类别的样本数确认没有类目为空。6. 验证不只是 accuracy混淆矩阵与注意力可视化的进阶收尾模型跑通、指标达标之后这个工程还能再做两件让分数和研究价值都明显提升的事分类评估细化和注意力可视化。前者让答辩有数据支撑后者让transformer 在关注什么这个黑匣子问题变得可解释。6.1 混淆矩阵与每类召回率accuracy 在类别不均衡时是骗人的。木薯叶 5 类里健康叶片占比偏高模型把所有样本都预测成健康类也能拿到不错的 accuracy但对病害识别毫无意义。我一般会额外输出混淆矩阵和每类 precision、recall# 分类评估混淆矩阵与每类指标 from sklearn.metrics import confusion_matrix, classification_report import seaborn as sns labels, preds [], [] for imgs, lbls in val_loader: with torch.no_grad(): out model(imgs.to(device)).argmax(dim1).cpu().numpy() preds.extend(out) labels.extend(lbls.numpy()) print(classification_report(labels, preds, target_names[CMD, CBSD, CGM, MCD, healthy])) sns.heatmap(confusion_matrix(labels, preds), annotTrue, cmapBlues)confusion_matrix 能直接暴露模型最容易混淆哪两类。木薯叶任务里 CGM 和 MCD 的早期症状视觉上很像如果这两类互相误判严重说明模型学到的是颜色纹理特征而不是病斑结构特征这时可以尝试把输入分辨率从 224 提到 384让模型看到更多细粒度纹理。6.2 注意力热力图看模型到底在看叶片哪里ViT 的一个天然优势是注意力权重可以直接可视化。把最后一层 encoder 的注意力矩阵取出来对 class token 那一行做 softmax再 reshape 回 14×14叠加到原图上就能看到模型分类时重点关注的叶片区域# 提取最后一层注意力并可视化 attn model.blocks[-1].attn_weights[0] # 拿到注意力矩阵 cls_attn attn[0, :, 0, 1:].mean(dim0) # class token 对所有 patch 的平均注意力 heatmap cls_attn.reshape(14, 14).detach().cpu().numpy() heatmap cv2.resize(heatmap, (224, 224)) # 还原到原图尺寸如果热力图集中在叶片边缘和病斑区域说明模型学到了正确的判别依据如果注意力分散在背景土壤上就要检查预处理阶段是否裁剪得过松让背景干扰进了模型。6.3 调参与 Checkpoint 版本管理的习惯这个工程跑到后期我最大的教训是 checkpoint 命名必须带参数信息。一开始我习惯存成 best_model.pt跑了几次实验之后自己都分不清哪份权重是 lr 1e-4 跑出来的哪份是加了 mixup 的。从那以后我每次都把学习率、输入尺寸、epoch 写进文件名比如 vit_224_lr1e-4_ep120.pt。调参的时候坚持每次只动一个变量先定学习率再动数据增强最后动模型深度记录表格比靠记忆靠谱得多。这套折腾下来调参的节奏稳了结果也可复现了希望帮到你。本文还有配套的精品资源点击获取
