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

基于CNN的垃圾识别分类系统:从数据集到部署的完整实战

基于CNN的垃圾识别分类系统:从数据集到部署的完整实战 ★ FEATURED ARTICLE
简介这份资源是面向高校学生与深度学习入门者的垃圾识别分类课程设计完整项目基于卷积神经网络实现图像分类可直接用于期末大作业或课程设计答辩。压缩包共约2000个文件以1196张jpg与789张jpeg图像构成训练与测试数据集另有13个Python源码文件负责模型搭建、训练与推理并附带json配置与md说明文档整体约564.69MB下载后无需修改即可运行。项目已获导师指导并通过取得97分的高分评价涵盖数据预处理、CNN模型构建、训练调参到分类预测的完整流程目录结构清晰便于按模块阅读与二次开发。目前已有263人学习下载适合希望快速掌握图像分类实战、需要现成数据集与模型参考的读者也可作为进一步优化网络结构与提升识别精度的起点。1. 垃圾识别分类系统从一张照片到四个桶CNN 到底做对了什么你拍一张外卖餐盒的照片系统告诉你「这属于其他垃圾」——听起来简单但背后要跑通一条完整的链路数据采集、图像预处理、CNN 模型搭建、训练调参、推理部署。这套「基于深度学习卷积神经网络实现垃圾识别分类系统」的课程设计核心就是用 Python 把这条链路串起来。它适合正在做课程设计的学生、想入门深度学习 CNN 的开发者以及需要一套可复现图像分类模板的工程师。热搜里「深度学习入门」「cnn卷积神经网络」「python深度学习教程」这些词恰好对应了这套系统涉及的三个层面框架选型、网络结构、训练流程。我见过太多人卡在环境配置或数据集格式上模型还没跑起来就放弃了。这篇笔记按「先立住原理、再动手复现、最后避坑」的顺序展开每一步都给出可抄的代码和参数说明。2. 数据集准备与 CNN 输入管线的搭建2.1 垃圾图像数据集的来源与目录结构常见做法是使用公开的垃圾分类数据集比如 TrashNet 或华为云垃圾分类大赛的数据集通常包含可回收物、厨余垃圾、有害垃圾、其他垃圾四大类部分版本会细分为玻璃、纸张、塑料、金属等子类。我一般会先确认三件事类别数量、每类样本量、图像分辨率。如果某类样本少于 200 张训练时很容易过拟合需要做数据增强。目录结构建议按dataset/train/类别名/图片和dataset/val/类别名/图片组织这样后续用ImageFolder或flow_from_directory可以直接读取不用自己写标签映射。下面是一个典型的目录树dataset/ ├── train/ │ ├── recyclable/ │ ├── kitchen_waste/ │ ├── hazardous/ │ └── other/ └── val/ ├── recyclable/ ├── kitchen_waste/ ├── hazardous/ └── other/划分比例一般按 8:2 或 7:3如果原始数据没有验证集用splitfolders库一行命令切分pip install splitfolders splitfolders --ratio 0.8 0.2 --group-prefix dataset/raw dataset/split--ratio 0.8 0.2表示训练集占 80%、验证集占 20%--group-prefix会按类别子目录自动分组避免随机切分导致某类全落在训练集里。切分后检查一下每个子目录的文件数确认没有空目录。2.2 用 torchvision 构建预处理与增强管线CNN 对输入尺寸和归一化很敏感。常见做法是把图像统一缩放到 224×224ResNet 系列的标准输入再按 ImageNet 的均值和标准差做归一化。训练阶段加入随机翻转、随机裁剪、颜色抖动等增强验证阶段只做缩放和归一化。import torch from torchvision import transforms, datasets from torch.utils.data import DataLoader # 训练集增强管线 train_tf transforms.Compose([ transforms.Resize((256, 256)), # 先放大再裁剪保留更多信息 transforms.RandomResizedCrop(224), # 随机裁剪到 224 transforms.RandomHorizontalFlip(p0.5), # 水平翻转 transforms.ColorJitter(0.2, 0.2, 0.2), # 亮度/对比度/饱和度扰动 transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) # 验证集只做确定性变换 val_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_ds datasets.ImageFolder(dataset/train, transformtrain_tf) val_ds datasets.ImageFolder(dataset/val, transformval_tf) train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4) print(train_ds.classes) # 输出类别名列表RandomResizedCrop(224)的默认缩放范围是 0.08~1.0如果垃圾图像主体占比很小可以改成RandomResizedCrop(224, scale(0.5, 1.0))避免裁掉关键区域。ColorJitter的参数不要设太大0.2 左右足够否则颜色失真会让模型学到错误的纹理特征。num_workers在 Windows 上如果报错改成 0 用主进程加载。注意ImageFolder要求每个类别一个子目录且目录名就是类别标签。如果数据集里混入了非图片文件如 .DS_Store、Thumbs.db加载时会直接报错先清理干净。3. 卷积神经网络选型与迁移学习策略3.1 从零搭一个轻量 CNN 还是直接用预训练模型课程设计里常见的两种路线一是自己堆几层 Conv-BN-ReLU-Pool 从零训练二是用 ResNet18/MobileNetV3 做迁移学习。从零训练的好处是结构透明、便于写进论文但垃圾图像类间差异小比如塑料瓶和玻璃瓶从零训练往往需要更多数据和更长的训练周期。我一般会先用 ResNet18 预训练权重跑一版 baseline再决定要不要自己搭。下面是一个从零搭建的四层 CNN适合理解卷积、池化、全连接的作用import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes4): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), # 224 - 112 nn.Conv2d(32, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2), # 112 - 56 nn.Conv2d(64, 128, 3, padding1), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(2), # 56 - 28 nn.Conv2d(128, 256, 3, padding1), nn.BatchNorm2d(256), nn.ReLU(), nn.AdaptiveAvgPool2d(1) # 全局平均池化 ) self.classifier nn.Linear(256, num_classes) def forward(self, x): x self.features(x) x x.flatten(1) return self.classifier(x)AdaptiveAvgPool2d(1)把任意空间尺寸压成 1×1避免全连接层参数量爆炸。BatchNorm2d放在卷积和激活之间能加速收敛。如果验证集准确率卡在 60% 上不去优先检查数据增强是否过强、学习率是否太大。3.2 迁移学习冻结与微调的参数怎么设用预训练 ResNet18 时常见做法是先把 backbone 冻结只训练最后的全连接层等 loss 稳定后再解冻最后几个 block 做微调。这样能在小数据集上快速拿到一个不错的起点。import torchvision.models as models import torch.nn as nn model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) # 阶段一冻结 backbone for param in model.parameters(): param.requires_grad False # 替换分类头 num_features model.fc.in_features model.fc nn.Sequential( nn.Dropout(0.3), nn.Linear(num_features, 4) ) # 阶段二解冻 layer4 做微调可选 # for param in model.layer4.parameters(): # param.requires_grad Trueweightsmodels.ResNet18_Weights.DEFAULT会自动下载 ImageNet 预训练权重。Dropout(0.3)在全连接前加正则如果训练集很小可以调到 0.5。解冻layer4时学习率要调小一般设为基础学习率的 1/10否则预训练权重会被快速破坏。优化器选 AdamW学习率 1e-3冻结阶段或 1e-4微调阶段权重衰减 1e-4。损失函数用CrossEntropyLoss如果类别不平衡可以加weight参数。import torch.optim as optim criterion nn.CrossEntropyLoss() optimizer optim.AdamW(filter(lambda p: p.requires_grad, model.parameters()), lr1e-3, weight_decay1e-4) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max20)filter(lambda p: p.requires_grad, ...)只把需要更新的参数传给优化器冻结的层不会浪费计算。CosineAnnealingLR让学习率按余弦曲线下降比 StepLR 更平滑。4. 训练循环、评估指标与模型导出4.1 一个可复用的训练与验证循环训练循环要记录 loss 和准确率每个 epoch 结束后在验证集上评估保存最佳模型。下面是一个最小可用的模板import torch from tqdm import tqdm device torch.device(cuda if torch.cuda.is_available() else cpu) model model.to(device) best_acc 0.0 for epoch in range(30): # 训练阶段 model.train() running_loss 0.0 for imgs, labels in tqdm(train_loader, descfEpoch {epoch1} train): imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() outputs model(imgs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * imgs.size(0) scheduler.step() train_loss running_loss / len(train_ds) # 验证阶段 model.eval() correct, total 0, 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels imgs.to(device), labels.to(device) outputs model(imgs) preds outputs.argmax(dim1) correct (preds labels).sum().item() total labels.size(0) val_acc correct / total print(fEpoch {epoch1}: train_loss{train_loss:.4f}, val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print(f - saved best model (acc{best_acc:.4f}))model.train()和model.eval()必须成对出现前者启用 Dropout 和 BatchNorm 的训练模式后者切换到推理模式。torch.no_grad()在验证时关闭梯度计算节省显存。scheduler.step()放在 epoch 结束后调用如果按 batch 更新则放在 batch 循环里。4.2 混淆矩阵与分类报告看清模型到底错在哪准确率只能看整体垃圾识别里更关心「有害垃圾有没有被误判成其他垃圾」。用sklearn的classification_report和confusion_matrix能定位问题类别。from sklearn.metrics import classification_report, confusion_matrix import numpy as np model.eval() all_preds, all_labels [], [] with torch.no_grad(): for imgs, labels in val_loader: imgs imgs.to(device) preds model(imgs).argmax(dim1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) print(classification_report(all_labels, all_preds, target_namesval_ds.classes, digits4)) print(confusion_matrix(all_labels, all_preds))如果某一类的 recall 明显偏低说明该类样本太少或特征不明显优先补充该类数据或调整CrossEntropyLoss的weight参数。混淆矩阵里如果「有害垃圾」大量被预测成「其他垃圾」检查一下两类图像在颜色和形状上是否过于相似。4.3 导出 ONNX 与推理脚本训练完的.pth只能在 PyTorch 环境里用导出 ONNX 后可以跨框架部署也方便集成到 Web 服务或桌面应用里。import torch model.load_state_dict(torch.load(best_model.pth, map_locationcpu)) model.eval() dummy torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy, garbage_cnn.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch}, output: {0: batch}}, opset_version11 )dynamic_axes让 batch 维度可变推理时可以一次传多张图。opset_version11兼容性较好如果部署环境支持更高版本可以调到 13 或 17。导出后用onnxruntime跑一遍验证输出是否一致import onnxruntime as ort import numpy as np sess ort.InferenceSession(garbage_cnn.onnx) dummy_np dummy.numpy() out sess.run(None, {input: dummy_np}) print(out[0].shape) # (1, 4)5. 避坑与排查垃圾识别 CNN 训练中最容易翻车的 5 个点5.1 现象训练 loss 不下降准确率始终在 25% 左右原因学习率过大导致梯度爆炸或者数据标签和目录名没有对上。ImageFolder按目录名排序生成标签如果训练集和验证集的类别顺序不一致验证准确率会随机波动。解决先把学习率降到 1e-4 试一轮确认 loss 有下降趋势。然后打印train_ds.classes和val_ds.classes确认两者完全一致。如果用了自定义 Dataset检查__getitem__返回的 label 是否和类别列表对应。5.2 现象验证集准确率远高于训练集准确率原因验证集太小或分布和训练集差异大也可能是 Dropout 在验证时没关闭。另一种常见情况是验证集里混入了训练集的图片导致「作弊」。解决检查model.eval()是否在验证前调用。用splitfolders重新切分数据确保同一张图不会同时出现在训练集和验证集。如果验证集少于 100 张考虑做交叉验证或扩充验证集。5.3 现象GPU 显存溢出报 CUDA out of memory原因batch_size 太大或者没有用torch.no_grad()导致验证阶段也计算梯度。图像分辨率从 224 提到 448 时显存占用会翻四倍。解决先把 batch_size 降到 16 或 8验证阶段加上with torch.no_grad():。如果还不够用torch.cuda.empty_cache()清理缓存或者把模型换成 MobileNetV3 这类轻量结构。5.4 现象模型在测试图上表现很好但实际拍照识别一塌糊涂原因训练数据是白底商品图实际场景是复杂背景、光照不均、角度倾斜。数据分布不一致是图像分类落地最大的坑。解决在训练集里加入实际场景拍摄的图片至少每类 50 张。增强管线里加入RandomRotation(15)、RandomAffine和更强的ColorJitter。如果条件允许用手机拍一批测试图做一次「真实场景评估」别只看验证集数字。5.5 现象ONNX 导出成功但推理结果和 PyTorch 不一致原因导出时模型没有切换到eval()模式Dropout 和 BatchNorm 仍在训练状态。或者输入数据的预处理方式和训练时不一致比如忘了归一化。解决导出前务必执行model.eval()。推理时确认输入张量的 shape 是(N, 3, 224, 224)且归一化参数和训练时完全相同。用同一张图分别跑 PyTorch 和 ONNX对比输出向量的余弦相似度低于 0.99 就说明有问题。6. 把模型塞进实际系统从单张推理到批量分类的一个技巧训练完模型只是第一步课程设计通常还要求做一个可交互的界面或脚本。我一般会写一个predict.py支持单张图片和整个文件夹的批量推理输出类别名和置信度。这样演示时不用每次都开 Jupyter Notebook。import torch from torchvision import transforms from PIL import Image import sys, os, json # 加载模型结构以 ResNet18 为例 import torchvision.models as models import torch.nn as nn def load_model(ckptbest_model.pth, num_classes4): model models.resnet18(weightsNone) model.fc nn.Sequential(nn.Dropout(0.3), nn.Linear(model.fc.in_features, num_classes)) model.load_state_dict(torch.load(ckpt, map_locationcpu)) model.eval() return model # 与训练一致的预处理 infer_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) CLASSES [hazardous, kitchen_waste, other, recyclable] def predict_image(model, img_path): img Image.open(img_path).convert(RGB) tensor infer_tf(img).unsqueeze(0) # 增加 batch 维度 with torch.no_grad(): logits model(tensor) probs torch.softmax(logits, dim1)[0] idx probs.argmax().item() return CLASSES[idx], probs[idx].item() if __name__ __main__: model load_model() target sys.argv[1] if os.path.isdir(target): results {} for fname in os.listdir(target): if fname.lower().endswith((.jpg, .png, .jpeg)): label, conf predict_image(model, os.path.join(target, fname)) results[fname] {label: label, confidence: round(conf, 4)} print(json.dumps(results, ensure_asciiFalse, indent2)) else: label, conf predict_image(model, target) print(f{target} - {label} ({conf:.2%}))CLASSES的顺序必须和训练时ImageFolder的classes一致否则标签会错位。unsqueeze(0)把单张图的(3, 224, 224)变成(1, 3, 224, 224)因为模型要求输入有 batch 维度。批量推理时用json.dumps输出结构化结果方便后续接入 Web 接口或写入数据库。这个脚本我一般会再包一层argparse加上--topk参数输出前三个类别及概率演示时更有说服力。如果要做成 Web 服务用 FastAPI 包一下predict_image函数接收上传的图片文件返回 JSON 即可。实际部署时注意图片大小限制和并发数CPU 推理单张 224×224 大约 50~100msGPU 上可以做到 10ms 以内。血泪经验是别等到答辩前一天才把模型导出和推理脚本串起来。训练和推理的预处理必须严格一致归一化参数、图像尺寸、通道顺序任何一个对不上结果就是玄学。我习惯在训练脚本里把预处理参数写进一个config.json推理脚本直接读同一个文件省得来回改。希望帮到你。本文还有配套的精品资源点击获取
阅读完成 · 觉得有帮助?
咨询建站