ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

用PyTorch构建轻量级CNN实现垃圾分类图像识别与桌面应用

用PyTorch构建轻量级CNN实现垃圾分类图像识别与桌面应用 简介这是一套基于卷积神经网络的垃圾分类系统实现资源包含完整的图像分类模型与可运行的GUI应用面向深度学习初学者、计算机视觉方向学生以及需要快速搭建垃圾分类演示项目的开发者。包内涵盖模型设计与实验过程包括卷积层深度调整、Batch Normalization、学习率调度、Dropout正则化等优化策略最终模型准确率约90%。同时提供基于PySimpleGUI的桌面小程序支持选择文件夹、加载图像并输出预测类别便于直观理解垃圾分类流程。资源包共2000个文件以1978张jpg图像数据集为主另有Python源码、XML配置文件、说明文档等整体大小146.47MB适合直接用于模型训练、测试和界面演示。已有202人学习下载可用于课程设计、毕业设计或科研预研。1. 为什么不堆深度也能把垃圾分类做到 90% 准确率垃圾分类这件事真正做过落地的人都知道痛点瓶、纸、金属、厨余垃圾在照片里的颜色、纹理、摆放角度差异极大同一类垃圾换个光线就完全变样。直接用传统视觉方案做分类特征工程能把人熬秃。这个项目没有选择 ResNet 级别的深层网络而是从两个卷积层起步通过 Batch Normalization、Dropout 和正确的学习率调度把模型压到 90% 左右的准确率然后封装成一个 PySimpleGUI 桌面应用。它解决的核心问题是在硬件条件一般、做演示系统或课程设计时如何用一套可复现的流程完成训练、评估和 GUI 预测。适合需要快速跑通整套流程的开发者——既能看懂卷积神经网络结构图也能自己动手调参而不是只停留在理论。2. 卷积核、池化与填充垃圾分类图像建模的第一步2.1 为什么垃圾图像分类要用 CNN 而不是传统特征垃圾图像与常规物体识别最大的区别在于类内差异大一个矿泉水瓶和一张揉皱的纸颜色都可能接近但纹理完全不同。传统方法先提取 HOG、颜色直方图这类手工特征再用 SVM 分类特征在不同垃圾类别之间的可分性很难保证。卷积神经网络通过卷积核自动学习局部纹理、边缘和颜色组合避免了手工设计特征的瓶颈。这类任务也不适合图卷积网络。图卷积处理的是非欧几里得结构数据比如社交关系、分子结构普通相机拍出来的 RGB 图像是规则的网格结构强行建模成图反而丢失了空间局部性。实践中标准卷积神经网络CNN仍然是最适合做法。2.2 卷积核、步长、填充 对特征图尺寸的影响拿到一张 64×64 的垃圾图像经过卷积层后特征图尺寸计算公式是[ output \lfloor (input 2 \times padding - kernel) / stride \rfloor 1 ]这个式子决定了网络每一层的输出大小实际工程中也直接影响显存占用。以两卷积层的小网络为例常见的参数组合如下表场景卷积核步长 stride填充 padding输入尺寸输出尺寸第一层提取边缘3×31164×6464×64压缩特征图3×32164×6432×32保持分辨率5×51232×3232×32快速降采样3×32032×3215×15实际项目里第一层卷积通常保留原始空间尺寸让网络充分学习边缘信息第二层再用 stride2 或池化降低分辨率。下面是 PyTorch 中定义卷积层的对应写法import torch.nn as nn # 第一层卷积输入1通道灰度图输出16个特征图 conv1 nn.Conv2d(in_channels1, out_channels16, kernel_size3, stride1, padding1) # 第二层卷积输入16通道输出32个特征图步长2实现下采样 conv2 nn.Conv2d(in_channels16, out_channels32, kernel_size3, stride2, padding1)代码说明in_channels为输入通道数灰度图是 1RGB 图是 3。out_channels是卷积核个数也是输出特征图的通道数。kernel_size3表示 3×3 的卷积核padding1保证 stride1 时输出尺寸不变。第二层卷积stride2配合padding164×64 输入会输出 32×32。参数的核心逻辑是通道数随深度翻倍空间尺寸随下采样减半这是 LeNet-5 风格的经典设计。2.3 Batch Normalization 与 Dropout 放在哪里很多初学者拿到卷积神经网络结构图会把 Batch Normalization 放在卷积之后、激活函数之后这个顺序看似不影响精度实际收敛速度有明显差异。常见做法是卷积 → BN → ReLU原因是 BN 需要在非线性变换之前稳定数据分布否则 ReLU 会把负值截断BN 统计的分布就不完整。Dropout 位置则不同一般只放在全连接层之后不放在卷积层后。卷积层的参数共享本身已经具备正则化效果过早加入 Dropout 会降低特征提取能力。一个典型的两卷积层结构如下self.features nn.Sequential( nn.Conv2d(1, 16, kernel_size3, stride1, padding1), nn.BatchNorm2d(16), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), nn.Conv2d(16, 32, kernel_size3, stride1, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), ) self.classifier nn.Sequential( nn.Linear(32 * 16 * 16, 128), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(128, 6) # 6类垃圾 )参数说明两个池化层把 64×64 输入逐步降到 16×16self.classifier的输入维度就是32 * 16 * 16修改输入图像大小时这里必须同步调整。最后输出维度为垃圾类别数二分类就改为 2多分类按实际数据集设置。Dropout(0.5)表示训练时随机丢弃 50% 的神经元预测时自动关闭。提示如果训练集较小Dropout 率建议从 0.3 开始调0.5 在类别多但每类样本少时容易欠拟合。3. 从数据目录到模型训练完整复现垃圾分类的训练流程3.1 数据集组织与预处理策略数据组织直接决定后续代码复杂度。推荐使用 ImageFolder 格式每个类别一个子目录目录名是类别名data/ ├── train/ │ ├── glass/ # 玻璃 │ ├── paper/ # 纸 │ ├── metal/ # 金属 │ ├── plastic/ # 塑料 │ ├── cardboard/ # 纸板 │ └── trash/ # 其他垃圾 └── val/ └── ...训练时使用torchvision.datasets.ImageFolder加载配合transforms做预处理。垃圾图片在用户拍摄时经常出现曝光不均、角度倾斜的情况数据增强环节如果做得不够验证集准确率会明显低于训练集准确率。这里给出完整的变换组合from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((64, 64)), transforms.RandomHorizontalFlip(p0.5), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485], std[0.229]) # 灰度图单通道 ]) val_transform transforms.Compose([ transforms.Resize((64, 64)), transforms.ToTensor(), transforms.Normalize(mean[0.485], std[0.229]) ])逻辑说明RandomHorizontalFlip模拟拍摄角度变化ColorJitter提高模型对光线的鲁棒性。灰度图只有一个通道所以Normalize的mean和std只给一个值如果使用 RGB 图则需要三个值比如mean[0.485, 0.456, 0.406]对应 ImageNet 统计值。提示训练集和验证集必须使用相同的Resize和Normalize否则验证阶段的输入分布与训练不一致准确率指标会失真。3.2 模型定义两卷积层加上 BN 的小网络这个项目最终选定的网络结构不是盲目的参考了 LeNet-5 的 conv-pool 交替设计。垃圾图像的纹理复杂度低于 ImageNet 场景两个卷积层加全连接层足够提取区分度特征。模型完整定义如下import torch 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, 32, kernel_size3, stride1, padding1), nn.BatchNorm2d(32), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), nn.Conv2d(32, 64, kernel_size3, stride1, padding1), nn.BatchNorm2d(64), nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), ) self.classifier nn.Sequential( nn.Linear(64 * 16 * 16, 256), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(256, num_classes) ) def forward(self, x): x self.features(x) x torch.flatten(x, start_dim1) x self.classifier(x) return x结构说明第一层卷积输出 32 个特征图池化后尺寸从 64×64 降到 32×32第二层卷积输出 64 个特征图再池化到 16×16。torch.flatten将[batch, 64, 16, 16]展平为[batch, 64*16*16]送入全连接层。这里的通道数为什么要依次翻倍因为池化会丢失空间信息通道数增加是为了保留足够的高层语义信息。3.3 训练循环与超参数基准值训练循环不复杂但几个细节决定了最终准确率是否稳定。一个是优化器选择AdamW比原生 Adam 的权重衰减实现更规范另一个是损失函数多分类任务直接使用CrossEntropyLoss它对原始 logits 计算 softmax 和交叉熵不需要手动加 Softmax 层。核心训练代码如下import torch.optim as optim from torch.optim.lr_scheduler import ReduceLROnPlateau model GarbageCNN(num_classes6) device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.AdamW(model.parameters(), lr0.001, weight_decay1e-4) scheduler ReduceLROnPlateau(optimizer, modemin, factor0.1, patience3, verboseTrue) for epoch in range(30): model.train() running_loss 0.0 for inputs, labels in train_loader: inputs, labels inputs.to(device), labels.to(device) optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * inputs.size(0) model.eval() val_loss 0.0 correct 0 total 0 with torch.no_grad(): for inputs, labels in val_loader: inputs, labels inputs.to(device), labels.to(device) outputs model(inputs) loss criterion(outputs, labels) val_loss loss.item() * inputs.size(0) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() val_acc correct / total scheduler.step(val_loss) print(fEpoch {epoch1}: train_loss{running_loss/len(train_loader.dataset):.4f}, fval_loss{val_loss/len(val_loader.dataset):.4f}, val_acc{val_acc:.4f})逻辑说明每个 epoch 先执行训练模式model.train()反向传播更新权重验证阶段切到model.eval()并关闭梯度计算减少显存占用并避免 BN 层更新均值和方差。torch.max(outputs, 1)返回每行最大值的索引即预测类别。scheduler.step(val_loss)会监控验证损失连续 3 个 epoch 不下降时学习率乘以因子 0.1。超参数基准值参考下表超参数推荐取值调参方向初始学习率0.001过大导致损失震荡过小收敛缓慢批量大小 batch_size32 或 64显存不足时减半Dropout 率0.5过拟合明显时增加到 0.6weight_decay1e-4L2 正则化强度过拟合时可调到 1e-3优化器AdamW任务简单可换 SGD Momentum4. 学习率调度与正则化把准确率从 80% 拉到 90% 的关键调整4.1 固定学习率为什么不行训练初期权重距离最优解较远0.001 的学习率能快速下降但训练到后期参数在最优解附近来回震荡固定学习率会导致验证损失降低到某个阈值后不再变化甚至反弹。垃圾数据集的类内差异大损失曲面陡峭且不平滑这个问题更明显。两个实用的调度器是StepLR和ReduceLROnPlateau。StepLR每固定步数将学习率乘以系数简单但不考虑实际训练状态ReduceLROnPlateau根据验证损失动态调整更适合这个项目。完整替换代码如下from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau # 方案一每 10 个 epoch 学习率乘 0.1 scheduler_step StepLR(optimizer, step_size10, gamma0.1) # 方案二验证损失连续 3 次未下降时学习率乘 0.1 scheduler_plateau ReduceLROnPlateau(optimizer, modemin, factor0.1, patience3)使用区别StepLR需要在每个 epoch 结束后调用scheduler_step.step()它不依赖损失值ReduceLROnPlateau需要传入当前验证损失scheduler_plateau.step(val_loss)。实际训练中如果第 15 个 epoch 时验证准确率徘徊在 85% 左右用 AdamW 配合 ReduceLROnPlateau 后学习率从 0.001 下降到 0.0001准确率通常能继续爬升 3-5 个百分点。4.2 AdamW 与 RMSprop 的取舍优化器不是越新越好。AdamW 在 PyTorch 中的实现与 Adam 的主要区别是把权重衰减从梯度更新中解耦对使用 weight_decay 的场景更稳定。RMSprop 的特点是自适应调节每个参数的学习率对梯度抖动敏感的任务有效但在多分类任务的收敛速度上通常不如 AdamW。对比见下表优化器收敛速度对超参数敏感度适用场景AdamW快较低大多数任务推荐优先尝试Adam快较低与 AdamW 效果接近但权重衰减实现有差异RMSprop中等中序列模型或梯度剧烈变化场景SGD Momentum慢高训练资源充足、需要精细调参追求极致精度垃圾图像分类的类别数通常不超过 10特征空间相对简单AdamW 是最高性价比的选择。4.3 过拟合的判断信号与正则化组合训练准确率持续上升而验证准确率停滞或下降是最典型的过拟合信号。此时优先检查 Dropout 是否只加在了全连接层、训练集图片数量是否与网络容量匹配。数据增强也是一种正则化手段RandomHorizontalFlip和ColorJitter组合能显著提升模型的泛化能力它们与 Dropout、weight_decay 并不冲突可以同时使用。提示正则化不是加得越多越好。当训练集本身只有几百张图片时过强的 Dropout 反而让模型欠拟合表现是训练准确率和验证准确率同步低于 80%。遇到这种情况先把 Dropout 率降到 0.3去掉 weight_decay 再观察。另一个容易踩坑的点是验证集图片尺寸。很多人在训练时用transforms.Resize((64, 64))但验证时写成Resize((64, 64))后忘记ToTensor导致输入类型不对报错或者精度异常。调试时可以先打印一组验证数据的形状for inputs, labels in val_loader: print(inputs.shape) # torch.Size([32, 3, 64, 64]) print(labels.shape) # torch.Size([32]) breaktorch.Size([32, 3, 64, 64])表示批量大小为 32、RGB 三通道、宽高 64 像素。如果输出是[32, 64, 64]说明ToTensor缺失Normalize也会失效。这类问题在模型结构正确时最难排查列出数据形状是最快的定位方式。5. 用 PySimpleGUI 做桌面前端加载模型、选图、预测5.1 GUI 布局设计终端里跑模型和做成一个可点击的应用体验差别很大。PySimpleGUI 不需要额外安装前端框架代码量比 Tkinter 原生写法少一半适合模型演示。核心布局是左侧文件浏览器选择图片文件夹中间列表展示文件夹内所有图片右侧显示选中图片和预测结果底部一个“开始预测”按钮。import PySimpleGUI as sg layout [ [sg.Text(垃圾分类识别系统, font(Microsoft YaHei, 16))], [sg.Text(选择文件夹), sg.InputText(key-FOLDER-), sg.FolderBrowse(浏览)], [sg.Listbox(values[], size(30, 10), key-FILE_LIST-, enable_eventsTrue)], [sg.Image(key-IMAGE-, size(300, 300))], [sg.Text(预测结果, key-RESULT-, font(Microsoft YaHei, 14))], [sg.Button(开始预测), sg.Button(退出)] ] window sg.Window(垃圾分类, layout)布局逻辑sg.FolderBrowse返回选中的文件夹路径存入-FOLDER-输入框。Listbox设置enable_eventsTrue后单击图片名会触发事件方便在左侧选中图片后立即显示预览。5.2 事件循环与预测调用事件循环是 GUI 应用的核心。用户点击“开始预测”后程序从-FILE_LIST-拿到选中文件名拼接完整路径读取图片做与训练时一致的预处理然后调用模型输出结果。预处理步骤这里不能省略这也是很多复现项目运行时准确率突然下降的原因——训练时有Normalize预测时忘记做。from PIL import Image import torchvision.transforms as transforms def predict_image(model, image_path, class_names, device): transform transforms.Compose([ transforms.Resize((64, 64)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ]) image Image.open(image_path).convert(RGB) tensor transform(image).unsqueeze(0).to(device) model.eval() with torch.no_grad(): outputs model(tensor) _, predicted torch.max(outputs, 1) return class_names[predicted.item()] while True: event, values window.read() if event in (sg.WINDOW_CLOSED, 退出): break if event -FOLDER-: import os folder values[-FOLDER-] files [f for f in os.listdir(folder) if f.lower().endswith((.jpg, .jpeg, .png))] window[-FILE_LIST-].update(files) if event -FILE_LIST-: file_name values[-FILE_LIST-][0] file_path os.path.join(values[-FOLDER-], file_name) window[-IMAGE-].update(filenamefile_path) if event 开始预测: file_name values[-FILE_LIST-][0] file_path os.path.join(values[-FOLDER-], file_name) result predict_image(model, file_path, class_names, device) window[-RESULT-].update(f预测结果{result})事件说明-FOLDER-事件在用户选择文件夹后触发过滤出图片格式文件-FILE_LIST-事件在鼠标单击列表项时触发更新右侧预览开始预测按钮触发模型推理。Image.open(image_path).convert(RGB)把灰度图或者带透明通道的 PNG 统一转为 RGB避免通道数不一致导致模型报错。5.3 进阶技巧用 TorchScript 加速 GUI 启动PySimpleGUI 程序启动时如果直接torch.load完整模型需要先构建模型类再加载权重代码一旦改动GUI 端也要跟着改。更稳妥的做法是训练完成后导出 TorchScript 格式GUI 只依赖torch.jit.load不用再 import 模型定义。导出命令model.eval() example_input torch.randn(1, 3, 64, 64).to(device) traced_model torch.jit.trace(model, example_input) traced_model.save(garbage_model.pt)在 GUI 代码中用model torch.jit.load(garbage_model.pt)替换原来的GarbageCNN()加载方式。这样做的好处是 GUI 启动时间更短且模型部署时不依赖训练脚本中的类和参数定义。如果后续把模型迁移到移动端或服务端TorchScript 也可以避免 Python 环境依赖问题。把这条加载方式换掉这个桌面应用就可以作为一个独立的演示程序分发给其他人。本文还有配套的精品资源点击获取
返回列表