☰
基于PyTorch的树叶识别系统实战:数据、迁移学习与调参全指南
2026/10/1 5:35:05 网站建设 项目流程

简介:基于Python语言的树叶识别系统是针对计算机相关专业毕业设计与大作业场景的完整实战项目,适合正在完成课程设计或希望进行项目练手的学习者。项目将Python脚本与图形化界面相结合,围绕树叶识别全流程展开,源码结构清晰,可运行且便于二次开发。压缩包共4个文件,约6.98MB,包含2个Python核心脚本、1个ui界面文件以及1个系统演示视频。核心脚本承担特征提取与识别调度,ui文件构建操作界面,演示视频则直观展示运行效果。该成果经导师指导并通过评审(98分),源码已完成本地编译与严格调试,目前已有71人学习下载,适合需要快速获取可运行项目并深入理解树叶识别实现思路的同学。

1. 树叶识别系统:从“肉眼分辨”到“模型分类”,问题远比拍一张照片复杂

想象一个很常见的场景:你在公园里用手机拍下一片叶子,想识别它是桂花、银杏还是红枫。树叶识别系统的核心任务就是把“看见叶子”变成“知道是哪棵树”,这本质上是一个图像分类问题,不是目标检测,更不是语义分割。网上关于基于Python语言的树叶识别系统源码通常跑通不难,但很多人在答辩或者换一批真实照片测试时发现,模型在数据集上准确率很高,一放到真实场景就翻车,准确率直接从“优秀”掉到“及格线”以下。我见过太多课程设计项目卡在这里:模型用的是现成框架,数据是公开数据集,唯独没人认真想清楚“模型到底学到了什么特征”。这篇笔记写给正在做树叶识别系统、手头有Python基础但没系统做过图像分类项目的开发者,我会按数据整理、模型搭建、调参、踩坑、验证这条完整链路,把一套可复现的方案讲清楚。

2. 从“拍树叶”到“喂模型”:树叶图片怎么整理才能让识别系统真正可用

2.1 别急着训模型,先把类别、拍图规范和目录结构定下来

很多树叶识别系统在启动阶段就埋了雷:拿到图片后不管三七二十一直接开始训练,最后模型精度上不去,回头查才发现是数据本身乱成一锅粥。图像分类任务里,一张训练样本就是“一张图片 + 一个类别标签”,而最简单的标签组织方式就是按类别建文件夹。几乎所有深度学习框架的数据加载工具都认这种结构,PyTorch的torchvision.datasets.ImageFolder、TensorFlow的image_dataset_from_directory,全部默认读取“根目录/类别/图片.jpg”这种层级。

一个规范的目录结构应该长这样:

leaf_dataset/ ├── train/ │ ├── ginkgo/ │ │ ├── ginkgo_001.jpg │ │ ├── ginkgo_002.jpg │ │ └── ... │ ├── osmanthus/ │ │ └── ... │ └── maple/ │ └── ... ├── val/ │ ├── ginkgo/ │ ├── osmanthus/ │ └── maple/ └── test/ ├── ginkgo/ └── ...

类别的命名建议用英文小写加下划线,比如ginkgo、osmanthus,而不是中文名。因为中文名在部分环境里会引发编码问题,而且ImageFolder默认按文件夹名的字典序生成类别索引,中文排序不可控。类别的数量取决于项目需求,但至少每类要有30到50张图,图像分类模型对样本量非常敏感,类别再多、每类只有10张图,训练出来的模型基本不具备泛化能力。

拍图规范方面,常见做法是在白纸或纯色背景下单片拍摄,尽量让叶片占画面主体。这种规范能减少背景干扰,让早期训练更稳定。但只拍白底图有一个隐患:模型可能学到“白色背景=叶子”而不是“叶子形状=叶子”,所以后面在数据增强里必须加入背景扰动,这一点在第四章细说。

2.2 数据切分脚本:把训练集和验证集分开再开始动手

拿到源码后,第一件事不是急着装依赖,而是看它怎么切分数据。最不靠谱的写法是把所有图片喂进去训练,再用同一批图片算准确率,这种“自评分”没有任何参考意义。我一般习惯先写一个独立的数据切分脚本,在每个类别内部随机抽取15%到20%的图片作为验证集,剩余作为训练集,测试集再单独留一部分不参与任何调参。

下面是这个切分脚本的完整写法:

import os import random import shutil from collections import defaultdict def split_train_val(src_dir, dst_dir, val_ratio=0.2, seed=42): """把 src_dir 下每个类别的图片按比例移动为验证集。 src_dir 结构应满足: src_dir/类别名/图片文件 """ random.seed(seed) for cls_name in os.listdir(src_dir): cls_path = os.path.join(src_dir, cls_name) if not os.path.isdir(cls_path): continue images = [f for f in os.listdir(cls_path) if f.lower().endswith(('.jpg', '.jpeg', '.png'))] if len(images) == 0: continue random.shuffle(images) val_count = max(1, int(len(images) * val_ratio)) # 创建目标目录 train_dst = os.path.join(dst_dir, 'train', cls_name) val_dst = os.path.join(dst_dir, 'val', cls_name) os.makedirs(train_dst, exist_ok=True) os.makedirs(val_dst, exist_ok=True) for img_name in images[val_count:]: shutil.move( os.path.join(cls_path, img_name), os.path.join(train_dst, img_name) ) for img_name in images[:val_count]: shutil.move( os.path.join(cls_path, img_name), os.path.join(val_dst, img_name) ) print(f"{cls_name}: total={len(images)}, val={val_count}") if __name__ == '__main__': split_train_val( src_dir='leaf_dataset/raw', dst_dir='leaf_dataset', val_ratio=0.2 )

这段代码的逻辑不难:遍历源目录下的每个类别,把图片文件名收集起来,随机打乱后按比例分成两份,一份移动到dst_dir/train/类别名,一份移动到dst_dir/val/类别名。有一点需要特别注意,脚本里用的是shutil.move而不是copy,移动文件不会额外占用磁盘空间,适合图片总量大的场景。但这也意味着脚本只能执行一次,重复执行会把本来已经分好的文件再次移动,导致验证集比例失真,所以建议在脚本开头加一个判断,如果dst_dir/val已经存在就直接退出,而不是每次覆盖执行。

seed=42这个参数也很关键。固定随机种子能保证每次运行得到相同的切分结果,这样你调模型时看到的精度变化只来自模型本身,而不是数据切分抖动。如果某次实验发现验证集上表现特别差,你可以确定不是数据分得不均匀,而是模型或者增强参数出了问题。

2.3 图像规范化:统一尺寸、通道与像素范围,避免模型学的是背景

树叶图像来自不同拍摄设备,分辨率、宽高比、色彩空间都不一致。模型训练时通常需要固定输入尺寸,因此所有图片都要经过一个预处理管道。以PyTorch为例,图像进入模型前要做三件事:解码成RGB三通道、缩放到统一尺寸(常用224×224)、归一化到模型预期的像素范围。

下面这段代码可以帮你快速检查一批图片的质量,判断是否存在灰度图、白边或不可读文件:

import cv2 import numpy as np import os def inspect_images(root_dir, target_size=(224, 224)): problems = [] for cls_name in os.listdir(root_dir): cls_path = os.path.join(root_dir, cls_name) if not os.path.isdir(cls_path): continue for img_name in os.listdir(cls_path): img_path = os.path.join(cls_path, img_name) img = cv2.imread(img_path) if img is None: problems.append((img_path, '无法解码')) continue # 灰度图检测 if len(img.shape) == 2: problems.append((img_path, '灰度图')) # 尺寸过小检测 h, w = img.shape[:2] if min(h, w) < target_size[0]: problems.append((img_path, f'尺寸过小: {h}x{w}')) if problems: print(f"共发现 {len(problems)} 个问题文件:") for path, reason in problems[:20]: print(f" [{reason}] {path}") else: print("未发现问题,所有图片均可正常读取") inspect_images('leaf_dataset/raw')

运行这段脚本时,最常见的三个问题是:图片本身是灰度图、读取失败、分辨率远低于训练尺寸。灰度图可以作为检测目标之一,但一般做法是将其转换为三通道图,即把灰度值复制到RGB三个通道上,而不是直接丢弃,因为样本本来就少。分辨率太低的图放入训练集只会给模型增加噪声,干脆在切分之前就把它们筛掉。

还有一类问题脚本不好直接发现:背景大面积同色。很多公开树叶数据集用的是深色或纯黑背景,模型在训练时会把这些背景特征一并学进去,导致真实照片上识别率骤降。要发现这类问题,可以随机抽取几十张图片、按类别拼成网格图,用肉眼看一遍训练集,这比任何统计数据都直观。这一步没办法自动化,但它能帮你对“模型学到什么”建立第一印象。

3. 用PyTorch迁移学习跑通第一版识别模型:MobileNetV2从数据管道到训练闭环

3.1 为什么选迁移学习而不是自己搭CNN:小数据实验下的最优解

树叶识别通常不是一个从头训练卷积神经网络的场景,公开树叶数据集的规模在每类几十到几百张之间,这点数据量支撑不了从头训练一个深层CNN。从头训练意味着网络要同时学习底层特征(边缘、纹理、颜色)和高层语义特征(叶形、叶脉),而迁移学习能直接把在ImageNet上预训练好的底层特征搬过来,我们只需要在目标任务上重新训练分类头。通俗说,预训练模型已经“见过”各种自然图像,知道什么是边缘、什么是纹理,我们让它在此基础上认识“桂花叶”和“银杏叶”的区别。

技术选型上,我一般用MobileNetV2。它在ImageNet上精度不错,参数量只有约350万,是ResNet50的四分之一左右,CPU也能跑推理,对课程设计和嵌入式部署都友好。更重要的是,MobileNetV2的深度可分离卷积结构在叶片这种纹理密集的对象上表现稳定,不会像一些轻量网络那样为了速度牺牲太多精度。

PyTorch实现迁移学习的标准姿势是:加载带预训练权重的MobileNetV2,把最后的分类层替换成我们自己的全连接层,然后分阶段训练。如果数据量极小,可以先把骨干网络全部冻结,只训练分类头;如果数据量够大,再解冻最后几层做微调。下面两节给出完整的代码。

3.2 PyTorch数据管道:ImageFolder与transforms的完整写法

数据管道是训练闭环的地基。PyTorch官方推荐的写法是用torchvision.datasets.ImageFolder配合torchvision.transforms。注意,训练集和验证集的transform必须不同:训练集要加入随机增强,验证集只做尺寸调整和归一化,否则验证集损失会被增强噪声污染。

import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader train_transform = transforms.Compose([ transforms.RandomResizedCrop(224, scale=(0.8, 1.0)), transforms.RandomHorizontalFlip(), 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_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]) ]) train_dataset = datasets.ImageFolder('leaf_dataset/train', transform=train_transform) val_dataset = datasets.ImageFolder('leaf_dataset/val', transform=val_transform) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4, pin_memory=True) val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4, pin_memory=True) print(f"训练集类别: {train_dataset.classes}") print(f"训练集样本数: {len(train_dataset)}, 验证集样本数: {len(val_dataset)}")

这段代码里有几个参数值得解释。RandomResizedCrop(224, scale=(0.8, 1.0))会随机裁剪一个比例在0.8到1.0之间的区域,再缩放到224×224,这相当于一种尺度不变性的模拟,让模型见过不同远近的叶片。ColorJitter的三个参数brightness、contrast、saturation分别控制亮度、对比度和饱和度的随机扰动范围,真实拍摄环境光照变化很大,加上这个变换能显著提升模型在不同天气条件下的鲁棒性。

Normalize里的均值和标准差是ImageNet数据集的统计值,使用预训练模型时必须沿用这套数值,不能自己重新计算。很多人把这一步漏掉或者改成/255.0,结果是模型推理时特征分布不匹配,精度莫名其妙下降。

num_workers=4是数据加载的并行进程数,Windows上建议设置成0或2,否则可能会因为多进程启动机制问题报错。pin_memory=True在GPU训练时可以加速数据从内存到显存的拷贝,纯CPU训练时设置这个参数没有实际帮助,反而会增加内存占用。

3.3 训练脚本:冻结骨干、只练分类头,最小可运行版本

迁移学习的最小可运行版本分三步走:加载预训练模型、替换分类头、定义训练循环。下面是一个完整的单文件训练脚本。

import torch import torch.nn as nn import torch.optim as optim from torchvision import models DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"使用设备: {DEVICE}") def build_model(num_classes, freeze_backbone=True): model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.DEFAULT) in_features = model.classifier[1].in_features # 替换分类头 model.classifier[1] = nn.Linear(in_features, num_classes) if freeze_backbone: for param in model.features.parameters(): param.requires_grad = False return model.to(DEVICE) def train_one_epoch(model, loader, criterion, optimizer): model.train() total_loss, correct, total = 0.0, 0, 0 for images, labels in loader: images, labels = images.to(DEVICE), 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.max(1) correct += preds.eq(labels).sum().item() total += labels.size(0) return total_loss / total, correct / total model = build_model(num_classes=len(train_dataset.classes)) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.classifier.parameters(), lr=1e-3) EPOCHS = 30 for epoch in range(1, EPOCHS + 1): train_loss, train_acc = train_one_epoch( model, train_loader, criterion, optimizer) # 验证 model.eval() val_loss, val_correct, val_total = 0.0, 0, 0 with torch.no_grad(): for images, labels in val_loader: images, labels = images.to(DEVICE), labels.to(DEVICE) outputs = model(images) loss = criterion(outputs, labels) val_loss += loss.item() * images.size(0) _, preds = outputs.max(1) val_correct += preds.eq(labels).sum().item() val_total += labels.size(0) print(f"Epoch {epoch:02d} | " f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} | " f"Val Loss: {val_loss / val_total:.4f} " f"Acc: {val_correct / val_total:.4f}") torch.save(model.state_dict(), 'leaf_model_v1.pth')

这段代码的核心设计是freeze_backbone=True,它就只更新最后分类层的权重。这样做的好处是训练速度快、不容易过拟合,适合树叶样本量不多的项目;坏处是骨干网络的特征未必完全适合树叶这种细粒度分类任务,后期可以把freeze_backbone改成False,对整个网络做低学习率微调。

训练循环里我用了交叉熵损失CrossEntropyLoss,这是多分类任务的默认选择,它的输出已经是logits,不需要在模型末尾额外接Softmax。优化器用Adam而不是SGD,因为Adam对学习率不敏感,起步阶段更省心。optimizer = optim.Adam(model.classifier.parameters(), lr=1e-3)注意到这里只传入了分类层的参数,冻结的骨干不会参与更新,省内存也省计算。保存模型时用state_dict(),保存的是权重字典,比保存整个模型对象更轻量,也方便后面换网络结构重新加载。

3.4 参数说明与本地CPU/GPU适配

训练脚本里最需要关注的是EPOCHS、batch_size和DEVICE。纯CPU环境训练MobileNetV2,30个epoch每轮在几百张图片上大约需要几分钟到十几分钟,总时长可以接受。但如果你的图片特别多,比如几千张,CPU训练会非常煎熬,这时有两个选择:调低EPOCHS到15左右,或者换用GPU。

GPU训练要注意一个常见问题:如果你的环境显存只有4GB,batch_size=32可能会爆显存,报CUDA out of memory。这时候不要急着换小模型,先把batch_size降到16或8,再把num_workers调低,往往就能解决。如果担心batch size降低影响模型效果,可以同步把学习率从1e-3降到1e-4左右,因为小的batch size训练时梯度噪声更大,偏小的学习率更稳定。

TensorFlow/Keras用户会看到不同的写法:在Keras里用MobileNetV2(weights='imagenet', include_top=False),再接GlobalAveragePooling2D和Dense,效果一样,只是在代码组织和训练方式上略有差别。如果你只是跑通一个课程项目,选自己最熟悉的一个框架即可,两个框架的训练结论和调参逻辑完全通用。

4. 把识别准确率从“掉到及格”拉回“高分”:学习率、batch size与数据增强的调参顺序

4.1 学习率:从1e-3起步,观察loss曲线再决定降温策略

迁移学习的第一次训练,把学习率设为1e-3是稳妥起点。这个值不会让分类头的权重更新过猛,也不会慢到半天看不到损失下降。训练过程中要盯的不是训练集准确率,而是验证集损失曲线。理想曲线是验证损失前几个epoch快速下降,然后进入平台期;如果验证损失在某个epoch后开始反弹上升,训练损失还在下降,那说明过拟合已经开始,需要提前停止。

Adam优化器自带自适应学习率,但“自适应”不等于不用调。训练后期我一般会用手动方式降温:每5个epoch把学习率乘以0.5,或者直接用torch.optim.lr_scheduler.ReduceLROnPlateau,当验证损失连续3个epoch不再下降时自动把学习率减半。代码写法如下:

from torch.optim.lr_scheduler import ReduceLROnPlateau scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=3) # 在每个epoch验证结束后调用 scheduler.step(val_loss_avg)

mode='min'表示监听验证损失的最小值,factor=0.5表示每次降低为原来的50%,patience=3表示容忍连续3个epoch不改善。这个策略比固定epoch数训练更可靠,因为它能在模型开始过拟合时自动踩刹车。如果你发现学习率从1e-3开始训练时损失不下降,可以先检查数据管道是否正常,而不是急着调学习率;确认数据没问题后,可以尝试1e-2,但要注意1e-2配合Adam在训练初期可能让损失震荡得很厉害。

4.2 batch size与步数:显存不够时不要先动模型,先动数据管道

batch size的决定因素不是模型精度,而是硬件资源。树叶数据集通常很小,batch size在16到64之间都能正常收敛。大batch size的梯度方向更平滑,训练更稳定;小batch size引入的噪声反而在某些任务上能帮助模型跳出局部最优。但这些都是次要的,显存不够时优先降batch_size和num_workers,不要一上来就换轻量模型。MobileNetV2已经是轻量模型了,再换到MobileNetV3或者改输入分辨率到160×160,损失的是细粒度纹理信息,得不偿失。

训练步数方面,每个epoch的迭代次数等于训练集样本数除以batch size。如果训练集只有300张图、batch size为32,一个epoch只有不到10次迭代,训练非常不充分。这种情况不需要纠结,直接加大EPOCHS到50甚至80,并配合早停和模型保存,等验证损失不再下降时停止即可。

还有一个容易忽略的batch size联动参数:PyTorch里DataLoader的drop_last。当训练集样本数不能被batch size整除时,最后一个batch会变小,如果不希望最后一个小batch干扰BatchNorm的统计,可以设置drop_last=True。树叶图像分类用MobileNetV2时,BatchNorm层在eval模式下会使用训练阶段累计的均值方差,这个参数影响不大,但能避免训练和验证阶段行为不一致。

4.3 数据增强:旋转、翻转、色彩抖动对付真实照片

数据增强是树叶识别系统提升泛化能力的关键手段。常见的增强策略包括随机水平翻转、随机旋转、随机裁剪和颜色抖动。水平翻转是零成本的增强,因为树叶没有方向性;随机旋转要小心90度倍数,如果某些树叶有明显的上下朝向,180度翻转可能让类别语义混淆。更稳妥的做法是把旋转范围限制在[-30, 30]度左右。

from torchvision import transforms augment_transform = transforms.Compose([ transforms.RandomRotation(degrees=30), transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)), transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.05), transforms.RandomHorizontalFlip(p=0.5), transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ])

RandomAffine的translate参数让图片在水平和垂直方向随机平移10%,这能模拟拍照时树叶没有严格居中的情况。hue=0.05的控制范围很小,因为色相偏移太大会让树叶颜色失真,模型会学到错误的颜色分布。要记住,增强虽然能提升泛化能力,但也可能引入噪声,尤其当某些类别本身外观相似时,过强的增强反而会让模型更难区分它们。所以增强参数要一点一点加,每次实验只改一个变量,不要同时调整三个增强项的强度。

5. 树叶识别系统避坑:过拟合、类别不均衡与推理阶段的三类翻车现场

5.1 过拟合:训练精度95%、验证精度60%,先查这三处

现象:训练集准确率一路冲到95%以上,验证集准确率却停在60%左右,损失曲线从第10个epoch开始明显分离。

原因:最常见的不是模型太复杂,而是训练集和验证集有数据泄漏。可能你在切分数据时没有按类别分层,导致同一棵树的不同照片被分到了两个集合;也可能你直接下载的公开数据集本身含有重复图片,训练时见过、验证时又遇见。其次是数据量太少,模型直接把训练集的噪声背了下来。

解决:第一步,检查切分脚本是否加了random.seed,没有的话先固定种子重新切分。第二步,对所有图片计算哈希值,删除完全重复的文件,这一步可以简单实现:

import hashlib from PIL import Image def file_hash(path): with open(path, 'rb') as f: return hashlib.md5(f.read()).hexdigest() duplicates = {} for cls in ['ginkgo', 'osmanthus', 'maple']: for img_name in os.listdir(f'leaf_dataset/train/{cls}'): path = f'leaf_dataset/train/{cls}/{img_name}' h = file_hash(path) if h in duplicates: print(f"重复图片: {path} 与 {duplicates[h]}") else: duplicates[h] = path

如果上述两步都没问题,再考虑给模型加一点约束:把freeze_backbone=False改成True,或者增大分类头的Dropout比例。MobileNetV2的原始分类头里有一个Dropout(p=0.2),这个比例可以适当提高到0.5,代价是训练收敛变慢但泛化能力更好。最简单有效的方法还是早停:每轮验证后比较当前验证损失,如果连续5个epoch没有改善就保存上一轮模型并终止训练。

5.2 类别不均衡:银杏叶子少,模型就永远猜桂花

现象:训练完后,模型对所有输入都倾向于输出样本数最多的那个类别,验证集准确率尚可,但看混淆矩阵会发现少数类几乎从未预测成功。

原因:交叉熵损失在类别不均衡时会偏向多数类,因为多数类样本对总损失的贡献更大,优化器为了降低总损失而牺牲少数类。

解决:最直接的办法是给每个类别加权,让少数类样本在损失计算中占有更高权重。PyTorch里的CrossEntropyLoss支持weight参数,权重可以按“样本数的倒数归一化”来设置:

import torch from torch.utils.data import WeightedRandomSampler class_counts = [310, 95, 140] # 按 train_dataset.classes 顺序 total = sum(class_counts) class_weights = [total / c for c in class_counts] weights = [class_weights[label] for _, label in train_dataset.samples] sampler = WeightedRandomSampler(weights, num_samples=len(weights), replacement=True) train_loader = DataLoader(train_dataset, batch_size=32, sampler=sampler)

WeightedRandomSampler会让每个epoch中少数类样本被采样到的概率变大,replacement=True表示可以重复采样,这是处理类别不均衡的标准做法。使用sampler之后要注意,DataLoader的shuffle参数必须设为False,因为sampler已经完成了打乱逻辑,两者同时开启会冲突。

5.3 推理阶段翻车:训练时224,推理时忘了resize

现象:模型在训练集和验证集上准确率都很高,但单独写一个推理脚本加载模型测单张图片时,准确率暴跌,甚至出现类别索引错乱。

原因:推理脚本没有完整复刻验证集的transform。最常见的是忘记Resize和CenterCrop,直接把原图缩放到224×224,破坏了图片的宽高比;或者忘了Normalize,把像素值以0到1范围送入模型,而模型预期的是标准化后的数值。另一个低级错误是模型用了GPU训练并保存state_dict,推理时在CPU加载却忘了加map_location='cpu',导致加载失败或者自动跳到GPU占用显存。

解决:推理脚本里把所有预处理步骤封装成一个函数,直接从之前定义好的val_transform复制,不要凭记忆重写。加载模型时加上设备映射:

import torch from torchvision import models, transforms from PIL import Image # 这里使用与训练时完全相同的预处理 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]) ]) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = models.mobilenet_v2(num_classes=len(train_dataset.classes)) model.load_state_dict(torch.load('leaf_model_v1.pth', map_location=device)) model.to(device).eval() def predict_image(image_path): img = Image.open(image_path).convert('RGB') tensor = val_transform(img).unsqueeze(0).to(device) with torch.no_grad(): outputs = model(tensor) _, pred = outputs.max(1) return train_dataset.classes[pred.item()] print(predict_image('test_samples/test_01.jpg'))

这段代码的关键是把model.eval()放在推理前,它会关闭BatchNorm和Dropout的随机行为,保证推理输出确定性。unsqueeze(0)是为了在图片张量前增加一个batch维度,因为模型接受的输入形状是(batch_size, 3, 224, 224),单张图片也要“伪装”成一个大小为1的batch。

5.4 没有GPU:显存不足、torch.cuda.is_available()为False时怎么办

现象:代码报了CUDA out of memory,或者明明装了CUDA版的PyTorch,torch.cuda.is_available()却返回False,一查发现PyTorch装在CPU版本下。

原因:显存不足通常不是模型太大,而是batch size、输入分辨率、num_workers共同叠加导致显存峰值过高。is_available()返回False的原因则是安装时用了CPU版本的PyTorch,或者显卡驱动和CUDA版本不匹配。

解决:如果没有NVIDIA GPU,直接走CPU训练路线就好,树叶识别系统这种小数据集CPU完全扛得住,只是训练轮次要放宽,每个epoch慢一点但能等。如果要装GPU版PyTorch,安装命令要指定CUDA版本,以PyTorch官方安装为例,安装CUDA 12.1版本的指令是pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121。安装完之后用torch.cuda.get_device_name(0)验证是否可见。还有一个易错点:即使显存不足,也优先尝试用torch.cuda.amp混合精度训练而不是直接放弃GPU。移动端或部分入门级显卡对混合精度支持有限,如果显存还是不够,最后才考虑换输入分辨率到160×160。

6. 最后一块拼图:用特征热力图验证模型真的在看树叶,而不仅是背景

模型训练完成、准确率看着不错,还不算真正的结束。有一个非常扎心的问题:如果模型是根据白色背景或花盆特征来分类的,测试集准确率照样能很高。要证明模型“认的是树叶”,需要借助可视化工具来看模型的注意力集中在图像的哪个区域。Grad-CAM是最常用的方案,它利用最后一层卷积输出的梯度值计算每个特征通道的权重,再加权求和得到注意力热力图。

import cv2 import torch import numpy as np from torchvision import models, transforms from PIL import Image def grad_cam(model, image_tensor, target_layer): activations = {} gradients = {} def forward_hook(module, input, output): activations['value'] = output def backward_hook(module, grad_input, grad_output): gradients['value'] = grad_output[0] hook_fwd = target_layer.register_forward_hook(forward_hook) hook_bwd = target_layer.register_full_backward_hook(backward_hook) model.zero_grad() output = model(image_tensor) pred_class = output.argmax(dim=1).item() score = output[0, pred_class] score.backward() act = activations['value'].squeeze(0) # (C, H, W) grad = gradients['value'].squeeze(0) # (C, H, W) weights = grad.mean(dim=(1, 2)) # 每个通道的平均梯度 cam = torch.zeros(act.shape[1:], dtype=torch.float32) for i, w in enumerate(weights): cam += w * act[i] cam = torch.relu(cam) cam = cam.numpy() cam = cv2.resize(cam, (224, 224)) cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8) hook_fwd.remove() hook_bwd.remove() return cam, pred_class # 使用示例 target_layer = model.features[-1] # MobileNetV2 最后一个反向残差块的输出 cam_map, pred = grad_cam(model, image_tensor, target_layer)

MobileNetV2中model.features[-1]是最后一个卷积阶段输出的特征图,空间分辨率是输入图像的1/32左右,Grad-CAM在这里的空间粒度足以框出叶片主体。如果热力图集中在叶片纹理区域,说明模型学到了有意义的叶子特征;如果集中在背景边缘或者图片角落,说明训练数据出了问题,需要回到数据增强或数据清洗环节。演示视频录制时把这个Grad-CAM可视化放在界面的侧边栏,会比单纯显示“识别为桂花,置信度95%”更有说服力,评审老师通常也认可这种验证思路。

还有一个值得养成的习惯:每次训练完都保存三样东西,模型权重、数据增强配置的完整代码、还有验证集上的混淆矩阵。混淆矩阵比单独一个准确率数字更能反映问题——它告诉你哪些类别最容易互相混淆,比如“银杏”和“鹅掌楸”的叶子外形接近,模型分不清是正常现象,下次可以针对性地收集这两类的更多样本,而不是盲目调参。把这三样东西固定下来,后续任何实验都有据可查。

我自己在任何图像分类项目里都保留一个笨办法:挑5张训练集里置信度最高的图片和5张置信度最低的图片,每次训练完都打印出来看一眼。置信度最高的如果都是背景干净的棚拍图,说明模型开始依赖背景;置信度最低的如果全是同一种类别的不同角度照片,说明这个类别的特征即使对人来说也难分辨。看清这些边界,才知道下一步该往哪个方向投入,而不是对着准确率数字瞎猜。希望这些踩坑记录和经验能帮你在做树叶识别系统时少走几段弯路。

本文还有配套的精品资源,点击获取

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询