基于PyTorch的垃圾分类系统实战:数据增强、迁移学习与部署全流程 简介一份基于Python与深度学习的垃圾分类系统毕业论文面向计算机科学与技术及相关专业本科生可用于毕业设计撰写、开题与降重参考。资源为单个docx文件约34KB内容涵盖完整论文框架包括研究背景与意义、Python及深度学习技术综述、系统需求分析与架构设计、数据集采集与模型训练、界面实现与测试以及实验对比和总结展望等章节。论文明晰地展示了从数据预处理到CNN模型调参、系统集成的完整流程并在实验结果中给出准确率等指标对比对于理解垃圾分类场景下深度学习的落地应用具有较高参考价值。文件已有584人学习下载适合需要快速把握论文章节脉络、撰写思路或进行二次修改的毕业生使用。1. 垃圾分类系统最难的不是模型而是这三件事把 Python 和深度学习落到垃圾分类系统的设计实现上多数人第一反应是选模型但真正让项目翻车的从来不是模型结构而是数据组织混乱、类别样本不均衡、训练集与真实拍摄环境分布不一致这三件事。你手上哪怕只有一份标注好的图片目录训练一个能用的分类器也只需要几十行代码可一旦验证集指标好看、现场识别却乱猜问题基本都出在数据管线而非网络结构。这篇文章按实际落地顺序展开先讲数据集怎么整理和增强再讲迁移学习的模型选型与训练参数然后用一整章列出我踩过的坑最后把模型导出并提供可调用的识别服务。适合正在做课程设计或毕业设计的同学也适合想快速验证垃圾分类识别可行性的开发者。照着做你得到的不是一个训练脚本而是一条能自测、能部署、能继续迭代的完整链路。2. 数据集准备目录结构、样本清洗与数据增强怎么做垃圾分类的数据准备比想象中花时间。先别急着写训练代码模型能不能扛住真实场景七成取决于这一章做没做扎实。分类任务的输入是图片输出是类别数据层面的两个核心问题是样本从哪里来以及模型看到的图片和真实拍摄的图片差距有多大。2.1 公开数据集与自建数据的取舍公开的垃圾分类数据集并不难找常见的构成是几十个类别每类样本量从几百到几千不等。类别通常覆盖纸类、塑料、玻璃、金属、织物、厨余和有害垃圾但不同数据集的类目划分口径并不一致有的把塑料瓶单独成类有的把塑料瓶和塑料袋归为同一类。这直接决定了你的分类器能回答什么问题。先确认类目定义再决定训练目标避免模型在“这个瓶子算可回收还是有害”这种边界定义上纠结。自建数据的成本主要在拍摄和标注。手机拍几百张图并不难难的是保持背景多样性。另一个现实问题是公开数据里很多图片是白底或单色背景的产品图而实际场景是桌面、地面、户外、光线不均。我一般会混合使用公开数据提供类别覆盖面自建数据补充真实背景和拍摄角度。这样既保证了样本量也让模型见过“脏”的图像。数据清洗这里要花点耐心。把每个类别目录打开翻一遍删掉明显错标、模糊到人眼都无法判断、以及同一张图在不同类别里重复出现的样本。重复图片看起来无害但在训练集和验证集里各出现一次时验证集指标会被悄悄抬高这就是一种数据泄漏。清洗阶段多花一小时训练阶段少熬夜。2.2 目录结构设计模型读数据的第一步PyTorch 的torchvision.datasets.ImageFolder要求一种固定目录布局根目录下先按train和val划分每个子目录再按类别分。这种约定俗成的结构可以让数据加载代码短到只剩两行也方便后来切换别的框架。目录长这样dataset/ ├── train/ │ ├── plastic_bottle/ │ ├── paper_box/ │ └── ... └── val/ ├── plastic_bottle/ ├── paper_box/ └── ...手头图片散乱时我习惯用一个小脚本把它们按类别和划分比例归置好。下面这个脚本读取原始图片目录子目录名即类别名按比例随机划分并复制到目标结构下# arrange_data.py import os import random import shutil from pathlib import Path random.seed(2024) src Path(raw_images) # 原始目录raw_images/类别/图片.jpg dst Path(dataset) # 目标目录dataset/{train,val}/类别/图片.jpg train_ratio 0.8 # 训练集占比剩余 20% 做验证 for cls_dir in src.iterdir(): if not cls_dir.is_dir(): continue images sorted(cls_dir.glob(*.jpg)) sorted(cls_dir.glob(*.png)) random.shuffle(images) n_train int(len(images) * train_ratio) for phase in (train, val): (dst / phase / cls_dir.name).mkdir(parentsTrue, exist_okTrue) for i, img_path in enumerate(images): phase train if i n_train else val shutil.copy(img_path, dst / phase / cls_dir.name / img_path.name) print(f{cls_dir.name}: train{n_train}, val{len(images) - n_train})这个脚本的时间复杂度在文件数量级几千张图几秒跑完。shutil.copy会复制文件图片总量大时磁盘占用翻一倍如果你只是本机实验可以把copy改成os.link建硬链接秒级完成且不占额外空间。注意脚本里glob(*.jpg)和glob(*.png)只匹配单层后缀如果你有.jpeg或.bmp的图需要补上相应匹配。参数层面值得留意两个地方。train_ratio 0.8在样本充足时没问题但如果每个类别只有一两百张图我建议把比例调到 0.9保证训练样本量。划分前执行random.shuffle(images)是为了打乱文件夹里的原始顺序如果你的图片本身按拍摄时间排列不打乱会让某一段时间的图片全部进验证集导致指标失真。最后切分粒度是“文件名”级别而不是“拍摄批次”级别这一点后面避坑章节还要展开。2.3 数据增强配置小样本下最划算的补充垃圾图片数据量通常不大每类几百张是常态。这时数据增强不是可选项而是提升泛化能力的核心手段。增强要模拟的是真实世界的变化拍摄距离不同、角度不同、光线不同、物体摆放方向不同。我在训练集上用下面这组变换验证集只做缩放和裁剪避免用增强后的图片评估模型from torchvision import transforms # 训练集增强 train_transform transforms.Compose([ transforms.RandomResizedCrop(224, scale(0.6, 1.0)), transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.3, contrast0.3, saturation0.2), transforms.RandomAffine(degrees0, translate(0.05, 0.05)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) # 验证集只做 resize center crop val_transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])逐项说作用于。RandomResizedCrop随机裁剪一块区域并缩放到 224×224scale(0.6, 1.0)的意思是裁剪区域占原图的 60% 到 100%下限不能太低——垃圾主体通常占画面不小裁剪太多会让模型看到一个残缺物体。RandomRotation(15)的 15 度是经验值旋转超过 20 度会让塑料瓶这种有明显垂直特征的物体变得不自然。ColorJitter调节亮度、对比度、饱和度对应不同光线环境。RandomAffine这里的translate(0.05, 0.05)表示图片水平垂直最多平移 5%模拟拍摄时物体不在画面正中间的情况。归一化为什么要用[0.485, 0.456, 0.406]这组数因为后面要加载在 ImageNet 上预训练好的权重ImageNet 数据的均值和标准差就是这三组值。输入分布和预训练时保持一致迁移效果才能发挥出来。如果你的类目里包含大量黑白或特殊色调的图片也不要改这两组数保持默认即可。验证集增强最常见的错误是把旋转、翻转也加到验证集上这会让你评估到的准确率偏高因为模型相当于见过这些变换后的图片。验证集要的是稳定、一致的处理越少变换越好。3. 模型选型与训练配置从迁移学习到可复用的训练脚本垃圾分类本质上是一个图像分类问题模型选型不需要多花哨。真正影响训练体验的是预训练权重的选择、超参数的设置以及你判断“训练已经够了”的能力。这一章给出一个可以直接跑通的训练脚本并解释每一个关键参数为什么这么设。3.1 为什么从预训练权重开始而不是从零训练垃圾图片的特征比如纹理、轮廓、材料质感和 ImageNet 里那些物体特征有大量重叠。预训练模型已经学会了通用的边缘、纹理、形状表示你要做的不是重新教它看图片而是让它把注意力转向垃圾分类的专属特征。从零训练一个 ResNet在每类只有几百张图的情况下结果往往是验证集准确率上不去训练集却已经拟合得很好典型的小数据过拟合。可选的模型集中在三个ResNet18、ResNet50、MobileNetV3。我的取舍逻辑是这样模型输入尺寸参数量级推理速度适用场景ResNet18224×224约 11M快数据量小、CPU 推理也能跑ResNet50224×224约 25M中等数据量较大追求更高准确率MobileNetV3224×224约 2.5M极快需要部署到边缘设备或追求低延迟课程设计或快速验证阶段我推荐 ResNet18。它的参数量足够表达垃圾分类这种中等复杂度的任务训练速度快调参省心。ResNet50 能带来一点准确率提升但训练时间和显存占用也会涨性价比不高。如果最终要部署到树莓派或安卓端再换 MobileNetV3 也不迟分类头替换方式和 ResNet 完全一样。迁移学习还有一个操作要点图片输入尺寸。上面数据增强已经把训练图统一成 224×224这和预训练权重要求的输入尺寸一致。不要自作主张改成 128 或 299预训练特征的统计特性是在 224 下学出来的乱改输入尺寸会削弱迁移效果。3.2 训练脚本与关键超参数下面这个训练脚本是完整可跑的核心逻辑包括加载预训练模型、替换分类头、处理类别不均衡、训练循环、验证并保存最优权重# train.py import numpy as np import torch import torch.nn as nn from torch.utils.data import DataLoader, WeightedRandomSampler from torchvision import datasets, models BATCH_SIZE 32 EPOCHS 30 DEVICE torch.device(cuda if torch.cuda.is_available() else cpu) train_dataset datasets.ImageFolder(rootdataset/train, transformtrain_transform) val_dataset datasets.ImageFolder(rootdataset/val, transformval_transform) # 类别不均衡时按样本数的倒数加权采样 class_counts np.bincount(train_dataset.targets) sample_weights 1.0 / class_counts[np.array(train_dataset.targets)] sampler WeightedRandomSampler(sample_weights, num_sampleslen(sample_weights), replacementTrue) train_loader DataLoader(train_dataset, batch_sizeBATCH_SIZE, samplersampler) val_loader DataLoader(val_dataset, batch_sizeBATCH_SIZE, shuffleFalse) # 加载 ImageNet 预训练权重 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) num_classes len(train_dataset.classes) model.fc nn.Linear(model.fc.in_features, num_classes) model.to(DEVICE) criterion nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr1e-3) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_maxEPOCHS) best_acc 0.0 for epoch in range(EPOCHS): model.train() total_loss 0.0 for images, labels in train_loader: images, labels images.to(DEVICE), labels.to(DEVICE) optimizer.zero_grad() logits model(images) loss criterion(logits, labels) loss.backward() optimizer.step() total_loss loss.item() scheduler.step() model.eval() correct total 0 with torch.no_grad(): for images, labels in val_loader: images, labels images.to(DEVICE), labels.to(DEVICE) outputs model(images) preds outputs.argmax(dim1) correct (preds labels).sum().item() total labels.size(0) acc correct / total print(fepoch{epoch1}/{EPOCHS}, loss{total_loss/len(train_loader):.4f}, val_acc{acc:.4f}) if acc best_acc: best_acc acc torch.save(model.state_dict(), best_cls.pth) print(fbest val_acc{best_acc:.4f})几个参数说明。BATCH_SIZE 32在 224×224 输入下ResNet18 大约占用 3-4G 显存显存不够就降到 16显存宽裕也可以升到 64。Adam优化器的lr1e-3是迁移学习的常见起点这个学习率对分类头合适但作用在预训练主干上会偏大。sample_weights的计算方式是每个样本的权重等于其类别样本数的倒数少数类样本权重更高让模型每个 batch 里更均衡地看到各个类别。训练脚本里有两个容易踩的细节。WeightedRandomSampler和DataLoader的shuffleTrue互斥设置了sampler后shuffle必须删掉或设为False否则会报错。num_samples这里设为len(sample_weights)表示每个 epoch 采样的样本数和训练集总样本数一致如果你想每个 epoch 让少数类被更多次采样可以调大num_samples效果等于额外补充了少数类样本。训练初期如果发现模型容易把预测推向某几个类别把CrossEntropyLoss改为带权重版本权重同样用类别倒数这是比加权采样更直接的干预手段。它的缺点是改变损失函数的梯度尺度学习率需要相应调小一点。项目初期先不加跑一个 baseline 观察混淆矩阵再决定要不要上。3.3 看损失曲线决定何时停训练过程不是跑完固定 epoch 就结束。我习惯每轮打印训练 loss 和验证准确率然后看两条曲线的相对关系。正常情况下训练 loss 持续下降验证准确率同步上升最终稳定在一个区间。这里有个经验标准如果训练 loss 还在下降但验证准确率已经连续 5 轮不涨说明模型开始过拟合早停比继续跑更划算。另一种常见形态是训练 loss 降得很慢验证准确率也上不去。先检查学习率是不是太大导致 loss 震荡不收敛再检查数据增强是不是过强比如裁剪比例太小把物体主体切掉了。还有一个易忽略的点是否忘记在验证模式下用torch.no_grad()验证阶段开启梯度计算会额外消耗显存但不影响正确性。训练过程的 loss 波动是正常现象尤其 batch size 只有 16 或 32 时单个 batch 的噪声会让 loss 出现尖峰。不要看到一次冲高就打断训练至少要观察连续几个 epoch 的趋势。这算是训练阶段最玄学的地方多数中途放弃的模型再跑 10 个 epoch 往往就稳下来了。4. 训练过程避坑与效果排查四类翻车与对应检查路径这一章讲的全是实际会遇到的坑每一条对应一个明确的检查路径。跑训练之前先通读一遍遇到症状可以直接对照。如果你已经跑出了不错的结果也可以把这一章当体检清单逐条排查你的训练管线是否有隐患。4.1 类别不均衡多数类“赢家通吃”准确率虚高现象训练完成后打印的验证准确率在 90% 以上但看混淆矩阵发现预测结果几乎全部集中到样本最多的几个类别样本少的类别召回率接近零。准确率虚高只是因为多数类占比大模型学到的实际上是一个“猜多数类”的策略。原因交叉熵损失在计算时每个样本对梯度的贡献是相等的。样本量大的类别贡献了绝大部分梯度模型优先拟合它们少数类被忽略。这是分类任务在数据不平衡时的默认行为。解决两条路同时走。第一条是用WeightedRandomSampler让采样更均衡代码已经在训练脚本里给出这是数据层面的干预不改变损失函数。第二条是对少数类做针对性增强比如某个类别只有 80 张图可以对这个类单独增加旋转、缩放、色彩扰动的强度相当于手动扩充少数类。注意加权的副作用少数类权重过高时模型可能在这几个类上过拟合表现形式是训练 loss 降得很慢验证准确率也不稳定。这时候可以给采样器加一个温度参数把权重做平滑比如sample_weights sample_weights ** 0.7让权重差异不那么极端。4.2 背景干扰比垃圾本身更容易翻转预测现象验证集准确率不错但拿手机在桌面、草地、水泥地上拍几张图测预测结果明显乱套。把图片拿出来看发现模型判断依据不是垃圾本身而是背景的颜色和纹理特征。原因训练集里大量图片是白底或纯色背景模型学到的是“白色背景 → 某类别”的相关性。真实拍摄背景复杂提出特征的分布完全偏离训练分布模型就开始乱猜。背景干扰是垃圾分类落地场景里最常见的问题因为垃圾出现的环境天然混乱。解决短期补救是在训练增强里加入RandomErasing随机遮挡图片的一部分强迫模型不要把注意力集中在单一区域上from torchvision import transforms train_transform transforms.Compose([ # 在之前增强链的末尾加上 transforms.RandomErasing(p0.3, scale(0.02, 0.15)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])p0.3表示 30% 概率触发擦除scale(0.02, 0.15)表示擦除区域占原图面积的 2% 到 15%。这样模型必须依赖多个局部特征做判断不至于被某个区域主导。长期的解决方式是补充真实背景的样本哪怕每类只加几十张效果也比单纯调增强参数好得多。这也是我始终建议自建数据的原因。4.3 验证集指标高现场识别却掉链子现象验证集准确率稳定在 92%把训练脚本原封不动换个数据目录重跑验证集指标却降了 5 个百分点以上。原因随机划分训练集和验证集时同一次拍摄的连续图片被分到了两边。这些图片相似度极高验证集实际上在“考试”时见过“答案”指标虚高。真实场景的图片和训练分布差距大成绩立刻暴露。这类问题在数据集来自同一个拍摄批次时特别容易发生。解决划分验证集时按“拍摄批次”切分而不是按“图片”切分。比如 A 同学拍的图片整体进训练集B 同学拍的图片整体进验证集。如果你用爬虫或公开数据按来源域名或采集时间分。判断方法很简单如果验证集准确率明显高于训练集且训练过程正常先怀疑是不是数据泄漏了。4.4 训练结果不稳定随机种子与数据增强的叠加效应现象同一个脚本连续跑两次验证集准确率差 2 到 3 个百分点学习率、batch size、epoch 都没动。原因深度学习训练里有三处随机性模型权重初始化、数据加载顺序、数据增强操作的随机采样。迁移学习模型的初始权重是固定的但数据加载顺序和增强的不确定性足以让最终结果产生可见波动。这种波动在调参时会让你误判哪个参数改了、哪个参数没改。解决在训练脚本开头固定随机种子让每次实验可复现import random import numpy as np import torch def set_seed(seed42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed)固定种子之后数据和增强顺序仍然有细微差异但结果基本稳定在 0.5 个百分点以内。调参时你会感谢这个习惯一个参数改动的效果能看得清不会因为随机波动而误判方向。需要说明的是固定种子不等于完全复现不同 GPU 型号和框架版本仍会导致微小差异但作为调参对照足够用了。5. 把模型变成能用的垃圾分类服务导出、接口与现场实测训练完的模型还是 PyTorch 的.pth权重文件离“可用”还有两步导出成通用推理格式再封装成一个能接收图片、返回类别的接口。这两步做完就可以进入真正的验收环节。5.1 导出 ONNX摆脱训练框架依赖PyTorch 模型直接部署需要带上整个训练框架体积大且不方便。常见做法是导出为 ONNX 格式然后只用onnxruntime做推理。导出代码这样写import torch from torchvision import models model models.resnet18(weightsNone) model.fc torch.nn.Linear(model.fc.in_features, 40) # 40 改成你的类别数 model.load_state_dict(torch.load(best_cls.pth, map_locationcpu)) model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, cls.onnx, input_names[input], output_names[output], opset_version13 )dummy_input是导出的模版输入形状必须是 (1, 3, 224, 224)和训练时的输入保持一致。opset_version越高导出的算子在较新版本onnxruntime上支持越好如果你手里的推理环境比较旧优先用 11 或 13兼容性更稳。导出的模型同时包含网络结构和权重部署时不再需要torchvision。5.2 一个最小的 Web 识别服务用 Flask 加onnxruntime可以写一个 60 行左右的推理服务。它接收 POST 上传的图片预处理后输入 ONNX 模型返回类别 ID、类别名和置信度# app.py import io import numpy as np import onnxruntime as ort from PIL import Image from flask import Flask, request, jsonify from torchvision import transforms app Flask(__name__) sess ort.InferenceSession(cls.onnx, providers[CPUExecutionProvider]) transform transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) class_names [plastic, paper, glass, metal, organics, other] # 按训练时的类别顺序 app.route(/predict, methods[POST]) def predict(): file request.files[image] img Image.open(io.BytesIO(file.read())).convert(RGB) x transform(img).unsqueeze(0).numpy() logits sess.run([output], {input: x})[0] idx int(np.argmax(logits[0])) return jsonify({ class_id: idx, class_name: class_names[idx], confidence: float(np.max(logits[0])) }) if __name__ __main__: app.run(host0.0.0.0, port8899)代码里最容易出错的是class_names的顺序。ImageFolder会按类别目录名的字母顺序分配类别索引所以这个列表不是随意写的而是在训练时通过train_dataset.classes打印出来的顺序。顺序错了模型输出的索引和名字就会对不上预测结果看着就像“乱猜”。启动服务后用一行命令做冒烟测试curl -X POST -F imagetest.jpg http://localhost:8899/predict这个接口对图片大小没有限制PIL 会自动读取但过大的图片会占用更多内存。生产环境应在Resize之前检测图片尺寸上限超过 2000×2000 的图片先等比缩放到短边 1000 以内再做后续预处理。5.3 现场实测的验证套路部署完成后真正有价值的验收不是再看一遍验证集指标而是用手机拍 20 到 30 张现场照片覆盖不同光线、角度和背景逐张调用接口测试。把这些结果记录下来做一个小混淆矩阵看看预测失败集中在哪些类别上。如果失败图片里大量出现“背景复杂的瓶子”说明背景干扰的问题没解决干净需要补这类样本如果失败集中在纸箱和塑料膜这种视觉相似的类别考虑合并类目或增加类别内的多样性。我的个人习惯是保留一个“现场测试集”这些图片不参与训练每次改完模型都跑一遍。这个习惯帮我挡过很多次教训——有些调整让验证集指标涨了现场测试却更差。只有两种指标都过了才算真正改进了。垃圾分类这类视觉任务模型的最终水平不是由训练脚本决定的而是由你愿不愿意花时间补齐数据盲区决定的。希望这篇笔记能帮你在自己项目里少走一段弯路。本文还有配套的精品资源点击获取