基于CNN的6分类垃圾分类实战：从数据准备到迁移学习

简介这份资源面向计算机、人工智能相关专业的毕业设计学生与深度学习入门者提供一套基于Python与卷积神经网络实现六类别垃圾分类的完整项目方案。数据集涵盖glass、cardboard、metal、paper、plastic、trash六类常见生活垃圾帮助读者在真实分类任务中理解CNN的卷积、池化与全连接结构并完成模型搭建、训练与调参的完整流程。压缩包共28个文件约16.71MB包含5个py脚本用于模型定义与训练1个ipynb交互式笔记本便于逐步调试另有14张png图片记录模型结构与学习曲线以及md说明、docx文档和h5权重文件等目录结构清晰便于按模块查阅。目前已有4456人学习下载。读者可从中获得可运行的代码、六分类数据集处理思路、模型图与学习曲线分析示例以及调整学习率等超参数对结果影响的对比实验适合作为毕业设计参考或CNN实战练手项目。1. 从一张垃圾桶照片说起6 分类 CNN 垃圾分类到底在做什么你拍一张外卖盒的照片模型要在 6 个类别里给出唯一答案可回收物、有害垃圾、厨余垃圾、其他垃圾再加上玻璃、纸类这种更细的二级类。听起来像图像分类的入门作业但真正动手做毕业设计时翻车点往往不在网络结构而在数据分布和类别定义上。我见过太多同学把 ResNet 一贴、准确率跑到 95% 就交差结果答辩老师随手拿一张带油渍的塑料盒模型直接判成厨余垃圾。这个方向适合两类人一类是计算机、电子信息、物联网专业的毕业设计选题需要一套能跑通、能写论文、能演示的完整方案另一类是想入门深度学习 CNN 的工程师拿垃圾分类当第一个端到端项目练手。它覆盖了 python 环境配置、卷积神经网络原理、数据增强、迁移学习、模型评估这一整条链路比 MNIST 手写数字有说服力得多。我一般会把整个方案拆成四块数据准备、模型搭建、训练调参、推理部署。下面按这个顺序讲每一步都给能直接抄的代码和参数说明。你不需要 GPU 也能跑CPU 上把图片尺寸压到 128×128、batch 设 16一个 epoch 大概两三分钟。2. 数据准备与 6 类别定义别让标签毁掉整个模型2.1 6 个类别怎么定直接决定模型上限垃圾分类的类别定义是整个项目的地基。常见做法是沿用四分类国标再拆出两个高频子类凑成 6 类。我一般这样定类别编号类别名典型样本易混对象0可回收物-塑料饮料瓶、塑料盒其他垃圾里的污染塑料1可回收物-纸类纸箱、报纸厨余里的纸巾2有害垃圾电池、药品、灯管其他垃圾里的化妆品瓶3厨余垃圾果皮、剩饭、菜叶其他垃圾里的骨头4其他垃圾烟头、陶瓷、尿不湿可回收物里的脏塑料5玻璃/金属玻璃瓶、易拉罐其他垃圾里的碎玻璃这个划分的好处是每类都有足够样本且边界相对清晰。注意纸巾不管多干净都算其他垃圾大骨头算其他垃圾这两个是答辩高频追问点提前在论文里写清楚。数据集来源上公开的垃圾分类数据集常见有 4 万张左右、40 个细类的版本你需要自己合并映射成 6 类。我一般写一个映射字典把细类名归到大类避免手工拖文件夹。2.2 用 python 把原始数据整理成 train/val/test假设你下载的原始数据是raw_data/下按细类分文件夹先做类别映射和划分。下面这段脚本直接抄import os import shutil import random from pathlib import Path # 细类到大类的映射按你实际数据集的名字改 LABEL_MAP { plastic_bottle: recyclable_plastic, plastic_bag: recyclable_plastic, paper_box: recyclable_paper, newspaper: recyclable_paper, battery: hazardous, medicine: hazardous, fruit_peel: kitchen, leftover: kitchen, cigarette: other, ceramic: other, glass_bottle: glass_metal, can: glass_metal, } RAW_DIR Path(raw_data) OUT_DIR Path(dataset) SPLIT {train: 0.7, val: 0.15, test: 0.15} random.seed(42) for split in SPLIT: for cls in set(LABEL_MAP.values()): (OUT_DIR / split / cls).mkdir(parentsTrue, exist_okTrue) for sub_dir in RAW_DIR.iterdir(): if not sub_dir.is_dir(): continue big_cls LABEL_MAP.get(sub_dir.name) if big_cls is None: continue imgs [p for p in sub_dir.iterdir() if p.suffix.lower() in (.jpg, .png, .jpeg)] random.shuffle(imgs) n len(imgs) n_train int(n * SPLIT[train]) n_val int(n * SPLIT[val]) buckets { train: imgs[:n_train], val: imgs[n_train:n_train n_val], test: imgs[n_train n_val:], } for split, files in buckets.items(): for f in files: shutil.copy(f, OUT_DIR / split / big_cls / f.name) print(done)逻辑说明先按映射字典把细类归并再对每个细类内部随机打乱后按 7:1.5:1.5 切分。参数上random.seed(42)保证可复现答辩时老师让你重跑结果一致。SPLIT可以按数据集大小调样本少于 500 张的类建议把 train 提到 0.8。注意一个坑不要先合并再切分。如果先把所有塑料瓶混在一起再随机切同一张图的近邻帧可能同时进 train 和 test导致测试准确率虚高。按细类切分能缓解这个问题。2.3 用 ImageFolder 和 DataLoader 搭数据管道整理完目录结构是dataset/train/类别名/图片直接用 torchvision 的 ImageFolderimport torch from torchvision import datasets, transforms from torch.utils.data import DataLoader IMG_SIZE 128 BATCH_SIZE 16 train_tf transforms.Compose([ transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.2, contrast0.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((IMG_SIZE, IMG_SIZE)), 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) test_ds datasets.ImageFolder(dataset/test, transformval_tf) train_loader DataLoader(train_ds, batch_sizeBATCH_SIZE, shuffleTrue, num_workers2) val_loader DataLoader(val_ds, batch_sizeBATCH_SIZE, shuffleFalse, num_workers2) test_loader DataLoader(test_ds, batch_sizeBATCH_SIZE, shuffleFalse, num_workers2) print(train_ds.classes)参数说明IMG_SIZE128是 CPU 能接受的折中有 GPU 可以上 224。Normalize用的是 ImageNet 统计量因为后面要用预训练权重这一步不能省。RandomRotation(15)对垃圾分类有用因为拍摄角度多变但RandomVerticalFlip不要加垃圾不会倒过来拍加了反而引入噪声。num_workers在 Windows 上设 0 更稳Linux 设 2 到 4。3. 搭一个能打的 CNN从三层卷积到迁移学习3.1 手写 CNN 的结构和每层参数怎么定毕业设计论文里通常要求你自己搭一个网络不能只调库。我一般给一个 4 层卷积的基线import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes6): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), # 128 - 64 nn.Conv2d(32, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2), # 64 - 32 nn.Conv2d(64, 128, 3, padding1), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(2), # 32 - 16 nn.Conv2d(128, 256, 3, padding1), nn.BatchNorm2d(256), nn.ReLU(), nn.AdaptiveAvgPool2d(1), # 任意尺寸 - 1x1 ) self.classifier nn.Sequential( nn.Flatten(), nn.Dropout(0.5), nn.Linear(256, num_classes), ) def forward(self, x): return self.classifier(self.features(x))逻辑说明每个卷积块是 Conv-BN-ReLU-Pool 的标准组合。padding1配合 3×3 卷积保证特征图尺寸不变尺寸只被池化层缩小。BatchNorm2d放在卷积后、激活前能明显加快收敛这是很多人漏掉的一步。最后用AdaptiveAvgPool2d(1)代替全连接展平好处是输入尺寸变了也不用改结构。参数上通道数 32-64-128-256 是逐层翻倍的经典做法显存不够就整体减半。Dropout(0.5)只加在全连接前卷积层里不加因为卷积本身参数共享过拟合风险低。3.2 训练循环损失、优化器和学习率import torch from torch import optim device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleCNN(num_classes6).to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3, weight_decay1e-4) scheduler optim.lr_scheduler.StepLR(optimizer, step_size10, gamma0.5) EPOCHS 30 best_acc 0.0 for epoch in range(EPOCHS): model.train() total_loss 0.0 for imgs, labels in train_loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() out model(imgs) loss criterion(out, labels) loss.backward() optimizer.step() total_loss loss.item() model.eval() correct, total 0, 0 with torch.no_grad(): for imgs, labels in val_loader: imgs, labels imgs.to(device), labels.to(device) pred model(imgs).argmax(dim1) correct (pred labels).sum().item() total labels.size(0) acc correct / total scheduler.step() print(fepoch {epoch1} loss {total_loss/len(train_loader):.4f} val_acc {acc:.4f}) if acc best_acc: best_acc acc torch.save(model.state_dict(), best_cnn.pth)参数说明Adam配lr1e-3是安全起点weight_decay1e-4做 L2 正则。StepLR每 10 个 epoch 把学习率砍半防止后期震荡。best_acc只在验证集上更新测试集全程不碰这是论文里必须写清楚的原则。失败时看什么如果 loss 一直不降先检查标签有没有对错位、Normalize 有没有漏如果 train_acc 高但 val_acc 低说明过拟合加数据增强或加大 Dropout如果两个都低把学习率降到 1e-4 再试。3.3 迁移学习小数据集上把准确率拉起来自己搭的网络在几千张图上通常到 85% 左右就卡住。想上 92% 以上用预训练模型微调是最省事的路子import torchvision.models as models import torch.nn as nn model models.resnet18(weightsmodels.ResNet18_Weights.DEFAULT) for param in model.parameters(): param.requires_grad False # 先冻结主干 model.fc nn.Linear(model.fc.in_features, 6) # 换掉最后的分类头 model model.to(device) # 只训练 fc 层 5 个 epoch再解冻最后两个 block 微调 optimizer optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr1e-3)逻辑说明先冻结主干只训分类头让随机初始化的 fc 快速收敛避免一开始就把预训练权重带偏。5 个 epoch 后把layer4解冻学习率降到 1e-4 继续训。这套两段式微调在小数据集上比直接全量微调稳得多。注意weightsmodels.ResNet18_Weights.DEFAULT会自动下载权重第一次跑需要联网。如果环境不能联网提前把权重文件放到~/.cache/torch/hub/checkpoints/下。4. 避坑与排查那些让准确率虚高或直接崩掉的问题4.1 测试集准确率 99% 但演示翻车现象验证集和测试集都跑到 98% 以上答辩现场拿手机拍一张真实垃圾模型乱判。 原因数据集里的图片大多是白底商品图和真实拍摄的光照、背景、角度差异巨大模型学到的是背景而不是物体。 解决训练时加更强的数据增强包括随机裁剪、背景替换、亮度对比度扰动另外自己用手机补拍 50 到 100 张真实场景图作为额外的测试集专门用来暴露这个问题。论文里把这个差距写出来反而是加分项。4.2 某个类别准确率永远是 0现象6 个类里有一类 recall 一直是 0混淆矩阵里全被分到隔壁类。 原因类别样本严重不均衡或者两个类的视觉特征太接近比如塑料盒和纸盒都是白色方形。 解决先看这一类有多少张图少于 300 张就做过采样或加类权重CrossEntropyLoss(weight...)。如果样本够但还是分不开检查标签映射有没有把两个类混到一起这种低级错误我见过不止一次。4.3 训练到一半 loss 变成 nan现象前几个 epoch 正常突然 loss 变 nan准确率归零。 原因学习率太大导致梯度爆炸或者数据里有损坏图片产生异常值。 解决先把学习率降到 1e-4加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm5.0)。然后写个脚本遍历所有图片用 PIL 打开一遍把报错的图删掉。损坏图片是黑匣子不主动查根本发现不了。4.4 CPU 训练慢到怀疑人生现象一个 epoch 要跑半小时30 个 epoch 根本等不起。 原因图片尺寸太大、batch 太小、num_workers 设成 0 导致数据加载串行。 解决把IMG_SIZE降到 96 或 64BATCH_SIZE提到 32num_workers设成 CPU 核数。另外把模型换成resnet18而不是自己搭的深网络预训练模型收敛快10 个 epoch 就能出结果。实在不行用 Google Colab 的免费 GPU但注意数据要提前传上去。4.5 保存的模型加载后预测结果全乱现象训练时准确率 90%重新加载best_cnn.pth后预测全是同一类。 原因保存的是state_dict加载时必须先实例化同样的网络结构类别顺序也要一致。ImageFolder 是按文件夹名首字母排序的换台机器如果文件夹名大小写不同顺序就变了。 解决把train_ds.classes存成classes.json推理时读这个文件做索引映射不要依赖 ImageFolder 的默认排序。5. 推理脚本与论文里能加分的评估细节训练完只是半成品答辩要现场演示你得有一个干净的推理入口。下面这个脚本读单张图输出类别和置信度import json import torch from PIL import Image from torchvision import transforms with open(classes.json, r, encodingutf-8) as f: classes json.load(f) model SimpleCNN(num_classes6) model.load_state_dict(torch.load(best_cnn.pth, map_locationcpu)) model.eval() tf transforms.Compose([ transforms.Resize((128, 128)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) img Image.open(test.jpg).convert(RGB) x tf(img).unsqueeze(0) with torch.no_grad(): prob torch.softmax(model(x), dim1)[0] idx prob.argmax().item() print(classes[idx], f{prob[idx].item():.4f})逻辑说明unsqueeze(0)把单张图变成 batch1softmax把 logits 转成概率。map_locationcpu保证在没 GPU 的机器上也能加载。置信度低于 0.6 时建议在演示界面提示「不确定」比硬给一个错答案体面。论文评估部分除了准确率把混淆矩阵和每类的 precision、recall、F1 都算出来。用 sklearn 三行搞定from sklearn.metrics import classification_report, confusion_matrix print(classification_report(all_labels, all_preds, target_namesclasses)) print(confusion_matrix(all_labels, all_preds))老师最爱问的就是「哪两类最容易混」混淆矩阵直接回答这个问题。我一般还会画一张错分样本图把判错的图贴出来分析原因这一页在答辩时比准确率数字更有说服力。最后一个习惯每次改完超参把配置写进一个config.yaml连同classes.json和best_cnn.pth一起归档。我吃过亏隔两周回头跑忘了当时用的学习率和图片尺寸只能重训一遍。这个方案从数据整理到推理大概两三天能跑通值不值得做取决于你要的是论文素材还是真实可用的分类器前者一周够后者得在数据上多花功夫。希望帮到你。本文还有配套的精品资源点击获取