ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

基于深度学习的水果图像识别系统:CNN原理、PyTorch实现与部署避坑指南

基于深度学习的水果图像识别系统:CNN原理、PyTorch实现与部署避坑指南 简介基于深度学习的水果识别系统完整源码与训练模型包面向计算机视觉方向的毕业设计学生和深度学习入门开发者。系统采用卷积神经网络构建分类模型覆盖数据预处理、模型训练、保存、推理及结果展示等环节可帮助读者掌握图像分类项目的完整落地流程。压缩包内共277个文件大小约17.49MB主要包含8个Python源码文件、114个JavaScript脚本、84个GIF动图、26个CSS样式表、多个HTML页面及图片资源等其中py文件对应核心训练与推理模块HTML/CSS/JS构成可视化交互界面训练好的模型可直接用于水果类别预测。目前已有59人学习适合作为课程设计或毕业设计的参考脚手架也可基于它扩展更多水果种类或接入摄像头实现实时识别。1. 拿到这个压缩包你其实拿到了一套图像分类的标准流水线一个标题里带python源码训练好的模型的zip解压出来里面基本就是这几样东西一个放着分类图片的data文件夹、一个train.py、一个predict.py、一个.pth或.pt结尾的权重文件。所谓基于深度学习的水果识别系统本质就是用一个卷积神经网络CNN对水果图像做分类输入一张照片输出这是什么水果、置信度是多少。它解决的是一类很具体的问题面对几十个水果类别、复杂的光照和遮挡让程序像人一样认出水果。这套东西特别适合三类人做课程设计需要交一个能跑通的完整项目的学生、刚啃完深度学习入门教程想动手做第一个CNN实战的人、以及要给客户或同学演示我这边已经跑起来了的工程师。接下来我把这套系统从模型选型、训练到加载推理一条线讲透并且把最容易翻车的那几个地方提前摆出来。2. 先讲清楚原理水果识别为什么绕不开深度学习模型怎么选2.1 传统CV在水果识别上的三个死穴光照、遮挡、相似类别水果识别看起来是个简单图像分类但真用老办法做一遍就知道有多痛苦。传统做法是先做特征工程用颜色直方图描述苹果的红色用纹理特征描述猕猴桃表皮的绒毛用边缘检测勾勒香蕉的轮廓最后丢给SVM或者随机森林。问题在于水果图像实在是太看天吃饭了。第一个死穴是光照。同一个红富士直射阳光下和树荫里拍出来的RGB值能差出好几个等级颜色直方图直接就偏了你根本没法写一套规则同时适配晴天和阴天。第二个死穴是遮挡水果长在枝头常有叶子挡住一半轮廓被切断、边缘特征全乱基于形状的方法当场失效。第三个死穴是相似类青苹果和梨的轮廓几乎一样柠檬和青柠颜色接近香梨和苹果切开以后更是难分靠手工设计特征很难在这些细粒度类别之间拉开距离。深度学习把怎么提取特征这件事交给了卷积层自己学习从颜色、纹理到局部形状逐层抽象不再需要人肉设计特征。这也是为什么现在几乎所有识别类系统都在用CNN而不是SVM加颜色直方图。你只需要提供足够多的标注图片网络自己会找到那些人眼说不出来但确实有效的判别线索。2.2 骨干网络怎么选ResNet18、MobileNetV3、EfficientNet的取舍解压源码后你会发现train.py里肯定有一个骨干网络的选项常见的是ResNet、MobileNet、EfficientNet这几个家族。它们是用来做图像分类的预训练网络在ImageNet那1000类上学到过通用的视觉特征拿到水果识别这种小任务上属于降维打击。选哪个取决于你的运行环境我通常按下面这张表来定网络参数量级训练设备推理场景适用阶段ResNet18约1100万普通办公电脑CPU能训CPU/GPU均可课设、入门demo首选ResNet50约2500万建议有独立显卡GPU追求更高精度MobileNetV3百万级别小于前两者任意设备树莓派、手机、嵌入式要部署到边缘设备EfficientNet-B0500万上下单卡可训CPU尚可精度和速度都想要水果识别往往只有几十个类别和ImageNet的1000类相比是小头所以不需要上几百层的超大网络。模型容量一旦过剩训练集上表现很好换到新图片上反而容易过拟合。我自己做课设或者快速验证方案时默认用ResNet18稳、快、调参成本低。如果要在树莓派或者Android设备上演示就换MobileNetV3。如果模型文件是项目里已经训练好的那压缩包里应该包含了对应的网络类型定义接手时先打开源码看它import了哪个models函数再决定要不要换。2.3 从图片到概率推理链路内部发生了什么不管哪个网络推理链路的骨架是一样的。一张224x224x3的RGB图片先做归一化把像素值从0-255缩放到模型期望的范围然后变成一个四维张量[batch, 3, 224, 224]进了卷积网络。卷积层是一组可学习的滤波器前面的层学到边缘、色块中间的层学到纹理、局部形状后面的层学到果梗表面光泽圆形轮廓这些高层概念。最后全局池化压成一个特征向量接一个全连接层输出一个长度等于类别数的logits向量。logits model(x) # 形状: [batch, num_classes] prob torch.softmax(logits, dim1) # 转成和为1的概率分布 pred torch.argmax(prob, dim1) # 取概率最大的类别下标这段逻辑几乎出现在每一个分类项目的预测脚本里。注意softmax的dim1是按类别维度做归一化不是按batch维度argmax之后拿到的是类别下标还要映射回真正的类别名字比如下标0对应apple、1对应banana。这里有个新手容易忽略的坑softmax永远能输出一个看起来自信的结果哪怕输入一张小狗照片它也会把概率分配给最接近的某个水果类。所以工程上正规的做法是设一个置信度阈值概率低于0.6就返回无法识别而不是硬报一个答案。3. 把训练脚本跑起来数据集目录、数据增强与三个必调参数3.1 数据集目录结构train/val分开类别名就是文件夹名源码里的数据加载几乎都是用PyTorch的ImageFolder这个类要求图片按类别放在文件夹里文件夹的名字就是类别名。标准布局长这样dataset/ train/ apple/ 1.jpg 2.jpg banana/ 1.jpg orange/ 1.jpg val/ apple/ 1.jpg banana/ 1.jpg训练目录和验证目录要严格分开每一类都对应一个子文件夹。ImageFolder会自动扫描目录按字母序给类别编号apple是0、banana是1、orange是2。这里有两个值得提前注意的地方图片格式jpg、png都行尺寸不需要统一DataLoader的transform会自动缩放一个类别最好至少准备50张训练图20个类别就得1000张以上如果只有十几张建议直接用第6章的迁移学习方案否则很容易过拟合。另外不要在图片文件名里写中文也不要把类别文件夹起成红富士苹果好的中文路径在Windows上会带来一串编码问题后面第5章细说。3.2 train.py拆解数据增强、模型替换、损失与学习率训练脚本的核心其实是三块数据预处理、网络初始化、训练循环。先看数据预处理部分我一般这样写import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, transforms, models # 训练集随机裁剪缩放模拟距离变化水平翻转颜色抖动模拟光照变化 train_tf transforms.Compose([ transforms.RandomResizedCrop(224), transforms.RandomHorizontalFlip(), transforms.ColorJitter(brightness0.2, contrast0.2, saturation0.2), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 验证集只做缩放、中心裁剪和归一化不做随机增强 val_tf transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])RandomResizedCrop会随机裁剪出一块区域再缩放到224模拟相机距离远近不同带来的尺度变化ColorJitter调整亮度、对比度、饱和度就是为了对抗2.1里说的光照问题。验证集不能用RandomResizedCrop否则每次评估结果都在变没法判断模型好坏。Normalize的均值方差用的是ImageNet的标准值[0.485, 0.456, 0.406]因为加载的是ImageNet预训练权重输入分布必须要和预训练时保持一致。接下来是模型的初始化和训练循环# 水果类别数和你的data/train下的子文件夹数量一致 num_classes 10 device torch.device(cuda if torch.cuda.is_available() else cpu) # 加载ImageNet预训练模型替换最后一层全连接 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(512, num_classes) model model.to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.SGD(model.parameters(), lr0.001, momentum0.9, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max30) for epoch in range(epochs): model.train() running_loss 0.0 for inputs, labels in train_loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * inputs.size(0) scheduler.step() val_acc evaluate(model, val_loader, device) print(fepoch {epoch1}/{epochs} loss {running_loss/len(train_dataset):.4f} val_acc {val_acc:.4f})model.train()这行必须写它告诉BatchNorm和Dropout当前处于训练模式如果不写BN层会一直用推理时的统计量训练发散是常有的事。每轮结束记得调用scheduler.step()做学习率衰减。evaluate函数需要包在torch.no_grad()里否则评估时期望也会建立计算图显存或内存直接被吃满这也是很多人训练时内存爆掉的直接原因。3.3 三个必须自己改的参数image_size、batch_size、epochs压缩包里的源码用的是别人调好的参数到你自己数据集上一定要改三个东西。第一是image_size最常见的是224或256。224是ImageNet标准尺寸ResNet18原生输入就是224x224水果识别这个分辨率够用。别一上来就设512训练时间暴涨但准确率提升有限小模型根本吃不下那么多细节。第二是batch_size取决于你的显卡。8G显存用32或64纯CPU训练建议用16再大内存占用高而且收敛速度也不见得快。batch_size太小比如4或8时BN统计量不稳定验证准确率会上下跳得厉害。第三是epochs迁移学习场景30到60轮足够从零开始训练才需要100轮以上。另外一个容易被忽略的学习率问题。初始学习率用0.001是标准做法配合CosineAnnealingLR慢慢衰减。如果发现训练loss飞起来了先把学习率降到3e-4再试。如果loss在0.3左右震荡不降大概率是数据问题不是学习率问题回去看看是不是类别严重不均衡。训练过程中一定要做保存最优模型这个动作别傻等到最后一轮才存权重——保存验证集准确率最高的那个模型这也是源码里训练好的模型那个文件的由来。4. 加载训练好的模型做推理predict.py的写法和批量评估4.1 模型文件里装的到底是什么state_dict和它背后的约束.pth文件用torch.save()保存有两种常见内容一种是只保存state_dict也就是网络的参数字典另一种是把整个模型对象序列化进去。正规项目里默认保存state_dict因为它不携带网络结构文件更小跨设备加载也更干净。state_dict是一个OrderedDictkey形如conv1.weightfc.weight每个key对应网络里一个权重张量。加载state_dict的时候有个硬性约束你必须先用代码创建出一个结构完全一样的模型再调用load_state_dict做填充。也就是说压缩包里只有权重文件还不够train.py里那个model.fc nn.Linear(512, num_classes)的改动也必须同步到预测脚本里来。如果你换了一个类别的数量或者把fc输出改成了别的维度load_state_dict直接报size mismatch。还有更常见的坑是训练时用了DataParallel保存下来的key全部带module.前缀直接加载会报unexpected key这个第5章专门讲。4.2 单张图片推理predict.py的标准写法理想情况下压缩包里自带的predict.py应该是能直接运行的。如果作者没写完整我建议你自己补一个这是这套系统最有价值的落地脚本。下面是我常用的写法兼容了上面说的module前缀坑import torch import torch.nn as nn from PIL import Image from torchvision import transforms, models def predict(image_path, model_path, class_names, devicecpu): # 推理预处理必须和训练验证时完全一致 tf transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) img Image.open(image_path).convert(RGB) tensor tf(img).unsqueeze(0).to(device) # 加batch维度 # 重建模型结构类别数从class_names列表长度推断 model models.resnet18(weightsNone) model.fc nn.Linear(512, len(class_names)) state torch.load(model_path, map_locationdevice) # 如果权重key带了module.前缀说明训练时用了DataParallel if list(state.keys())[0].startswith(module.): state {k.replace(module., ): v for k, v in state.items()} model.load_state_dict(state) model.eval() with torch.no_grad(): logits model(tensor) prob torch.softmax(logits, dim1) idx torch.argmax(prob, dim1).item() return class_names[idx], prob[0][idx].item()这段代码里最容易被忽略的是最后那三行model.eval()切到推理模式让BN层用训练时学到的全局统计量torch.no_grad()关掉梯度计算既省内存又让速度提升不少Image.open之后要convert(RGB)否则遇到PNG带透明通道的图会直接报错或者维度对不上。类别的顺序是class_names这个列表给定的它必须和训练时的文件夹顺序一致最保险的办法是从训练时的类别字典里导出来见4.3。4.3 批量和FPS评估别拿单张速度当部署速度单张推理跑通之后还要做一件事把整个验证集跑一遍算出准确率和每类准确率。只测一张图就宣布系统完成是典型的课设思维真正投入前必须看统计结果。批量评估脚本核心是这样的def evaluate(model, val_loader, device, class_names): model.eval() correct 0 total 0 per_class_correct {name: 0 for name in class_names} per_class_total {name: 0 for name in class_names} with torch.no_grad(): for inputs, labels in val_loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) _, preds torch.max(outputs, 1) for label, pred in zip(labels, preds): per_class_total[class_names[label]] 1 if label pred: correct 1 per_class_correct[class_names[label]] 1 total 1 return correct / total, per_class_correct, per_class_total这个脚本不但算总准确率还统计每个类别的命中情况。比如总准确率90%看着不错一拆开发现香蕉类98%、柠檬类只有60%说明柠檬和其他类混了得单独加数据或者增加柠檬类的样本权重。至于FPS建议用100张图片连续推理后取平均值别拿单张图除以耗时算所谓FPS第一次加载模型、CPU频率波动都会把结果带偏。CPU上ResNet18一张图大约几十到两百毫秒想跑到实时得换MobileNetV3或者转ONNX这部分放在第6章讲。5. 避坑训练和加载模型时我真实翻过的车这套系统看似简单但我在不同数据集上反复折腾过下面的坑不是猜测都是现象明确、原因可查的真实经历。每一条都按现象、原因、解决三步说明你遇到同样的报错可以直接照方抓药。5.1 加载权重报错unexpected key字典里全是module.前缀现象predict.py执行到load_state_dict时抛错提示unexpected key module.conv1.weight或者报了missing keys。原因训练时用了torch.nn.DataParallel把模型包了一层保存出来的state_dict所有key都被加上了module.前缀。预测脚本里没有做兼容于是加载时找不到对应的层名。解决加载后对state_dict做一次字符串替换把前缀剥掉。我上面4.2的代码里已经写好了关键就一行state {k.replace(module., ): v for k, v in state.items()}顺便提醒这个坑在单卡训练时不出现但如果你是在实验室多卡机器上训练、然后拷回自己电脑推理十有八九会碰上。判断方法很简单打印list(state.keys())[0]看看有没有module.。5.2 训练loss死活不降准确率停在类别数分之一现象网络结构没问题、学习率也调了但训练了大几十轮loss一直在1.0以上震荡验证准确率大约等于最常见的那个类别占比。比如3类水果各占60%、20%、20%准确率就一直卡在60%附近。原因类别严重不均衡模型学到的策略是把所有样本都判成占比最多的那一类就能拿到低loss于是它懒得去学少数类的特征。解决第一优先是对损失函数加权。统计每个类别的样本数算出权重作为CrossEntropyLoss的weight参数让少数类的错误被放大class_counts torch.tensor([500, 200, 50], dtypetorch.float) class_weights 1.0 / class_counts class_weights class_weights / class_weights.sum() criterion nn.CrossEntropyLoss(weightclass_weights.to(device))更彻底的办法是用WeightedRandomSampler对少数类做上采样让每个batch里各类别数量均衡。我一般先试loss加权不行再上采样。这个坑的隐蔽之处在于loss下降曲线看着很平滑容易误以为训练正常其实模型已经变成了拒绝对多数类之外的类别做预测。5.3 训练集准确率98%验证集只有60%被过拟合闷头一棒现象训练的loss一直在降、训练集准确率逼近一百但每次跑验证集准确率都上不去两个数值差距越来越大。原因典型过拟合。诱因通常是三类训练数据太少模型又把细节背下来了数据增强开得太弱epochs跑太多。还有一个容易被忽视的情况是验证集预处理错了比如训练时用了RandomResizedCrop验证集也照着写随机性导致评估结果上下乱跳让你误以为模型不行。解决先固定预处理训练用随机增强、验证用Resize加CenterCrop这是标准对称写法然后增加ColorJitter和RandomHorizontalFlip强度相当于把数据量撑大最后做早停保留验证集准确率最高的epoch权重而不是最后一轮。我现在的做法是把保留最优模型写进训练循环里每一轮如果验证准确率创新高就覆盖保存一次到best_model.pth这样就算epochs设超了也不用怕。5.4 中文图片路径和中文类别名导致推理结果全乱现象Windows下图片路径带中文ImageFolder加载时报错或者图片读取出来是黑的或者训练时类别是苹果香蕉这类中文名predict.py返回的结果对不上。原因PyTorch在Windows上和文件编码的配合很脆弱中文字符串的编码方式不一致容易出现UnicodeDecodeError另外ImageFolder按字典序给类别编号如果不同脚本之间对类别排序的规则不完全一样编号就会错位模型输出下标0对应的类别在另一个脚本里可能是下标1。解决内部统一用英文做文件夹名和类别名比如apple、banana只在最终显示结果时再映射成中文。训练结束后把类别列表原样保存成一个json文件import json with open(class_names.json, w, encodingutf-8) as f: json.dump(class_names, f, ensure_asciiFalse, indent2)推理时用json.load加载类别列表不要用任何方式重新排序。这个json文件要和模型文件放在一起zip包里如果缺了它你自己补一个按训练文件夹的字母序把类别名手写进列表。5.5 CPU推理速度慢到没法演示一两秒才出一张图现象训练好的模型在GPU上挺快但换到普通笔记本或者演示机上一张图要几百毫秒甚至一秒以上根本没法做实时演示。原因ResNet18本身就是百毫秒级的网络CPU上跑就是这么慢再加上如果机器没启用AVX指令集或者跑在虚拟机里速度会雪上加霜。很多人以为代码优化一下就好其实瓶颈在网络本身的浮点运算量。解决优先把模型换成MobileNetV3推理速度比ResNet18快一个数量级水果识别这个任务精度损失很小。其次是用ONNX Runtime跑转换后的模型转换脚本几年前就成熟了一行python命令的事情。动态量化对全连接层有效但卷积层收益不大真想提速还得是ONNX路线。演示场景如果要求不高还有个土办法把输入尺寸从224降到160速度明显变快准确率只掉一两个点适合现场demo应付过去。6. 进阶迁移学习调优、模型瘦身与结果验证方法如果这套系统已经能跑到90%以上的准确率下一步值得做的是往能投入使用的方向靠拢。这里说三个我常用的进阶操作。第一个是增量加类别的迁移学习比如现在系统能认10种水果你摘了20张杨梅的图想加一类。不要重训整个网络正确的做法是加载已有权重把fc层的输出改成11然后冻结backbone只训练新层for name, param in model.named_parameters(): param.requires_grad False if not name.startswith(fc) else True optimizer torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr0.0001)注意优化器要用filter过滤掉requires_gradFalse的参数否则冻结不生效更新还是会发生。这样只需要几百步迭代就能学会新类别而不会把老类别的特征破坏掉。第二个是模型瘦身如果最终目标是塞进树莓派或者App里可以把权重转成半精度GPU推理显存减半速度还提升CPU场景优先转ONNX Runtime这是目前性价比最高的方式。第三个是验证方法别只看总准确率把每类的准确率列成表再看一眼哪些类别互相混淆多数时候你会发现柠檬被认成了青柠这类细粒度错误那就针对性地补充这类样本。我现在的习惯是每次训练完除了保存best_model.pth一定把class_names.json、预处理参数、image_size写进同一个目录。这两个星期后回来看着文件就能无缝加载推理要是当初只留一个裸权重到时候对着报错猜归一化参数是什么十有八九翻车重来这是我最血泪的一条经验。这套技术方向本身是稳定、可复现的照着上面的步骤把目录和数据准备好你也能交付一个真正能跑、能演示、能延伸的水果识别系统。希望帮到你。本文还有配套的精品资源点击获取
返回列表