简介一套基于Python CNN卷积神经网络实现六类别垃圾分类的毕业设计完整资料面向计算机与人工智能相关专业学生尤其适合需要完成课程设计、毕业设计或入门图像识别实战的初学者。资源围绕glass、cardboard、metal、paper、plastic、trash六类生活垃圾图像覆盖数据加载、模型搭建、训练评估与调参分析完整链路。压缩包共28个文件大小约16.71MB以多个Python模型脚本mynet.py、resnet.py、kerasmodel.py等、预训练权重res.h5和Jupyter Notebook为主另有README说明文档、config.py配置脚本、网络结构示意图及训练结果可视化图片便于边读代码边复现。压缩包内还包含可直接加载的权重文件和测试图片方便快速验证预测效果。目前已有4458人学习浏览。读者可以拿到可直接运行的CNN分类模型对比不同网络结构和学习率等超参数的影响根据TensorBoard日志和学习曲线分析训练过程快速搭建并理解垃圾分类项目。1. 基于 python CNN 卷积神经网络模型实现6类别垃圾分类毕设到底难在哪我最近帮一个师弟收拾毕业设计题目是“基于 python CNN 卷积神经网络模型实现6类别垃圾分类”听起来是个典型的深度学习图像分类项目但真正上手后才发现最难的不是搭一个卷积神经网络而是把数据管好、把训练过程调到不翻车。这个项目要解决的事情很具体输入一张垃圾照片输出它属于可回收物、有害垃圾、厨余垃圾等 6 类中的哪一类。对计算机、人工智能方向的毕设来说它的价值在于既有完整的 CNN 原理能讲清楚又能在普通笔记本上用 CPU 跑完成本低、演示效果好。这篇内容不打算重现什么官方教程就按我实际带项目的过程把标签体系、数据增强、模型结构、训练配参和五个高频踩坑点全部摆出来新手可以照着代码一步步跑。2. 先把数据管好6 类别垃圾分类的标签体系与预处理脚本2.1 六类标签体系怎么定才不被 ImageFolder 坑做 CNN 图像分类第一步不是急着写网络而是先定 6 个类别。垃圾分类在各城市标准不完全一样毕设里我一般定成纸类、塑料、金属、玻璃、有害垃圾、厨余垃圾。目录结构直接按英文做garbage6/ ├── paper/ ├── plastic/ ├── metal/ ├── glass/ ├── hazardous/ └── kitchen_waste/每个文件夹的名字会被torchvision.datasets.ImageFolder直接当作类别标签类别的 index 按目录名字母顺序排列所以paper0, plastic1, metal2, glass3, hazardous4, kitchen_waste5。很多人后面画混淆矩阵时想当然按自己认为的顺序填标签结果矩阵对角线不在一条线上排查了半天才发现是类别顺序错位。这步看起来基础但能避免后面所有指标对不上。数据量方面6 类每类最好不低于 300 张再多不嫌多少了才麻烦。我会把原始图片按 8:1:1 分成 train、val、test 三个子集测试集从训练开始就冻结绝不参与调参否则答辩时老师拿新图一测就暴露了。下面的脚本完成拆分并把随机种子固定住保证每次实验结果可复现# split_data.py import os import random import shutil random.seed(42) source garbage6 target garbage6_split ratio {train: 0.8, val: 0.1, test: 0.1} for class_name in os.listdir(source): class_dir os.path.join(source, class_name) if not os.path.isdir(class_dir): continue images os.listdir(class_dir) images [img for img in images if img.lower().endswith((.jpg, .jpeg, .png))] random.shuffle(images) total len(images) n_train int(total * ratio[train]) n_val int(total * ratio[val]) subsets { train: images[:n_train], val: images[n_train:n_train n_val], test: images[n_train n_val:], } for split_name in subsets: out_dir os.path.join(target, split_name, class_name) os.makedirs(out_dir, exist_okTrue) for img_name in subsets[split_name]: shutil.copy( os.path.join(class_dir, img_name), os.path.join(out_dir, img_name) ) print(split finished)逻辑说明先过滤掉非图片文件避免.txt或.DS_Store混进数据集导致读取报错。用copy而不是move是为了保留原始数据万一后面拆分比例不对还能重来。参数说明比例可以按数据量调整如果样本总数很少train 提高到 0.9val 和 test 各留 0.05但这样评估结果波动大最好还是去多采集一些图片。2.2 用 python 写数据集加载与增强脚本分类模型不挑图片是 jpg 还是 png但垃圾照片往往在光线、背景、拍摄角度上差异很大训练时不做数据增强模型很容易把某种背景当特征。下面是最小可用版本的 dataset.py# dataset.py from torchvision import datasets, transforms from torch.utils.data import DataLoader from PIL import Image train_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.Lambda(lambda img: img.convert(RGB)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) val_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.Lambda(lambda img: img.convert(RGB)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) train_data datasets.ImageFolder(garbage6_split/train, transformtrain_transforms) val_data datasets.ImageFolder(garbage6_split/val, transformval_transforms) train_loader DataLoader(train_data, batch_size32, shuffleTrue, num_workers0) val_loader DataLoader(val_data, batch_size32, shuffleFalse, num_workers0)逻辑说明Resize((224, 224))先把所有图统一成正方形这是后面 CNN 输入张量的固定形状。Lambda(lambda img: img.convert(RGB))把 RGBA 四通道图片统一转成 RGB 三通道避免读图片时通道数不一致。RandomHorizontalFlip和RandomRotation(15)是训练集专用的增强让模型看到更丰富的姿态验证集不用随机增强否则每次评估结果都会抖。ToTensor()把 PIL 数组从 HWC 变成 CHW并除以 255 缩到 0-1后面的Normalize再用 ImageNet 统计量做标准化。参数说明num_workers0在 Windows 上最省心不会因为多进程启动方式卡死batch_size32适合显存 4G 以上的环境如果 CPU 内存紧张就降到 16。2.3 预处理参数怎么调尺寸、归一化、batch size 的边界Resize 不一定非要 224用 128 会快很多但小垃圾上的边缘纹理容易丢用 256 会占用更多内存对毕设这种规模来说性价比不高。我一般固定在 224因为后面换预训练模型时ResNet 等模型默认也是这个尺寸不需要二次改造。Normalize 的均值和标准差不需要自己统计直接沿用 ImageNet 的[0.485, 0.456, 0.406]就没有问题。有人喜欢“自己算一遍数据集的 mean/std”算出来往往和 ImageNet 差不多但浪费时间除非图片整体色调特别统一比如全是暗光照片才值得自己统计。batch size 的影响更直接32 在几百张图的小数据集上可能让梯度方向偏陡我会先观察 loss 是否上下波动。如果训练 loss 跳得很厉害就把 batch size 减半到 16这时候梯度方向更平滑如果减到 16 后每个 epoch 时间翻倍就用增大数据增强来补偿。另外一定记住预处理顺序增强 →ToTensor→Normalize。我见过有人写反先Normalize再ToTensor输入还是 0-255 的原始数模型直接不收敛。这个小坑后面还会单独拿出来说。提示数据管好后先打印一个 batch 的images.min()和images.max()确认值在 0-1 之间再往下走。这个习惯能省掉后面大量排查时间。3. 构建 CNN 卷积神经网络模型从结构到最小可跑通的 PyTorch 代码3.1 为什么我选 PyTorch 而不是 TensorFlow搬这个项目时我第一时间就锁定 PyTorch。对 CNN 图片分类PyTorch 的动态图机制让打印每层输出形状、打断点调试都变得非常自然torchvision自带ImageFolder、预训练模型和常用 transform省掉大量造轮子的时间。TensorFlow 当然也能做但 Keras 的高封装在“只改某一层输出”这类需求上反而要绕弯对新手不够友好。毕设答辩更看重你对原理的讲解PyTorch 的代码结构和论文里的conv - relu - pool几乎一一对应讲起来很顺。另外现在深度学习相关的环境配置PyTorch 的安装和报错信息在社区里讨论最多遇到问题搜起来也快。3.2 一个适合 CPU 跑的三层卷积结构垃圾分类图片没有 ImageNet 那么细的粒度三层卷积足够提取颜色、纹理、边缘这类视觉特征。我常用的自建网络长这样# model.py import torch.nn as nn class GarbageCNN(nn.Module): def __init__(self, num_classes6): super(GarbageCNN, self).__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(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), ) self.avgpool nn.AdaptiveAvgPool2d((1, 1)) self.classifier nn.Sequential( nn.Dropout(p0.5), nn.Linear(64, 128), nn.ReLU(inplaceTrue), nn.Dropout(p0.3), nn.Linear(128, num_classes) ) def forward(self, x): x self.features(x) x self.avgpool(x) x x.view(x.size(0), -1) return self.classifier(x)逻辑说明输入是3x224x224第一层卷积核数量从 16 开始后面每层翻倍这是 CNN 的常见做法通道数增加、空间尺寸减小让网络能表达更丰富的语义特征。每个卷积后面接BatchNorm2d它能把中间特征分布拉回稳定区间训练时不那么玄学也能支持更大的学习率。MaxPool2d(2)负责把特征图长宽各减半三次池化后 224 变成 28但后面接了一个AdaptiveAvgPool2d((1,1))所以不用手工算最终尺寸它会把任意空间大小压成 1x1全连接层输入维度永远是64。Dropout(0.5)对几千张样本的小数据来说非常关键能明显缓解过拟合。参数说明kernel_size3, padding1保证卷积不会改变图片长宽如果想提高容量可以把通道数调成32/64/128但 CPU 训练时间会明显变长。这个模型的参数量只有几百万普通笔记本 CPU 一个 epoch 大概几分钟非常适合毕设。如果你有 CUDA 显卡可以在train.py里自动选用 GPU训练会快很多。3.3 训练主循环损失函数、优化器与保存最优权重分类任务用CrossEntropyLoss它把输出层 logits 转成概率再算损失不需要自己在最后一层加 softmax。优化器先用 Adam学习率 0.001 起步。下面这个 train.py 可以直接抄# train.py import torch import torch.nn as nn import torch.optim as optim from model import GarbageCNN from dataset import train_loader, val_loader device torch.device(cuda if torch.cuda.is_available() else cpu) model GarbageCNN(num_classes6).to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) scheduler optim.lr_scheduler.StepLR(optimizer, step_size8, gamma0.5) epochs 40 best_acc 0.0 for epoch in range(epochs): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) scheduler.step() model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() val_acc correct / total if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_model.pth) print(fepoch {epoch1}/{epochs}, loss {running_loss / len(train_loader.dataset):.4f}, val_acc {val_acc:.4f})逻辑说明训练阶段必须用model.train()这样 BatchNorm 和 Dropout 才会按训练行为工作验证阶段用model.eval()Dropout 关闭、BatchNorm 用累计统计量结果才稳定。torch.no_grad()让验证时不计算梯度省显存也加快速度。参数说明StepLR每 8 个 epoch 将学习率乘 0.5防止后期 loss 震荡torch.save只保存权重后面你想换一个网络结构重新加载也能兼容。epochs40对这个小网络够用如果到 40 轮验证准确率还在上升就把epochs加到 60。训练过程中 loss 持续下降但验证准确率上不去基本就是过拟合或学习率衰减时机不对下一章集中处理。4. 训练配参、评估指标与可视化让垃圾分类 CNN 真正收敛4.1 学习率与优化器调节别迷信默认值Adam 默认lr0.001在大多数数据集上都能跑但垃圾图片的表现差异很大。如果图片干净、背景单一0.001 可能在后期震荡如果光照差异大0.001 又太低第一个 epoch 结束 loss 还在 3.0 以上。我一般看第一个 epoch 的 loss 判断如果 loss 低于 2.0说明学习率合适如果还在 3.0 以上把 lr 调到 0.01 试一轮如果 loss 飞出了 NaN马上降到 0.0001。手动调学习率太累我更推荐用ReduceLROnPlateau直接监控验证准确率卡住就降学习率from torch.optim import lr_scheduler scheduler lr_scheduler.ReduceLROnPlateau( optimizer, modemax, factor0.5, patience5 ) # 在每个 epoch 的验证集准确率算出来后调用 # scheduler.step(val_acc)逻辑说明modemax表示监控的指标是越大越好这里对应验证准确率patience5允许连续 5 个 epoch 准确率不创新高后才降学习率避免刚进入平台期就降得太频繁。factor0.5是衰减比例。使用它之后原来的StepLR就去掉保存最优模型的逻辑保留。4.2 输出准确率和混淆矩阵答辩要的硬指标训练完只打印 val_acc 不够老师更想知道每个类到底分得好不好。下面用 sklearn 和 matplotlib 生成混淆矩阵# evaluate.py import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader import matplotlib.pyplot as plt import seaborn as sns from sklearn.metrics import accuracy_score, confusion_matrix from model import GarbageCNN device torch.device(cuda if torch.cuda.is_available() else cpu) model GarbageCNN(num_classes6).to(device) model.load_state_dict(torch.load(best_model.pth, map_locationdevice)) model.eval() transform transforms.Compose([ transforms.Resize((224, 224)), transforms.Lambda(lambda img: img.convert(RGB)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) test_data datasets.ImageFolder(garbage6_split/test, transformtransform) test_loader DataLoader(test_data, batch_size1, shuffleFalse) y_true, y_pred [], [] with torch.no_grad(): for images, labels in test_loader: images images.to(device) outputs model(images) _, preds torch.max(outputs, 1) y_true.extend(labels.tolist()) y_pred.extend(preds.cpu().tolist()) acc accuracy_score(y_true, y_pred) cm confusion_matrix(y_true, y_pred) class_names test_data.classes print(test accuracy:, acc) print(class order:, class_names) plt.figure(figsize(8, 6)) sns.heatmap(cm, annotTrue, fmtd, xticklabelsclass_names, yticklabelsclass_names) plt.xlabel(Predicted) plt.ylabel(True) plt.tight_layout() plt.savefig(confusion_matrix.png)逻辑说明map_locationdevice让模型权重在无 CUDA 的机器上也能加载。batch_size1虽然慢但方便逐张追踪错误样本数据量不大时完全可以接受。混淆矩阵对角线越亮越好如果“有害垃圾”预测成“厨余垃圾”的格子很亮说明这两个类在视觉上本来就接近需要补数据或提高类别权重。class_names来自test_data.classes保证标签顺序和训练时一致。参数说明annotTrue把具体数值标在格子里fmtd表示显示整数而不是科学计数法。4.3 过拟合判断与早停别让训练白跑垃圾图片数据集通常只有几百到几千张模型很容易把训练集里的背景纹理背下来。判断过拟合的经典信号是训练 loss 继续下降验证 loss 在第 10 轮开始反弹。除了加 Dropout、做增强最直接的方法是早停并保存最佳权重。前面的训练循环只在 val_acc 变好时保存模型这就是最简单的早停。更稳的做法是监控验证 loss连续多个 epoch 不下降就停patience 7 wait 0 best_val_loss float(inf) for epoch in range(epochs): # 前面的训练循环 model.eval() val_loss 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() * labels.size(0) val_loss / len(val_loader.dataset) if val_loss best_val_loss: best_val_loss val_loss wait 0 torch.save(model.state_dict(), best_model.pth) else: wait 1 if wait patience: print(fearly stop at epoch {epoch1}) break逻辑说明验证 loss 对过拟合比准确率更敏感模型在准确率还没掉的时候验证 loss 可能已经开始回升。wait变量记录“已等待多少个 epoch 没有创新低”超过patience就中断。参数patience7适合数据量小的情况如果数据量大这个值可以降到 5避免等太久。如果把损失曲线画出来横轴是 epoch纵轴是 train/val loss答辩时放这张图比只说准确率更有说服力。画图代码就是普通的 matplotlib 两条折线注意先收集train_losses和val_losses两个 list。5. 避坑/常见问题/排查CNN 垃圾分类里最容易翻车的 5 个问题5.1 类别不平衡导致准确率虚高现象训练集里“塑料”类有 3000 张“有害垃圾”只有 200 张训练完 val_acc 达到 92%但看混淆矩阵有害垃圾大量被预测成塑料。原因交叉熵默认每个类别的权重相同多数类贡献了绝大部分梯度决策边界会偏袒多数类。准确率虚高是因为即使在有害垃圾全错的情况下只要塑料类预测得足够好整体准确率依然很高。解决用加权采样器让少量类在 batch 里出现得更频繁。代码改动很小from torch.utils.data import WeightedRandomSampler # 按 train_data.classes 的顺序统计每个类的图片数量 class_counts [len(train_data.targets[train_data.targets i]) for i in range(6)] weights [1.0 / count for count in class_counts] sample_weights [weights[label] for label in train_data.targets] sampler WeightedRandomSampler( sample_weights, num_sampleslen(sample_weights), replacementTrue ) train_loader DataLoader( train_data, batch_size32, samplersampler, num_workers0 )逻辑说明样本少的类每个样本被抽中的概率更高这样每个 batch 里小类数量会被拉上来。replacementTrue允许重复采样保证均衡效果。注意加了sampler后DataLoader不能再写shuffleTrue否则程序报错。5.2 图片读取路径与通道编码问题现象代码在 Linux 上跑得好好的放到 Windows 上一读到某些图片就报错或者目录路径带中文程序直接找不到文件。原因部分手机或相机拍出来的图是 RGBA 四通道有些还是 16 位深度的 PNGPIL默认读取方式可能返回 4 通道Windows 中文用户目录还会让路径编码不统一。解决在 transform 里强制转 RGB并把数据集放到英文、无空格的路径下。transform 里加一行transforms.Lambda(lambda img: img.convert(RGB))然后在ImageFolder加载前先用 Python 的os.path.exists()检查路径确认没有中文。如果图片出现损坏用try-except在加载时跳过坏图for img_name in images: try: Image.open(os.path.join(class_dir, img_name)).load() except (OSError, Image.DecompressionBombError): continue这能避免某张损坏图让整个 epoch crash。5.3 归一化顺序写反导致 loss 不收敛现象训练第一轮 loss 直接变成 NaN或者准确率卡在 20% 左右怎么调学习率都没用。原因把Normalize放在ToTensor之前此时输入还是 0-255 的整数数组Normalize 的均值假设输入是 0-1这样标准化等于把整张图随机放缩模型看到的数据分布完全错乱。解决严格遵守顺序增强 →ToTensor→Normalize。写完 transform 后手动从 DataLoader 里取一个 batch 打印一下images.min()和images.max()确认在 0-1 之间。我还会打印前几个像素值看起来是正负小数就说明标准化生效了。5.4 Windows 上 DataLoader 卡死现象设置num_workers4后程序启动就卡在第一个 batchCPU 占用很高但图片加载不出来。原因PyTorch 在 Windows 上的多进程 DataLoader 需要if __name__ __main__保护否则子进程无法正确导入模块会进入假死状态。这在 Linux 上不会发生所以很容易被忽略。解决把训练主逻辑包进if __name__ __main__:或者直接设置num_workers0。毕设数据量不大num_workers0的耗时差异可以忽略但稳定性提升很明显。5.5 单独预测新图时忘记做同样的预处理现象测试集准确率 95%但自己拍一张垃圾图片拿去预测输出结果永远是同一类或者 softmax 概率几乎全是 0.5。原因训练时的 transform 包含了Resize、ToTensor、Normalize而单独预测时只做了Resize和ToTensor少了标准化输入分布和训练时不一致或者图片尺寸不是 224模型全连接层接不上。解决写一个单张图片预测函数强制走和验证集完全一样的 transform并在预测前打印图片 shape 校验from PIL import Image def predict_one_image(model, image_path, transform, device): img Image.open(image_path).convert(RGB) x transform(img).unsqueeze(0).to(device) with torch.no_grad(): logits model(x) probs torch.softmax(logits, dim1) pred probs.argmax(dim1).item() confidence probs[0, pred].item() return pred, confidence这个小函数建议放在inference.py里答辩演示时直接调用不要在交互式环境里临时拼代码。6. 给毕设加分迁移学习、可解释性验证与单张推理习惯6.1 用预训练 ResNet 替换自建 CNN如果时间允许把自建 CNN 升级成torchvision里的resnet18(pretrainedTrue)代码只需要改几行import torchvision.models as models model models.resnet18(pretrainedTrue) model.fc torch.nn.Linear(model.fc.in_features, 6) for name, param in model.named_parameters(): if fc not in name: param.requires_grad False说明冻结前面的卷积主干只训练最后一层全连接训练速度快、不容易过拟合通常能在自建 CNN 基础上提升几个百分点。注意迁移学习的学习率要降到 0.0001 左右因为预训练权重已经在 ImageNet 上学到通用特征学习率太大会把这些能力冲掉。如果想进一步提升可以只冻结前 80% 层让最后两个 block 也跟着微调。6.2 用 Grad-CAM 让模型解释分类依据答辩时老师最常问“你的模型为什么这么分”。用 Grad-CAM 生成热力图能直观回答它把最后一个卷积层的梯度回传到特征图得到模型关注区域的响应和原图叠在一起后就能看出模型是看菜叶还是看盘子来决定厨余垃圾。pytorch_grad_cam这个库封装得很完整只需在模型 forward 后注册钩子读取特征和梯度最后保存一张叠加图。这一步不需要花太多时间但确实让毕设的完整度上一个档次。6.3 我每次训练结束坚持做的一组验证我的个人习惯是训练完成后一定写一份单张推理脚本随机从 test 集挑 10 张图片跑一遍打印真实标签、预测标签和置信度并保存到表格里。这么做的原因很简单测试集的整体准确率高不一定代表演示环节不出丑单独预测时最容易被预处理不一致、模型文件路径写错这类问题绊倒。把best_model.pth、inference.py和evaluate.py都放到同一个目录演示前先确认路径中没有中文、模型权重能够正常加载。从数据拆分、CNN 结构、训练调参到避坑这条路我已经走过很多遍最深的体会是这类毕设的难点从不在模型结构本身而在数据细节和训练耐心。希望这篇笔记能帮你的垃圾递归类项目少走几个弯路也希望你最后不只是一个能跑通的模型还有一套能解释、能演示、能经得起追问的完整工程。本文还有配套的精品资源点击获取