简介基于深度学习卷积神经网络实现的垃圾识别分类系统提供完整Python源码、近两千张标注图片数据集及训练好的模型文件。项目源自导师指导下的高分课程设计流程完整下载解压即可运行适合作为深度学习、计算机视觉方向课程设计与期末大作业的参考范本。包内共有2000个文件其中1196张jpg与789张jpeg构成可用于模型训练与验证的垃圾分类数据集13个py脚本覆盖数据预处理、模型训练、推理预测与可视化等核心环节另有json配置和md说明文档辅助快速理解项目结构。压缩包总大小约564.69MB已有263人浏览学习。通过该资源可掌握卷积神经网络图像分类任务的完整工程实现思路包括数据集组织、网络搭建、模型保存与调用方式并可直接基于现有代码扩展或迁移到其他分类场景。1. 垃圾识别分类系统的完整链路从数据集到CNN推理一份拿到97分的课程设计通常不是因为模型结构有多前沿而是整条链路完整、可复现、能自圆其说。这套基于深度学习卷积神经网络的垃圾识别分类系统数据集中是类似img_电池_452.jpeg、img_塑料瓶_244.jpeg这样真实拍摄的生活垃圾照片靠CNN完成电池、塑料瓶等类别的端到端识别。对准备课程设计或期末大作业的人来说它的价值在于数据集是现成的模型是训练好的源码拿到就能改能跑适合用来理解图像分类的完整工程流程而不是只看一个孤立的模型文件。本文按数据预处理、模型搭建、训练评估、推理部署的顺序拆解这套系统的实现细节。2. 数据集分析与图像预处理从原始 JPEG 到可训练张量2.1 先看清数据长什么样打开压缩包后建议先别急着写训练代码用一段脚本把数据集画像拉出来确认图片尺寸、通道数、类别分布再做后续处理。import os from PIL import Image from collections import Counter data_dir data/images samples [] for fname in os.listdir(data_dir): if not fname.lower().endswith((.jpg, .jpeg, .png)): continue label fname.split(_)[1] # 文件名格式 img_类别_编号.jpeg samples.append((fname, label)) print(样本总数:, len(samples)) print(类别分布:, Counter(label for _, label in samples)) # 抽查前5张图的尺寸 for fname, _ in samples[:5]: with Image.open(os.path.join(data_dir, fname)) as im: print(fname, im.size, im.mode)文件名的命名格式是img_类别_编号.jpeg可以直接把split(_)[1]作为标签。实际项目里如果文件名不规范常见做法是维护一个映射表把文件名前缀或子目录名映射到类别ID。这段代码的意义在于训练前确认两个关键信息类别是否均衡、图片尺寸是否统一。不均衡会导致后面模型偏向多数类尺寸不统一则需要在预处理阶段统一缩放。2.2 图像预处理的关键参数设计CNN要求输入张量尺寸固定但手机拍出来的照片尺寸五花八门。常见做法是统一缩放到224×224这也是ResNet系列的标准输入尺寸。缩放策略上直接resize会拉伸变形对电池、塑料瓶这类形状特征明显的物体影响不大但如果要更严谨应该用Resize加CenterCrop的组合。归一化参数用ImageNet的均值和标准差是通用做法因为预训练模型基于ImageNet统计量即使从零训练这一组数值也能加快收敛。from torchvision import transforms train_transforms transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(p0.5), 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_transforms transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])训练集和验证集的预处理必须分开。训练集用RandomCrop和ColorJitter做数据增强提升泛化能力验证集只做缩放和归一化保证评估结果稳定可复现。ColorJitter的颜色抖动幅度不要太大垃圾图片本身颜色是重要特征过度扰动会让模型学不到颜色信息。2.3 Dataset与DataLoader的实现定义好变换后需要把图片读入Dataset交给DataLoader做批量加载和多进程读取。import torch from torch.utils.data import Dataset, DataLoader from PIL import Image class GarbageDataset(Dataset): def __init__(self, data_dir, label_map, transformNone): self.samples [] for fname in os.listdir(data_dir): if not fname.lower().endswith((.jpg, .jpeg, .png)): continue label_str fname.split(_)[1] label label_map[label_str] self.samples.append((os.path.join(data_dir, fname), label)) self.transform transform def __len__(self): return len(self.samples) def __getitem__(self, idx): img_path, label self.samples[idx] image Image.open(img_path).convert(RGB) if self.transform: image self.transform(image) return image, label label_map {电池: 0, 塑料瓶: 1} train_loader DataLoader( GarbageDataset(data/images, label_map, train_transforms), batch_size32, shuffleTrue, num_workers4 )标签映射是关键步骤。建议显式用字典做字符串到整数的映射而不使用LabelEncoder这类自动编码工具因为可读性和下游推理时的可解释性都更好部署时也能直接对照。batch_size32是显存和收敛速度的平衡点显存不够降到16num_workers4让数据读取和模型训练并行避免GPU等待CPU读图。3. CNN 模型搭建卷积栈与迁移学习的选型实践3.1 卷积神经网络的核心组件与参数含义这套系统本质是一个图像分类任务CNN的作用是自动从像素中提取层次化特征。浅层卷积核捕捉边缘、纹理深层卷积核捕捉语义部件全连接层把特征映射到类别得分。组件典型配置作用注意事项Conv2din_channels3, out_channels64, kernel_size3, stride1, padding1提取局部特征padding保持尺寸不变时设为kernel_size//2BatchNorm2dnum_features64加速收敛、稳定训练训练和推理行为不同eval模式下用累计统计量ReLU无参数引入非线性避免用Sigmoid做隐藏层激活容易梯度消失MaxPool2dkernel_size2, stride2降维、扩大感受野会丢失位置信息位置敏感的检测任务慎用Dropoutp0.5抑制过拟合只在全连接层后用卷积层后少用Linearin_features, out_featuresnum_classes输出类别得分最终输出不需要Softmax损失函数内部处理3.2 从零搭建一个轻量CNN作为基线课程设计场景下从零写CNN能更直观地展示对卷积神经网络结构的理解也方便写进报告。以224×224输入为例一个三层卷积栈加两层全连接的结构就足够支撑这个小规模分类任务。import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes2): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, kernel_size3, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 112 nn.Conv2d(32, 64, kernel_size3, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 56 nn.Conv2d(64, 128, kernel_size3, padding1), nn.BatchNorm2d(128), nn.ReLU(inplaceTrue), nn.MaxPool2d(2), # 28 ) self.classifier nn.Sequential( nn.Dropout(0.5), nn.Linear(128 * 28 * 28, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.3), nn.Linear(256, num_classes) ) def forward(self, x): x self.features(x) x x.view(x.size(0), -1) return self.classifier(x)每经过一次MaxPool2d(2)特征图尺寸减半224经过三次池化变成28×28所以全连接层输入维度是128 * 28 * 28。如果改了输入尺寸或卷积层数这个数值必须同步调整这也是初学者最容易报错的地方。inplaceTrue能省一点显存但要注意不能对需要保留原始值的张量使用。3.3 迁移学习ResNet更稳的进阶选择自己的卷积栈参数量小、训练速度快但精度上限有限。如果想让模型效果更好加载在ImageNet上预训练过的ResNet18是常见做法。垃圾分类图片和ImageNet的自然图像分布接近冻结浅层参数、只训练最后几层也能在一个小时内完成训练。import torchvision.models as models model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) # 冻结前4层参数 for name, param in model.named_parameters(): if layer4 not in name and fc not in name: param.requires_grad False # 替换最后一层全连接适配本任务的类别数 in_features model.fc.in_features model.fc nn.Sequential( nn.Dropout(0.3), nn.Linear(in_features, 512), nn.ReLU(inplaceTrue), nn.Linear(512, 2) )这里的关键操作是把fc替换成一个两层的小分类头并只对layer4和fc层计算梯度。冻结浅层能大幅减少反传播的计算量同时保留预训练模型对边缘、纹理等通用特征的提取能力。原始ResNet18的fc是2048到1000的单层结构替换成512维中间层能提升非线性表达能力但如果数据量小中间层维度可以降到256甚至直接单层输出防止过拟合。4. 训练与评估优化器、损失函数和混淆矩阵的工程化配置4.1 训练循环与关键超参数训练脚本的核心是优化器、损失函数和学习率调度器的配合。垃圾图片分类是标准的单标签多分类问题用CrossEntropyLoss配合AdamW优化器是PyTorch里最稳妥的组合前者内部包含Softmax和负对数似然的计算不需要在网络输出层额外加Softmax。import torch.optim as optim from torch.optim.lr_scheduler import StepLR device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleCNN(num_classes2).to(device) criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr1e-4, weight_decay1e-2) # 不加 lr_scheduler scheduler StepLR(optimizer, step_size10, gamma0.1) best_acc 0.0 epochs 30 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) outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() running_loss loss.item() * images.size(0) # 验证 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() acc correct / total scheduler.step() if acc best_acc: best_acc acc torch.save(model.state_dict(), best_model.pth) print(fEpoch {epoch1}/{epochs} | Loss: {running_loss/len(train_loader.dataset):.4f} | Val Acc: {acc:.4f} | Best: {best_acc:.4f})lr1e-4对迁移学习是合适的起步值如果想从零训练自定义CNN可以提到1e-3。weight_decay是L2正则帮助抑制过拟合。学习率调度用StepLR每10个epoch降到原来的0.1倍避免训练后期在最优解附近震荡。model.train()和model.eval()别忘切换BatchNorm和Dropout在两种模式下行为完全不同忘切eval()是验证精度异常高的常见原因。4.2 早停与模型保存策略训练30个epoch如果验证集精度连续5轮不涨继续训练大概率只会过拟合这时候应该早停。保存模型时要注意state_dict()只保存参数不保存网络结构加载时需要先实例化同样结构的模型再load_state_dict。patience 5 early_stop_counter 0 if acc best_acc: best_acc acc early_stop_counter 0 torch.save({ epoch: epoch, model_state_dict: model.state_dict(), optimizer_state_dict: optimizer.state_dict(), best_val_acc: best_acc, }, checkpoint_best.tar) else: early_stop_counter 1 if early_stop_counter patience: print(fEarly stop at epoch {epoch1}) break直接保存state_dict的问题是没有优化器状态中断后想恢复训练只能从头开始。完整checkpoint保存到本地.tar文件里包含epoch、模型参数和优化器参数恢复训练时只要重新加载回来就可以接着上次的进度跑。保存完整checkpoint的磁盘占用一般也就几十到几百MB代价可以忽略。4.3 混淆矩阵与分类报告在课程设计答辩时评估模型只用准确率是不够的因为类别不均衡很容易靠多数类刷高准确率。这时至少要看混淆矩阵和Precision/Recall/F1这些细粒度指标用Scikit-learn现成的classification_report就能算出来。from sklearn.metrics import classification_report, confusion_matrix import numpy as np all_preds [] all_labels [] model.eval() 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) all_preds.extend(predicted.cpu().numpy()) all_labels.extend(labels.cpu().numpy()) target_names [电池, 塑料瓶] print(classification_report(all_labels, all_preds, target_namestarget_names)) print(confusion_matrix(all_labels, all_preds))classification_report输出每个类别的查准率、查全率和F1值电池这类数量较少的类别在这里的问题会比准确率暴露得更明显。混淆矩阵能看出哪些类别容易被互相混淆如果电池和塑料瓶大量互混说明颜色和形状特征不足以区分需要增加数据或加深网络。这两个指标对写课程设计报告来说也是必要素材答辩时老师一定会问。5. 推理部署进阶输入一致性、置信度阈值与模型导出5.1 训练和推理的输入一致性陷阱模型训练完最容易踩的坑是训练时做了Normalize(mean[0.485, ...], std[0.229, ...])推理时却忘记对单张图片做同样处理导致模型精度大幅下降。训练和推理的输入分布必须完全一致单张图片推理的预处理要与验证集保持一致不能夹带RandomCrop和ColorJitter这些数据增强操作。from PIL import Image def predict_single_image(model, image_path, label_map_reverse, device, threshold0.7): image Image.open(image_path).convert(RGB) tensor val_transforms(image).unsqueeze(0).to(device) # 必须转成四维 model.eval() with torch.no_grad(): logits model(tensor) probs torch.softmax(logits, dim1) conf, pred torch.max(probs, 1) if conf.item() threshold: return 无法判断, conf.item() return label_map_reverse[pred.item()], conf.item() label_map_reverse {0: 电池, 1: 塑料瓶} result, confidence predict_single_image(model, test_images/img_电池_517.jpeg, label_map_reverse, device) print(f识别结果: {result}, 置信度: {confidence:.4f})unsqueeze(0)把单张图片从三维张量扩展成[1, C, H, W]的四维批量输入这是推理新手最容易忽略的点漏掉就会报维度不匹配的错误。softmax把logits转成概率后取出最大项置信度过低时宁可拒绝识别也不要强行输出垃圾图片里出现背景杂乱、物体不完整的情况很常见阈值通常设在0.6到0.8之间。5.2 递归扫描目录并批量推理单张图片能跑通之后改成遍历目录做批量推理就只是包装一层循环了。这个脚本可以用在验收演示时把测试文件夹里的图片全部走一遍输出预测表格。import glob image_paths glob.glob(test_images/**/*.jpeg, recursiveTrue) for img_path in image_paths: result, confidence predict_single_image(model, img_path, label_map_reverse, device) print(f{os.path.basename(img_path):30s} - {result} ({confidence:.4f}))预测结果会直接显示在控制台或重定向到日志文件里。批量推理脚本对调试很有用能快速看到哪些图片被错分从而反向检查是预处理问题还是模型能力不足。5.3 导出ONNX加速推理如果要把模型交给答辩演示或集成到其他端侧应用把PyTorch模型导出成ONNX格式是性价比最高的做法。ONNX可以在CPU上获得比PyTorch原版更高的推理速度还能跨框架部署。导出时要把模型切到eval()模式并且输入尺寸要固定成训练时的[1, 3, 224, 224]。model.eval() dummy_input torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( model, dummy_input, garbage_classifier.onnx, input_names[input], output_names[output], opset_version11 )opset_version决定了ONNX算子集的兼容性较老的环境用11比较保险新环境可以用更高的版本。导出成功后会生成一个二进制模型文件后续可以用ONNX Runtime加载并推理就不再需要依赖PyTorch环境了。推理前仍然要走val_transforms预处理ONNX只负责从张量到张量的计算不做图像解码和归一化。整个流程跑通后这套项目就完成了从数据到模型再到服务的完整闭环。资源获取方式本项目完整源码、数据集及训练好的模型权重文件已整理打包。项目由导师指导完成并获得97分高分评价代码可直接运行。关注公众号「阿汤爱编程」回复「垃圾识别」即可获取下载链接。本文还有配套的精品资源点击获取