首页 / 资讯中心 / 文章详情

基于CNN的垃圾图像分类系统落地实战:从数据到部署

基于CNN的垃圾图像分类系统落地实战:从数据到部署 ★ FEATURED ARTICLE
去年我配合导师带了一个本科毕设小组选的就是“基于卷积神经网络的垃圾图像分类系统”。这个题目乍看简单真正做下来才发现从数据整理、模型调优到最后拼出一个能现场演示的界面坑一个接一个。这篇文章就把这个项目的完整落地过程梳理一遍包括需求怎么拆、CNN到底在学什么、数据怎么准备、模型怎么训练以及最后的系统集成和部署。不管你是正准备开题的本科生还是想快速搭一套图像分类Demo的初学者都可以把这篇当作一份可复用的参考案例。1. 项目背景与需求拆解1.1 为什么选“垃圾图像分类”这个方向很多同学选题时容易走入两个极端要么选一个太抽象的方向比如“基于深度学习的智能系统”听上去高大上但实际上不知道做什么要么选一个过于具体的冷门方向数据都凑不齐。垃圾图像分类刚好处于一个比较舒服的区间第一垃圾分类本身有明确的社会应用场景答辩时可以说清楚“解决什么问题”第二图片数据相对容易获得网上有不少公开的垃圾分类图片集第三图像分类是卷积神经网络最成熟的应用方向之一可参考的资料很多又不至于简单到没有工作量。更重要的是这个题目能够把本科阶段学到的数据结构、Python编程、机器学习基础、软件工程流程全部串起来。学生不是只调一个模型还需要做数据处理、训练评估、界面设计和论文撰写。对毕设来说这种“完整闭环”非常加分。我见过太多学生答辩时只展示一张准确率曲线一问到系统怎么部署就哑火而这个项目天然要求你做一个能跑的Demo所以展示效果会好很多。1.2 从题目到功能需求系统到底要干哪些事拿到题目后第一步不是直接写代码而是把题目翻译成一条条可验收的功能需求。这里很容易犯的错误是功能定义得太大比如“实现智能垃圾分类”结果什么都想做最后什么都做不成。我们的做法是拆成三个核心模块数据模块能够读取本地图片目录自动划分训练集、验证集和测试集支持常见的图片格式。模型模块负责CNN模型的构建、训练和评估训练完成后能够导出权重文件。交互模块提供上传图片-点击识别-显示结果类别置信度的界面并支持单张图片的实时推理。按这个拆法每个模块都有明确的输入输出也方便分工和调试。实际开发中我们还加了两个隐性需求一是训练过程要打印日志并保存损失曲线方便写论文时引用二是模型要能在普通笔记本电脑的CPU上跑通预测因为并不是每个答辩现场都有独立显卡。1.3 技术选型的内在逻辑关于框架选择我们最终用了PyTorch。原因很简单调试方便报错信息直观而且社区里图像分类的教程大多用PyTorch写的。另一个原因是毕业设计时间紧张学生往往只有三四个月PyTorch的“定义即运行”模式对新手更友好。如果用TensorFlow的静态图模式光调试数据管道就可能浪费一两周。模型方面初版使用自己搭建的轻量CNN用于跑通流程后续引入迁移学习用预训练的ResNet18做微调。为什么不直接一步到位上预训练模型因为如果不亲手搭一个简单的CNN学生很难理解卷积、池化、全连接这些层的作用。先把baseline跑通再谈准确率提升这是做深度学习项目最基本的心态。2. 核心原理拆解CNN怎么学会区分垃圾2.1 从像素到特征卷积层到底在提取什么很多学生第一次接触CNN时会把卷积理解成“某种很高级的数学变换”。其实没那么玄。一张图片输入模型后本质是一个三维数组比如224×224×3高和宽对应像素位置3对应R、G、B三个颜色通道。卷积层做的就是拿一个小窗口比如3×3在整张图上滑动每个位置计算窗口内像素和卷积核的加权求和。你可以把这个过程想象成用一组放大镜去照图片有的放大镜专门找横线有的专门找竖线有的专门找颜色突变的地方。一个卷积核学到的权重决定了它“关注”什么模式。在浅层卷积核提取的是边缘、纹理这类低级特征到了深层组合起来就能识别瓶子的轮廓、易拉罐的反光、纸箱的褶皱。这也是为什么网络要做很深——浅层特征太基础直接拿去分类效果很差。2.2 池化、激活函数与全连接一个完整的识别流水线单靠卷积层还不够中间还要穿插池化层和激活函数。池化层的作用是“降采样”比如2×2最大池化取窗口里最大的那个值图片尺寸直接缩小一半。这带来两个好处一是降低计算量二是让模型对物体位置的微小偏移不那么敏感。就好比你站在很远的地方看一个苹果苹果稍微移动几厘米你依然能认出来。激活函数通常用ReLU公式很简单大于0保留原值小于0直接变0。它的作用不是“激活”而是给网络注入非线性。如果没有非线性不管堆多少层卷积本质上还是线性变换根本学不了复杂模式。最后经过一系列卷积和池化后特征图被展平成一维向量送进全连接层全连接层就像是把提取到的所有“证据”汇总起来打分输出每个类别的得分。2.3 为什么是CNN而不是普通全连接网络如果不用卷积神经网络而是直接把图片拉成一维向量丢进全连接网络会发生什么以224×224×3的图片为例输入维度就是150528。第一层如果有256个神经元那这一层的参数量就是150528×256将近3850万个参数。这样的模型在垃圾图片数据集上几乎必然过拟合而且训练速度极慢。CNN通过“权值共享”解决了这个问题。同一个卷积核在图片所有位置使用可以理解为同一套特征检测器被反复复用参数量因此大幅下降。同时卷积的局部连接特性更符合图像本身的规律图像中相邻像素关联强相隔很远的像素之间几乎没有直接关系。所以说CNN不是简单比拼参数数量而是用更聪明的结构去匹配图像数据的天然特点。3. 数据准备与预处理图片质量决定模型上限3.1 数据集的构建与标注格式我们使用的数据来自一个公开的垃圾分类图片集总共约6000张分为六大类纸板、玻璃、金属、塑料、厨余垃圾和其他垃圾。拿到原始图片后第一步不是急着训练而是把数据按类别放进不同的文件夹。文件夹结构就是最简单的标注方式data/ train/ cardboard/ glass/ metal/ plastic/ organic/ other/ val/ ... test/ ...这里有个容易被忽视的细节一定要先检查图片完整性。有些图片可能是损坏的或者实际上是黑白图、带水印、尺寸极小。我们写了一个小脚本把所有图片统一转成RGB三通道并过滤掉宽度或高度小于50像素的图片。这些“脏数据”如果不清理训练时经常会出现莫名其妙的报错或者某个类别的准确率异常低。3.2 图像增强让模型见过“更多的世界”图像增强是提升模型泛化能力最有效的手段之一成本几乎为零。我们用的增强手段包括随机水平翻转、随机旋转±15度、随机亮度对比度调整、随机缩放裁剪。用PyTorch的transforms可以很方便地组合这些操作。这里有一个常见误解增强越多越好。有一年一个学生把旋转角度调到±90度结果把印着“可回收”字样的标签卡片旋转到倒置模型反而学歪了。需要注意增强的“合理性”。垃圾图片在真实场景中可能是任意角度摆放的所以旋转和翻转是合理的但如果你的应用场景是固定的传送带摄像头正对垃圾桶那么过度增强反而会让模型学到不存在的失真。做完增强后还要确认经过transform之后图片的尺寸我们统一缩放到224×224这是后面迁移学习模型的默认输入尺寸。3.3 数据集划分训练/验证/测试的讲究很多初学者直接使用torchvision.datasets.ImageFolder时会默认把数据分成训练和验证两部分但划分方式大有讲究。最关键的一点是必须按类别分层划分不能把所有图片混在一起随机切分。举个例子如果某个类别的图片只有200张纯随机切分时验证集里可能恰好只有10张这个类别的图片训练集却有190张这样评估结果波动会非常大。我们的划分比例是训练集:验证集:测试集 7:2:1并且保证每个类别在三份中的比例基本一致。另外如果同一个物体在多张照片中重复出现比如同一个矿泉水瓶从不同角度拍了5张这5张图最好只放在一个集合里不要一部分在训练集一部分在测试集。否则模型可能靠“记忆”记住这个瓶子测试准确率虚高实际部署时立刻露馅。3.4 标签编码与DataLoader把图片变成模型能吃的张量模型不认文件夹名只认张量。我们要做两件事一是将类别名映射成整数标签二是将PIL图片转换为归一化后的张量。PyTorch的ImageFolder会自动按照文件夹名称的首字母排序生成标签映射比如0: cardboard, 1: glass, ...。为了避免踩坑手动检查一遍映射关系确保和论文里的类别顺序一致。归一化使用的是ImageNet数据集的均值[0.485, 0.456, 0.406]和标准差[0.229, 0.224, 0.225]。如果后续要用预训练模型这个参数必须保持一致因为预训练权重是在这些统计量下学出来的。DataLoader设置batch_size32开启shuffleTrue和num_workers2。在Windows上如果num_workers设得过大有时会报DataLoader worker进程崩溃所以保守一点用2就行。4. 模型训练与调优实战从过拟合到收敛4.1 搭建CNN模型第一版基线架构我们没有直接上重型网络而是先搭建了一个轻量级的CNN用于验证数据管线和训练流程是否正确。网络结构如下卷积层13通道输入16个3×3卷积核输出16224×224接BN和ReLU然后2×2池化。卷积层232个3×3卷积核接BN和ReLU池化。卷积层364个3×3卷积核接BN和ReLU池化。展平后接全连接层64×28×28 → 256 → 6。代码如下import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes6): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 16, kernel_size3, padding1), nn.BatchNorm2d(16), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(16, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), nn.Conv2d(64, 32, kernel_size3, padding1), # 实际应写成64输入详见下文 nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.classifier nn.Linear(64 * 28 * 28, num_classes) def forward(self, x): x self.features(x) x x.view(x.size(0), -1) return self.classifier(x)这里故意留了一个小错误也是很多同学常犯的第三个卷积层的输入通道应该是64而不是32上面代码第二个卷积输出是32通道后面却写了64训练时维度会对不上。实际写代码时一定要逐个核对每层的输入输出维度。基线模型参数量约100万在一张普通显卡上训练很快但在本项目上只能跑到85%左右的准确率已经可以作为后续优化的起点。4.2 训练超参数设置先跑通再调优训练超参数的选择直接影响收敛速度和最终效果。我们采用Adam优化器初始学习率1e-3batch_size为32训练50个epoch。为什么不用SGDSGD对学习率和初始化的要求更高Adam对新手更友好而且在这个规模的数据集上两者差距不大。训练时记录训练集和验证集的损失并画出曲线。一个重要技巧是使用学习率衰减。到了第30个epoch左右损失曲线往往进入平台期此时把学习率降到1e-4可以进一步收敛。我们使用torch.optim.lr_scheduler.StepLR(optimizer, step_size15, gamma0.5)每15个epoch学习率减半。另外训练过程中每个epoch结束后在验证集上评估保存验证损失最低的模型权重而不是最后一个epoch的权重。这个习惯能避免后期模型在验证集上过拟合。4.3 过拟合与迁移学习用预训练模型救场第一版CNN在测试集上大概是82%但训练准确率已经接近95%这说明模型“记住了”训练集没有真正学出泛化规律。针对过拟合我们做了两件事一是增加数据增强强度比如加入随机光照扰动和灰度变换二是引入迁移学习加载预训练模型。我们选用了ResNet18作为baseline因为它的结构相对简单预训练权重容易加载而且参数量只有约1100万适合毕设场景。迁移学习的做法是保留预训练模型的特征提取部分替换最后一层全连接并分两个阶段训练冻结所有特征层参数只训练新分类器学习率1e-3。解冻最后若干层卷积层使用更小的学习率1e-4整体微调。这种方法非常有效在同样的训练集上测试准确率直接提升到93%左右。一个常见误区是解冻所有层然后用高学习率从头微调这会迅速破坏预训练学到的特征导致模型崩掉。正确做法是从小学习率开始并且监控验证集损失一旦回升就停止。4.4 评估指标准确率不是唯一标准毕业设计论文里如果只写“模型准确率达到93%”答辩老师一定会追问各类别的准确率如何是不是平均我们用混淆矩阵和每类Precision、Recall、F1-score来评估。垃圾分类数据集通常类别不均衡比如“其他垃圾”类图片数比其他类多30%此时只看总准确率会产生错觉。以我们的结果为例模型对塑料瓶的召回率达到97%但对玻璃碎片只有88%原因是玻璃和透明塑料在视觉上非常相似。这时候不能单纯追求总准确率而是要根据系统应用场景决定优化目标。如果系统更关注“不能把有害垃圾分错”那么就应该提高有害类别的Recall哪怕牺牲一点其他类别的准确率。这个取舍过程写进论文里是很能体现思考深度的部分。5. 系统集成与成果展示从模型到可用的小系统5.1 模型导出与单张图片推理流程训练完成后模型要能脱离训练环境独立运行。我们保存了完整模型加参数的权重文件但更推荐只保存状态字典torch.save(model.state_dict(), garbage_model.pth)加载时先创建相同的模型结构再load_state_dict。这里有一个特别容易踩的坑如果训练时模型在GPU上保存的是model.state_dict()加载时如果本机没有GPU需要在加载时指定map_locationcpu否则会报错。推理前必须调用model.eval()和with torch.no_grad()否则会占用大量显存且BatchNorm层的行为不一致影响预测结果。单张图片推理流程如下import torch from PIL import Image import torchvision.transforms as transforms def predict(image_path): img Image.open(image_path).convert(RGB) transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) tensor transform(img).unsqueeze(0) with torch.no_grad(): outputs model(tensor) probs torch.softmax(outputs, dim1) conf, cls torch.max(probs, dim1) return class_names[cls.item()], conf.item()注意Resize直接拉伸可能改变物体比例更稳妥的方案是Resize((256, 256)) CenterCrop((224, 224))这样能保留更多有效信息。我们最终使用了带中心裁剪的方式。5.2 GUI/Web界面让非技术人员也能用为了在演示现场让评委感觉“这是个完整系统”我们用轻量级网页框架搭了一个简单的上传界面。核心逻辑只有两部分前端一个文件上传框和结果显示区后端接收图片后调用上面写的predict函数。视频播放演示时学生把一张塑料瓶的图片拖进去界面显示“塑料 置信度0.98”视觉效果比黑乎乎的终端好太多。这里不推荐用tkinter做桌面界面虽然代码简单但打包和跨平台部署都很麻烦。用网页框架的好处是答辩时可以直接在浏览器里跑还可以顺便展示手机拍照上传的扩展空间。界面实现时要注意线程阻塞问题不要让图片推理阻塞网页主线程否则处理一张图片时整个页面卡住会给人很不专业的感觉。简单做法是使用单线程任务队列甚至只用Flask的异步模式也能缓解。5.3 部署常见坑把模型从训练机搬到演示机上的时候我们遇到了至少四个问题。第一个是路径问题训练时代码里写的是相对路径换个目录运行就报错。第二个是模型文件和类别映射不一致训练用的类别顺序和加载后的顺序对不上导致预测结果完全错误。第三个是图片推理时误操作导致图片被旋转最终kindle。第四个是显存问题有人用批量推理时没有手动清空缓存运行几次后GPU内存爆满。对于前三个问题我们的解决办法是建立一个固定的项目目录结构并在代码中使用绝对路径将类别映射保存为JSON文件每次推理时读取同一个文件。第四个问题在CPU上不会遇到但如果有显卡需要在每次推理后用torch.cuda.empty_cache()释放缓存虽然这不是必须的但对长时间的演示系统有帮助。6. 常见问题与避坑实录6.1 训练Loss不降的排查思路训练前几个epoch损失纹丝不动这是几乎每个深度学习项目都会遇到的现象原因多种多样。我们排查的顺序是先看数据加载是否正常——把增强后的图片打印出来确认没有全黑、全白或者标签错位再看学习率是否合适——如果损失在前几个step就变成了NaN通常是学习率太大如果损失下降特别慢可能是学习率太小或者模型初始化有问题最后看标签映射——检查一个batch的标签是否和图片内容一致。一个容易忽略的点是损失函数的选择。多分类问题一般用nn.CrossEntropyLoss这个函数内部已经包含Softmax所以模型输出层不要再额外加Softmax否则训练时梯度会不稳定。如果发现验证集准确率高于训练集准确率多半是训练过程中忘记切换到训练模式或者BatchNorm的设置有问题。6.2 显存不足与训练太慢怎么办毕设阶段不是每个人都有高性能显卡很多人只能用笔记本CPU跑。如果显存不足最简单的办法是把resize从224降到160输入尺寸变小计算量和显存占用都会下降准确率损失通常在2%以内。另一个办法是减小batch_size比如从64降到16。如果CPU上训练实在慢可以先用一个小规模训练集比如每个类别300张验证流程再全量训练。也可以用混合精度训练在PyTorch中使用torch.cuda.amp.autocast和GradScaler能把显存占用降低近一半训练速度也能提升30%。不过这个功能对CPU无效而且如果模型中有自定义层需要注意兼容性。在预算允许的情况下直接用在线GPU平台是最省事的方案但要注意上传数据和训练时间限制。6.3 预测置信度普遍偏低有些时候模型在验证集上准确率不低但实际上传一张真实照片时置信度只有0.5左右。这个问题的根源通常是训练数据太干净了背景单一、物体居中、光线均匀而真实照片里背景杂、角度乱、光线暗。解决办法是增加背景干扰比如把垃圾图片随机粘贴到不同场景背景上或者使用目标检测先裁剪出物体区域再做分类。如果暂时不想做检测也可以在界面提示用户“尽量让物体占据画面中心”。另一个原因是类别分布问题。如果测试图片中的物体不属于训练集的任何一类模型也会给一个“最高置信度”的类别。对于垃圾回收场景故意加入一个“其他/背景”类可以有效降低误判。我们在最后一版中特意加入了100张不放垃圾的背景图片作为额外类别结果系统的误判率下降明显。6.4 写论文时最容易犯的错毕设论文和工程报告不同需要强调“研究思路”和“对比实验”。学生最容易犯的错是直接抄别人的结果或者不写实验环境。我们要求论文里必须包括硬件环境CPU/GPU型号、内存大小、软件环境Python版本、PyTorch版本、数据集来源和划分方式、训练超参数、不同模型对比表。这些细节才是答辩时证明“你真做了”的关键。图表方面除了准确率曲线一定要附上混淆矩阵和部分错误分析比如“哪些图片被分错了错成了什么类”这能体现学生对模型局限的思考。避免使用“提高了20%”这种模糊表述要说明是相对提升了20%还是绝对提升了20%。还有一个常见问题把训练集的损失曲线和验证集的准确率曲线画在一张图里坐标轴不对齐这种图在答辩时会被一眼看穿。这个项目后续还可以扩展的方向有很多比如把分类模型接到摄像头上做成实时检测或者用模型蒸馏把网络压缩到手机上跑。我个人的体会是毕业设计最重要的不是模型分数有多高而是你是否理解每一个环节为什么这么做。A同学最后花了两周时间整理实验记录和写文档这一步帮他拿到了不错的答辩成绩。如果你也在做类似的项目建议从第一天开始就保存中间结果不要等训练全部结束后再回头补。先把baseline跑通再慢慢调优这是我一直坚持的做法。
阅读完成 · 觉得有帮助?
咨询建站