基于卷积神经网络的垃圾识别分类系统实战：数据集到部署全流程

简介：基于深度学习卷积神经网络实现的垃圾识别分类系统，提供完整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(p=0.5), transforms.ColorJitter(brightness=0.2, contrast=0.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, transform=None): 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_size=32, shuffle=True, num_workers=4 )

标签映射是关键步骤。建议显式用字典做字符串到整数的映射，而不使用LabelEncoder这类自动编码工具，因为可读性和下游推理时的可解释性都更好，部署时也能直接对照。batch_size=32是显存和收敛速度的平衡点，显存不够降到16；num_workers=4让数据读取和模型训练并行，避免GPU等待CPU读图。

3. CNN 模型搭建：卷积栈与迁移学习的选型实践

3.1 卷积神经网络的核心组件与参数含义

这套系统本质是一个图像分类任务，CNN的作用是自动从像素中提取层次化特征。浅层卷积核捕捉边缘、纹理，深层卷积核捕捉语义部件，全连接层把特征映射到类别得分。

3.2 从零搭建一个轻量CNN作为基线

课程设计场景下，从零写CNN能更直观地展示对卷积神经网络结构的理解，也方便写进报告。以224×224输入为例，一个三层卷积栈加两层全连接的结构就足够支撑这个小规模分类任务。

import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self, num_classes=2): super().__init__() self.features = nn.Sequential( nn.Conv2d(3, 32, kernel_size=3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.MaxPool2d(2), # 112 nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.MaxPool2d(2), # 56 nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.MaxPool2d(2), # 28 ) self.classifier = nn.Sequential( nn.Dropout(0.5), nn.Linear(128 * 28 * 28, 256), nn.ReLU(inplace=True), 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。如果改了输入尺寸或卷积层数，这个数值必须同步调整，这也是初学者最容易报错的地方。inplace=True能省一点显存，但要注意不能对需要保留原始值的张量使用。

3.3 迁移学习ResNet：更稳的进阶选择

自己的卷积栈参数量小、训练速度快，但精度上限有限。如果想让模型效果更好，加载在ImageNet上预训练过的ResNet18是常见做法。垃圾分类图片和ImageNet的自然图像分布接近，冻结浅层参数、只训练最后几层，也能在一个小时内完成训练。

import torchvision.models as models model = models.resnet18(weights=models.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(inplace=True), 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_classes=2).to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-2) # 不加 lr_scheduler scheduler = StepLR(optimizer, step_size=10, gamma=0.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(f"Epoch {epoch+1}/{epochs} | Loss: {running_loss/len(train_loader.dataset):.4f} | Val Acc: {acc:.4f} | Best: {best_acc:.4f}")

lr=1e-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(f"Early stop at epoch {epoch+1}") 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_names=target_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, threshold=0.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, dim=1) 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", recursive=True) 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_version=11 )

opset_version决定了ONNX算子集的兼容性，较老的环境用11比较保险，新环境可以用更高的版本。导出成功后会生成一个二进制模型文件，后续可以用ONNX Runtime加载并推理，就不再需要依赖PyTorch环境了。推理前仍然要走val_transforms预处理，ONNX只负责从张量到张量的计算，不做图像解码和归一化。整个流程跑通后，这套项目就完成了从数据到模型再到服务的完整闭环。

资源获取方式：本项目完整源码、数据集及训练好的模型权重文件已整理打包。项目由导师指导完成并获得97分高分评价，代码可直接运行。关注公众号「阿汤爱编程」回复「垃圾识别」即可获取下载链接。

本文还有配套的精品资源，点击获取