ARTICLE DETAIL

资讯详情

深耕网站建设、视觉设计与SEO优化的一线实战洞察。

基于PyTorch的垃圾分类识别系统:从CNN训练到GUI部署实践

基于PyTorch的垃圾分类识别系统:从CNN训练到GUI部署实践 简介基于卷积神经网络的垃圾分类系统完整工程针对大学生毕业设计、课程设计或机器学习实践场景解决从数据集准备到图形化识别应用的完整链路需求。资源共2000个文件以1978张垃圾图像数据为核心辅以6个Python源码含模型构建、训练与GUI程序、5个XML配置文件、2个TXT实验说明等压缩包体积约146.47MB。已有202人学习下载兼具研究与复现价值。内容详细记录了模型优化实验增加卷积层数量、引入Batch Normalization加速收敛、调整学习率和Dropout等超参数最终准确率约90%同时提供基于PySimpleGUI的桌面应用支持选择图片并实时输出垃圾分类结果便于直观理解深度学习落地流程。适合想快速跑通项目或参考其架构进行改进的学习者。1. 垃圾分类识别系统卷积神经网络负责看见GUI负责被用起来一个训练好的分类模型离一个能交给别人使用的系统中间还差一层交互包装。标题把卷积神经网络、垃圾分类和GUI应用绑在一起意味着要交付的不只是一组准确率数字而是一个非技术用户打开就能用的桌面程序。典型场景是社区宣传点、校园课程设计、单位内部的回收辅助识别对象通常从可回收、有害、厨余、其他四分类细化到玻璃、塑料、金属、纸类、电池等具体类别。后面的内容按「CNN的卷积、池化、步长、填充参数怎么定 → 垃圾分类数据怎么组织、训练参数怎么调 → 模型怎么接进GUI完成预测 → 部署后置信度怎么验证」推进每一节都有可直接运行的PyTorch代码代码后面都有参数说明和失败时的排查方向。适合课程设计、毕业设计以及想把模型落成小工具的一线开发者和运维人员。这里不讨论分布式训练和边缘端部署目标是在个人电脑上把一个中小规模的垃圾分类识别系统完整跑通。2. 卷积神经网络的参数设计与PyTorch基础实现2.1 卷积核、步长、填充特征图尺寸的一体两面卷积层的本质是滑窗扫描窗口本身叫卷积核它的大小决定了单次能看到多少局部特征。3×3 是当前主流两个 3×3 堆叠等价于一个 5×5 的感受野但参数量只有后者的 18/25所以 VGG 之后很少再有网络用大核。1×1 卷积不做空间聚合只做通道混合常用于降维。早期经典网络 LeNet-5 用的就是 5×5 卷积加池化的组合现在看这种搭配在尺寸小、通道少的数据上依然够用。处理垃圾分类这类视觉差异较为明显的任务时核大小对结果的影响远不如网络深度和通道数明显不需要在 3×3 和 5×5 之间反复纠结。步长控制滑窗每次移动的距离。步长为 1 时特征图尺寸基本不变步长为 2 时尺寸减半等价于一次可学习的下采样。填充解决边缘像素参与计算次数偏少的问题padding 为 1、核为 3、步长为 1 时输出尺寸等于输入尺寸。特征图尺寸计算公式是 output (input 2 * padding − kernel) / stride 1。训练过程中报 shape mismatch绝大多数不是网络结构写错而是全连接层输入维度没用这个公式推导后面用自适应池化可以直接规避这类手算这是 2.3 节会重点讲的做法。关于填充还有一种说法需要区分paddingsame 表示输出尺寸与输入一致但它只对步长为 1 成立步长为 2 时下采样必然发生输出尺寸照样减半。不少实训平台比如头歌的卷积神经网络关卡会先让人手算特征图尺寸再写代码说明这些参数直接关系到能否读懂一张卷积神经网络结构图。常见结构图上每一层标注的 112×112×32 就是这一层输出的高、宽、通道数输入经过池化后宽高折半通道数翻倍最后接全连接层给分类结果。参数常见取值对特征图的影响典型用途卷积核 kernel_size1×1 / 3×3 / 5×5核越大感受野越大参数量越多1×1 降维3×3 主流5×5 早期网络步长 stride1 / 2stride2 时宽高减半替代池化做下采样填充 padding0 / 1 / same补偿边缘信息核 3 配 padding1、stride1 保持尺寸池化 poolingMaxPool / AvgPool宽高减半通道不变下采样与防过拟合2.2 池化层为什么不能省池化没有可学习参数作用是下采样和增强平移不变性。最大池化取窗口内最大值保留最显著响应平均池化保留整体统计特征。2×2 窗口、步长为 2 的最大池化把空间尺寸减半后续计算量降为原来的四分之一。近年不少网络改用步长为 2 的卷积做下采样让降采样过程也参与学习ResNet 就是这么处理的。在小型垃圾分类数据集上两种方案差异不大我倾向于使用 MaxPool2d理由是结构图好画、调试时特征变化一目了然给课程设计答辩讲解也更直观。池化层还有一个实际工程作用把全连接层的输入维度稳定下来。如果不加池化224×224 的输入经过三层步长 2 的池化后变成 28×28、通道 64全连接层输入必须写成 64×28×28一旦换成别的输入分辨率这个数字又要重算。用 AdaptiveAvgPool2d 把任意尺寸的特征压到 1×1全连接输入就固定等于通道数模型也能接受任意分辨率输入。GUI 推理时用户随手拖进来的图片尺寸各不相同这个技巧能少改很多代码。2.3 用PyTorch搭一个可直接改分类数的骨架网络下面这个网络是六个类别的垃圾分类骨架卷积层后面跟 ReLU 和池化通道数从 16 翻到 64特征图从 224 依次减半到 28。import torch.nn as nn class GarbageCNN(nn.Module): def __init__(self, num_classes6): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 16, kernel_size3, stride1, padding1), # 224x224x16 nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), # 112x112x16 nn.Conv2d(16, 32, kernel_size3, stride1, padding1), # 112x112x32 nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), # 56x56x32 nn.Conv2d(32, 64, kernel_size3, stride1, padding1), # 56x56x64 nn.ReLU(inplaceTrue), nn.MaxPool2d(kernel_size2, stride2), # 28x28x64 ) self.classifier nn.Sequential( nn.AdaptiveAvgPool2d((1, 1)), # 1x1x64 nn.Flatten(), # 64 nn.Linear(64, 128), nn.ReLU(inplaceTrue), nn.Dropout(0.5), nn.Linear(128, num_classes), ) def forward(self, x): return self.classifier(self.features(x))注释里标注的尺寸就是 2.1 节公式的实际应用调试时逐层对照出错很快能定位。分类器用 AdaptiveAvgPool2d 把任意尺寸特征压成 1×1再 Flatten 成 64 维向量好处是输入分辨率改变时不用重算全连接维度GUI 里换图也不用改模型。Dropout 0.5 放在全连接前配合后面 3.3 节的早停能明显抑制小数据集的过拟合。如果手头数据量超过每类一千张也可以换 torchvision 里预训练的 ResNet18只把最后一层全连接改成nn.Linear(resnet.fc.in_features, 6)然后微调。自己从零搭一遍骨架、再切到预训练模型对理解结构图和迁移学习都会有更具体的体会数据量小的时候预训练模型通常比从零训练的自建网络高五到十个百分点的准确率。3. 垃圾分类训练集的组织与训练参数调优3.1 用ImageFolder按目录划分类别标签顺序保存好训练垃圾分类模型的第一步不是写网络而是把图片按类别放进目录用 torchvision 的 ImageFolder 直接读取。标准目录结构如下dataset/ ├── train/ │ ├── glass/ 玻璃 │ ├── plastic/ 塑料 │ ├── paper/ 纸张 │ ├── metal/ 金属 │ ├── kitchen/ 厨余 │ └── other/ 其他 └── val/ ├── glass/ └── ...ImageFolder 会把子目录名按字典序排序后编号glass 不一定是 class 0目录名排序后的位置才决定标签值。所以训练完保存 label_mapping 时一定从dataset.class_to_idx导出并随模型一起存成 JSON不要自己硬编码顺序。GUI 里分类结果和图片对不上十有八九是这一环节写死了标签顺序。公开的垃圾分类数据集常见四十类版本包含旧衣服、一次性餐具、指甲油等细类类别越多混淆越严重。我一般建议先用四到六类跑通完整链路确认模型、GUI、预处理都没问题后再扩类有些类别图片只有几百张宁可先合并到上一个层级也不要强行撑类别数量。from torchvision.datasets import ImageFolder train_ds ImageFolder(dataset/train) val_ds ImageFolder(dataset/val) # 把这个映射存下来GUI 推理时要用 idx_to_label {v: k for k, v in train_ds.class_to_idx.items()} print(idx_to_label)class_to_idx是目录名到整型标签的映射反向得到标签到类别名的字典。打印出来核对一遍确认 glass、plastic 这些名称和预期一致再进入训练环节。3.2 数据增强的合理边界别把玻璃杯变成塑料瓶垃圾分类依赖颜色、纹理和形状增强强度要克制。旋转角度过大或者色彩抖动过猛会把玻璃瓶的透光特征改成塑料质感模型反而学不到真实分布。下面是我在垃圾分类上常用的一组增强配置from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomResizedCrop(224, scale(0.7, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(degrees15), transforms.ColorJitter(brightness0.3, contrast0.3, saturation0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) val_transform transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ])RandomResizedCrop 的 scale 下限设 0.7太小的裁剪会让模型只看到局部纹理而丢失整体形态。ColorJitter 里特意不调 hue色相因为玻璃和塑料的区分很大程度依赖颜色本身。Normalize 用的是 ImageNet 的均值方差迁移预训练模型时保持这个数值能直接复用预训练权重从零训练自建网络时继续用它也不是错误只是需要知道这组数值的来源。验证集只做 Resize 和 Normalize不做任何随机增强否则验证准确率会忽高忽低早停判断也会失真。3.3 训练循环、最优模型保存与超参数表训练部分的核心是每个 epoch 更新参数、在验证集上评估只保存验证准确率最高的权重。代码如下import torch import torch.nn as nn from torch.utils.data import DataLoader train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4) model GarbageCNN(num_classes6) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, modemax, factor0.5, patience5 ) best_acc 0.0 for epoch in range(50): model.train() running_loss 0.0 for images, labels in train_loader: optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) model.eval() correct total 0 with torch.no_grad(): for images, labels in val_loader: outputs model(images) _, preds torch.max(outputs, 1) correct (preds labels).sum().item() total labels.size(0) val_acc correct / total scheduler.step(val_acc) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_garbage_cnn.pth) print(fepoch {epoch 1}: train_loss{running_loss / len(train_ds):.4f}, fval_acc{val_acc:.4f}, best{best_acc:.4f})loss 用 CrossEntropyLoss内部已经包含 softmax网络最后一层不需要额外接 softmax。优化器用 Adam稳定后可以观察 loss 是否持续下降如果验证集连续多个 epoch 不涨ReduceLROnPlateau 会把学习率减半再试。保存的 state_dict 只含权重加载时需要先实例化和训练时完全一样的网络结构。超参数推荐值说明输入分辨率224×224与预训练模型对齐方便迁移batch size32~64按显存调整太小梯度抖动大优化器Adamlr1e-3后期可换 SGDmomentum 精调学习率衰减factor0.5patience5验证集不涨时降一半早停轮数10~15连续不涨就停止训练Dropout0.5全连接层前防过拟合num_workers 在 Windows 上设置为 0 可以避免多进程数据加载报错Linux 下可以按 CPU 核心数调大。显存不够时优先降 batch size 而不是降分辨率244×224 的输入尺寸和预训练模型的兼容性最好。4. 模型接入GUI应用的推理链路4.1 推理预处理必须和训练严格一致GUI 里最常见的错误是推理时用了另一套预处理。训练时做了 Normalize推理时不归一化或者训练用的 Resize 是 224推理时传原图模型输出就会变成随机结果。推理预处理要和验证集完全一致from PIL import Image from torchvision import transforms INFER_TRANSFORM transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) def preprocess(path: str): img Image.open(path).convert(RGB) return INFER_TRANSFORM(img).unsqueeze(0) # [1, 3, 224, 224]PIL 读图默认通道顺序就是 RGB和训练一致。如果改用 OpenCV 的 cv2.imread读进来是 BGR必须cv2.cvtColor(img, cv2.COLOR_BGR2RGB)转回 RGB否则颜色通道错位玻璃和塑料这类颜色敏感的类别会乱掉。unsqueeze(0) 的作用是给单张图补一个 batch 维度让输入变成四维张量这是模型 forward 所要求的形状。4.2 用Tkinter实现选图、识别、结果展示把模型封装进一个类里预测逻辑和界面逻辑分开写后期加批量识别或者换界面框架都不用动推理代码。import tkinter as tk from tkinter import filedialog from PIL import Image, ImageTk import torch class GarbageApp: def __init__(self, model, idx_to_label, devicecpu): self.model model.to(device).eval() self.idx_to_label idx_to_label self.device device self.root tk.Tk() self.root.title(垃圾分类识别系统) self.root.geometry(560x520) self.image_label tk.Label(self.root, text请选择图片, bg#f0f0f0, width48, height16) self.image_label.pack(pady10) self.choose_btn tk.Button(self.root, text选择图片, commandself.open_file, width20) self.choose_btn.pack(pady5) self.result_label tk.Label(self.root, text, font(微软雅黑, 16)) self.result_label.pack(pady10) self.conf_label tk.Label(self.root, text, font(微软雅黑, 12)) self.conf_label.pack() def open_file(self): path filedialog.askopenfilename( filetypes[(图片文件, *.jpg *.jpeg *.png *.bmp)]) if not path: return display self.load_display_image(path) self.image_label.config(imagedisplay) self.image_label.image display # 防止被垃圾回收 category, confidence self.predict(path) self.result_label.config(textf识别结果{category}) self.conf_label.config(textf置信度{confidence:.2%}) def load_display_image(self, path): img Image.open(path).convert(RGB) img.thumbnail((320, 260)) return ImageTk.PhotoImage(img) def predict(self, path): tensor preprocess(path) with torch.no_grad(): logits self.model(tensor) prob torch.softmax(logits, dim1)[0] score, index torch.max(prob, dim0) return self.idx_to_label[index.item()], score.item() def run(self): self.root.mainloop()predict 里用 softmax 把 logits 转成概率再用 torch.max 取最大项和对应下标返回的是类别名和置信度。界面那一行self.image_label.image display不能省Tkinter 的 PhotoImage 如果没有 Python 变量持有引用会被垃圾回收导致图片不显示这是 Tkinter 开发里非常经典的坑。这个骨架 CNN 在 CPU 上单张推理约几十毫秒放进按钮回调里不会明显卡顿。4.3 模型变大时用线程保护界面响应如果换成 ResNet18单张推理耗时会到两三百毫秒界面上会出现明显的点击无响应。处理方式是把预测丢到后台线程预测完再通过 after 回调更新界面import threading def open_file(self, path): threading.Thread(targetself.predict_async, args(path,), daemonTrue).start() def predict_async(self, path): category, confidence self.predict(path) self.root.after(0, lambda: self.result_label.config( textf识别结果{category} 置信度{confidence:.2%}))Tkinter 不是线程安全的子线程里不能直接改控件属性root.after(0, ...)把界面更新请求塞回主线程的事件循环执行这是 Tkinter 跨线程更新界面的标准做法。daemonTrue 保证关窗口时后台线程随进程退出不会悬挂。5. 部署后必做的两个验证手段5.1 置信度阈值把「不确定」当成一种正常输出softmax 输出的概率分布不代表真实置信度小数据集上训练出来的模型经常对错误类别也给出 0.9 以上的高分。直接把这个分数展示给用户一旦模型认错用户会认为系统完全不可用。常见做法是设一个阈值低于阈值就提示用户重新拍摄把「我不确定」也变成系统的一种正常输出。def predict_with_threshold(self, path, threshold0.6): tensor preprocess(path) with torch.no_grad(): prob torch.softmax(self.model(tensor), dim1)[0] score, index torch.max(prob, dim0) if score threshold: return f无法确定请换角度重新拍摄{score:.2f}, score return self.idx_to_label[index.item()], score阈值怎么定在验证集上跑一遍统计每个类别的平均置信度和最差置信度把阈值定在「能覆盖大多数正确样本、又能挡住大部分错误样本」的位置。六分类系统一般从 0.5 试到 0.7每调一档看一次验证集上的拒绝率。拒绝率太高说明数据质量问题严重或者类别本身容易混淆这时候补数据比调阈值有效。5.2 用TorchScript固化模型部署端不再依赖训练代码把模型和权重一起固化成单个文件GUI 加载时不需要再实例化 GarbageCNN也不依赖任何训练相关的类定义。这类固化对交付课程设计和内部工具特别省事对方拿到一个 .pt 文件加一段加载代码就能跑。import torch model GarbageCNN(num_classes6) model.load_state_dict(torch.load(best_garbage_cnn.pth, map_locationcpu)) model.eval() dummy torch.randn(1, 3, 224, 224) traced_model torch.jit.trace(model, dummy) traced_model.save(garbage_cnn.pt) # 部署端加载无需再 import GarbageCNN deployed torch.jit.load(garbage_cnn.pt) result deployed(dummy)这套骨架网络没有 if 分支和动态控制流适合用 trace 方式固化如果模型里含依赖输入尺寸的分支逻辑追踪出的路径可能不完整换成 script 方式更稳妥。固化后再跑一遍 5.1 节的阈值函数确认输出完全一致模型转换这一步就算验证通过。5.3 把批测脚本绑进GUI做交付前体检最后给一个实践技巧把 inference 函数复用成批测脚本。遍历验证集所有图片用同一套 predict_with_threshold 统计每个类别的准确率和平均置信度输出到 CSV按准确率升序排列。这个脚本不用写新模型逻辑只是复用按钮回调里的 predict 函数跑完就能看到哪几个类别在拖后腿。针对排在最前的混淆类别回去翻数据增强参数和样本数量通常比继续调学习率更快见效。本文还有配套的精品资源点击获取
返回列表