简介:这是一份基于Python与EfficientNet的植物叶片病害图像识别完整项目,面向计算机相关专业学生完成毕业设计、课程设计或课题初期演示,也适合有一定基础的开发者学习迁移。包内共82个文件,涵盖7个py源码(训练、预测、工具模块)、60个jpg叶片样本图像、标签文件、类别索引文件,以及4个xml配置和工程结构文件,整体约1MB,目录按训练、预测、模型等模块划分,便于快速定位与二次开发。项目从数据准备、模型训练到图像预测构成完整链路,并配有详细文档与说明,可帮助读者理解EfficientNet在图像分类任务中的实际应用,掌握数据集组织、模型调用、结果预测等关键环节。目前已有136人学习,适合作为课堂项目、课程设计或答辩演示的基础工程,也可在此基础上扩展其他植物种类或病害类型,快速完成功能迭代。
1. 为什么植物叶片病害图像识别选了 EfficientNet 而不是 ResNet
刚接触植物叶片病害图像识别,第一反应通常是拿 ResNet 或 VGG 直接跑。但真正做完一轮,问题不在准确率,而在各类别样本不均衡、叶片正反面光照差异大、训练耗时高。这份基于 python + EfficientNet 的完整资源里,train.py 负责训练闭环,predict.py 负责单张与目录批量推理,model.py 把分类头替换成实际病害类别数,plot_img 存放训练过程可视化图片,class_indices.json 和 label.txt 做索引到可读名的映射。EfficientNet 在 ImageNet 上精度接近 ResNet50,参数量却只有后者的三分之一左右,配合迁移学习后,课程设计和毕业设计场景很容易在短时间内跑出可展示的结果。
2. EfficientNet 复合缩放原理与叶片数据预处理管线
2.1 复合缩放到底在缩什么
EfficientNet 的基线是 EfficientNet-B0,它的核心思想不是单纯把网络加深或加宽,而是同时缩放宽度、深度和输入分辨率。每个维度用一个复合系数 φ 控制,典型关系是 depth = α^φ,width = β^φ,resolution = γ^φ,且约束 α·β²·γ² ≈ 2。这句话读起来偏理论,但落地的含义非常实际:参数和计算量增长可控,精度提升却比单独加层或加宽度更明显。
对叶片病害这种细粒度识别任务,叶斑、霉层、枯萎区域的纹理差异往往集中在几十个像素里。深网络负责捕获语义特征,宽网络负责保留细节特征,复合缩放同时放大两者,比只加深 ResNet 更能保留叶片边缘和病斑纹路的区分信息。这也是为什么同一份数据用 EfficientNet-B0 就能达到 ResNet50 的效果,而推理速度更快,显存压力更小。
2.2 数据组织、label.txt 与 class_indices.json 的分工
压缩包里看到的 EfficientNet-main 目录并不是数据全集,而是把模型源码、训练脚本、推理脚本和可视化输出分开组织。标准做法是准备一个带类别子目录的数据根目录,torchvision 的 ImageFolder 会自动按子目录名生成标签:
data/ ├── train/ │ ├── Tomato___Early_blight/ │ │ ├── 001.jpg │ │ └── ... │ └── Potato___Late_blight/ └── val/ ├── Tomato___Early_blight/ └── ...label.txt 保存的是人可读的类别全名,class_indices.json 保存的是训练时生成的索引映射,推理时 predict.py 不会去读训练目录名,而是直接加载这个 JSON 文件,所以两者必须保持同步。class_indices.json 的内容结构通常是:
{ "0": "Tomato___Early_blight", "1": "Tomato___Late_blight", "2": "Tomato___healthy", "3": "Potato___Early_blight", "4": "Potato___Late_blight" }注意这里的 key 是字符串而不是整数,因为 JSON 的 key 只能是字符串。predict.py 解析出来之后需要按字符串索引,或者先统一转成 int,否则推理时class_indices[idx.item()]很容易因为类型不匹配报 KeyError。
2.3 数据增强:叶片识别最容易被忽略的一块
叶片识别的难点在于同一个病害在不同品种、不同光照条件下颜色差异很大。训练阶段的预处理建议分成两条管线:训练集做随机裁剪、水平翻转、旋转和颜色抖动,验证集只做 Resize 和归一化,不掺入任何随机增强。常见配置如下:
# utils.py 中常见的数据增强配置 from torchvision import transforms train_transforms = transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(p=0.5), transforms.RandomRotation(15), transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) val_transforms = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ])Resize 到 256 再 RandomCrop 到 224,是为了在裁剪过程中引入位置扰动,比直接 Resize 到 224 更能抑制过拟合。ColorJitter 的三个参数分别控制亮度、对比度和饱和度,叶片任务里饱和度扰动要克制,超过 0.2 会把黄化和缺素这两种不同病害混在一起。归一化 mean 和 std 用 ImageNet 的通用值,只要加载的是 ImageNet 预训练权重就必须沿用这组数值,换成其他值会导致第一层输入分布不匹配,迁移学习效果明显下降。
2.4 model.py 里的模型构建逻辑
model 目录下的 model.py 主要负责加载 EfficientNet 预训练模型,并把最后全连接层替换成实际病害类别数。常见的写法是用 timm 库,它对 EfficientNet 的 B0 到 B7 支持更完整:
# model/model.py import timm import torch.nn as nn def build_efficientnet(num_classes: int, model_name: str = 'efficientnet_b0'): model = timm.create_model(model_name, pretrained=True) in_features = model.classifier.in_features model.classifier = nn.Linear(in_features, num_classes) return model关于这层替换,有两个容易忽略的细节。第一,in_features 必须从原模型的 classifier.in_features 读取,不能写死成 1280,因为 B0 的嵌入维度是 1280,B3 是 1536,写死会在换模型时直接崩。第二,model/init.py 只是 Python 包标记文件,不需要修改,但所有涉及模型导入的地方都要保证 model 目录在 Python 路径里。
3. train.py 训练脚本拆解:从数据加载到权重保存
3.1 用 ImageFolder 加载数据并生成索引
train.py 的第一步是把文件夹结构变成可训练的 Dataset。直接用 torchvision.datasets.ImageFolder 即可,它会自动扫描子目录并生成 class_to_idx 映射。数据加载部分通常是这样写的:
# train.py 中加载部分 import json from torchvision import datasets from torch.utils.data import DataLoader train_dataset = datasets.ImageFolder('data/train', transform=train_transforms) val_dataset = datasets.ImageFolder('data/val', transform=val_transforms) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4) val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4) # 把 {类名: 索引} 反转为 {索引: 类名},推理时直接读取 class_to_idx = train_dataset.class_to_idx class_indices = {str(v): k for k, v in class_to_idx.items()} with open('class_indices.json', 'w', encoding='utf-8') as f: json.dump(class_indices, f, indent=2)这里的关键是反转逻辑。ImageFolder 的 class_to_idx 是“类名在前,索引在后”,但推理时拿到的是模型输出的整数索引,需要快速找到类名,所以保存时要倒过来,存成"0": "类名"的结构。json 文件里的 key 不能是整数,dump 之前必须 str() 转一次。许多人在推理时报 KeyError,问题就出在这一步。
num_workers 在 Windows 下建议设为 0 或 2,Linux 下可以设 4 或 8。如果你是在 Linux 服务器上用虚拟环境跑,先确认 Python 和 torch 版本和保存权重时一致,否则加载 checkpoint 会抛出版本兼容性相关告警。
3.2 优化器、损失函数与学习率策略
病害识别任务里决定收敛速度的参数有三个:优化器、初始学习率、学习率衰减方式。迁移学习场景下 AdamW 配 1e-4 起步是最稳妥的组合,SGD 动量法需要更长的热身期,对一两百张样本的小数据集并不友好。
# train.py 中优化器与调度器 import torch import torch.nn as nn import torch.optim as optim model = build_efficientnet(num_classes=len(class_to_idx)) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = model.to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4) scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=30, eta_min=1e-6)CrossEntropyLoss 内部已包含 softmax,所以模型最后一层不需要额外加 log_softmax。T_max=30 表示余弦周期覆盖 30 个 epoch,学习率会从 1e-4 平滑降到 1e-6。如果换成阶梯式衰减,可以用 StepLR(step_size=10, gamma=0.1),这类参数在课程设计里调整空间直观,答辩时也容易讲清楚为什么这么设。
3.3 训练循环:零梯度必须放在反向传播之前
训练循环是整份脚本里出错率最高的地方,顺序错一个位置,结果就完全不一样。标准顺序是零梯度、前向、算损失、反向传播、更新权重:
# train.py 核心训练逻辑 def train_one_epoch(model, loader, optimizer, criterion, device): model.train() total_loss = 0.0 correct = 0 total = 0 for images, labels in loader: images = images.to(device) labels = labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() * images.size(0) preds = outputs.argmax(dim=1) correct += (preds == labels).sum().item() total += labels.size(0) return total_loss / total, correct / totaloptimizer.zero_grad() 放在 loss.backward() 前,是为了清掉上一个 batch 累积的梯度。如果不做这一步,梯度会跨 batch 累加,相当于变相放大学习率,训练 loss 会出现锯齿状震荡。outputs.argmax(dim=1) 取每个样本预测概率最高的下标,再和 labels 逐位比较,得到这个 epoch 的训练准确率。
3.4 模型保存策略与 plot_img 的可视化
每个 epoch 结束后要在验证集上算准确率,只有验证准确率超过历史最佳时才保存权重。这样即使后面过拟合,留下的仍是最能泛化的版本,而不是最后一个 epoch 的模型。权重保存通常同时写模型参数和优化器状态:
torch.save({ 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), }, 'model_best.pt')load_state_dict 时只需要 state_dict 字段,optimizer 字段留给断点续训。训练中断后可以用它恢复学习率位置,而不是从头再来。plot_img 目录里的那些 3.JPG、5.JPG、8.JPG 就是在训练结束后生成的,常见做法是把 train_loss、val_loss、train_acc、val_acc 四条曲线画在一张图里保存,避免训练日志刷屏后找不到对应关系。
| 文件/目录 | 训练阶段作用 | 推理阶段作用 |
|---|---|---|
| train.py | 训练入口,负责数据加载和训练循环 | 不参与推理 |
| model/model.py | 构建 EfficientNet 并替换分类头 | 推理时需重建相同结构 |
| class_indices.json | 由训练数据生成 | 预测结果映射为类名 |
| label.txt | 记录类别可读名 | 打印到控制台 |
| plot_img/ | 保存 loss/acc 曲线图 | 答辩展示训练过程 |
4. predict.py 推理链路:class_indices 映射与 inference 排错
4.1 加载 checkpoint 前先重建模型结构
推理的第一步是重建和训练时完全一致的模型结构,再加载权重。常见错误是把 checkpoint 直接 load 进一个没改分类头的预训练模型,导致 size mismatch。规范的推理初始化如下:
# predict.py 中初始化与加载 import json import torch with open('class_indices.json', 'r', encoding='utf-8') as f: class_indices = json.load(f) num_classes = len(class_indices) model = build_efficientnet(num_classes=num_classes) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') checkpoint = torch.load('model_best.pt', map_location=device) model.load_state_dict(checkpoint['state_dict']) model.to(device) model.eval()map_location=device 是便携性的关键。在 GPU 上训练的权重拿到只有 CPU 的机器上推理,不加 map_location 会直接报 CUDA 不可用。load_state_dict 要求层名和维度完全一致,推理时构造的类别数和训练时不一致会报 size mismatch,这时优先检查 class_indices.json 是不是训练时同一次生成的文件。
4.2 单张叶片图片预测与置信度输出
推理阶段不需要计算梯度,也不能加随机增强。PIL 打开图片后要确认转成 RGB,灰度图通道数不对会直接报错。单张预测函数可以这样写:
# predict.py 单张预测函数 from PIL import Image def predict_single(img_path, model, class_indices, transform): img = Image.open(img_path).convert('RGB') x = transform(img).unsqueeze(0).to(device) with torch.no_grad(): logits = model(x) probs = torch.softmax(logits, dim=1) confidence, idx = probs.max(dim=1) class_name = class_indices[str(idx.item())] return class_name, confidence.item()最后全连接输出的 logits 不能直接当置信度,必须经过 softmax 才有概率含义。confidence 反映的是模型对当前预测的内部置信度,不完全等价于真实概率,但在工程上足够用来排序和过滤低置信结果。注意单张图片要加 unsqueeze(0) 补齐 batch 维,否则模型会报期望 4D 输入。
4.3 批量推理:遍历 test_img 输出结果文件
predict.py 里出现 test_img 目录,通常就是用来做整目录推理的。遍历所有图片,把文件名、预测类别、置信度写进一个文本文档,比一张张看的效率高很多:
# predict.py 批量推理 from pathlib import Path def predict_directory(img_dir, model, class_indices, transform): results = [] for img_path in sorted(Path(img_dir).glob('*.jpg')): name, conf = predict_single(str(img_path), model, class_indices, transform) results.append(f'{img_path.name}\t{name}\t{conf:.4f}') report = '\n'.join(results) print(report) with open('predict_result.txt', 'w', encoding='utf-8') as f: f.write(report) return reportsorted 排序是为了保证输出顺序稳定,glob 只匹配 jpg。如果你的测试图混入 png 或 bmp,可以改成glob('*')后用 suffix 过滤,但不要只过滤后缀却忘了检查图片格式是否和训练数据一致。批量推理时如果出现某张图损坏导致中断,推荐在循环里加 try-except,跳过坏图并记录文件名,而不是让整个目录推理停住。
4.4 推理阶段高频报错对照
predict.py 的报错绝大多数是类型和维度问题,不涉及模型本身。
| 报错现象 | 原因 | 处理方法 |
|---|---|---|
| size mismatch for classifier.weight | 训练类别数与推理类别数不一致 | 核对 class_indices.json 长度 |
| Missing key(s) in state_dict | 模型结构不同或加载了别的网络权重 | 检查 build_efficientnet 参数 |
| Expected 3D input, got 4D | 单张图缺少 batch 维 | 在 transform 后加 unsqueeze(0) |
| KeyError: '0' | class_indices 的 key 类型不是字符串 | json 读取后统一转 str |
| CUDA out of memory | 推理输入 batch 过大 | 降低 batch 或换 CPU 推理 |
5. 迁移学习微调技巧与样本均衡处理
5.1 分层冻结:先训分类头,再解冻深层
直接用 1e-4 全量微调不是最优方案。叶片背景相对简单,浅层卷积提取的是通用边缘和颜色特征,不需要重新学。常见做法是加载预训练权重后,把浅层 block 冻结,只训练最后分类头和最后两个 stage,先把分类头收敛,再解冻全部层用 1e-5 微调。冻结操作通过 requires_grad 控制:
# 冻结前几个 block,只训练后半部分 cnt = 0 for name, param in model.named_parameters(): if 'block' in name and cnt < 6: param.requires_grad = False cnt += 1这个两阶段策略在叶片数据集上通常能带来 2 到 5 个百分点的提升,尤其是病斑边界和健康组织颜色接近的类别,浅层特征保留原样反而更稳定。
5.2 类别不均衡:用采样器而不是改损失函数
病害数据里健康叶片总是比重病叶片多,模型容易把所有样本都推给数量大的类别。与其在 loss 上加 class weight,我更推荐 WeightedRandomSampler,它按类别样本数的反比调整采样概率:
from collections import Counter from torch.utils.data import WeightedRandomSampler class_counts = Counter(train_dataset.targets) weights = [1.0 / class_counts[t] for t in train_dataset.targets] sampler = WeightedRandomSampler(weights, num_samples=len(weights), replacement=True) train_loader = DataLoader(train_dataset, batch_size=32, sampler=sampler)这里的 weights 长度等于样本总数,每个样本被采到的概率与其所属类别的样本数成反比。少样本类别不要只靠过采样硬撑,还应该在增强里对病斑区域做更小幅度的旋转和裁剪,比简单复制原图更有效。
5.3 用混淆矩阵和 CAM 热力图验证模型真正学到了什么
准确率到 90% 不代表模型关注的是病斑。真正可靠的验证是看混淆矩阵中哪两个类别最容易互相混,再配合梯度加权类激活映射确认模型关注区域是不是叶片中心或病斑边缘。plot_img 目录下可以单独建一个 visualization 子目录,保存测试集的 top1 概率、top5 概率和 CAM 热力图。
最后把每个类别的准确率、推理耗时、参数量三个指标用文本方式输出,而不是只盯着整体 loss 曲线。这样在答辩或交付时,你能直接指出哪个类别是误差的主要来源,而不是只说整体准确率到了某个值。
本文还有配套的精品资源,点击获取