简介面向深度学习初学者和垃圾分类项目开发者资料包内含硬纸板、纸、塑料瓶、玻璃瓶、铜制品与不可回收垃圾共6类图片数据集并配齐从数据预处理、模型训练到图片分类预测的Python代码。压缩包共7个文件以5个py脚本为主干分别负责数据准备、训练、预测、图片预处理和结果可视化另有1个h5权重文件可直接加载以及1个zip格式数据集压缩包整体大小161.27MB。因原始训练数据量庞大资源未直接附带超大数据集但zip包中已提供6类图片数据可先跑通完整流程如需提升模型精度可自行补充数据后使用训练脚本。已有30656人浏览学习适合课程设计、毕业设计或垃圾分类识别项目快速落地运行预测脚本并输入图片路径即可输出类别便于在此基础上扩展新类别。1. 垃圾分类数据集及代码是图像分类项目里最常被低估的一环我接过一个社区智能回收站的需求摄像头拍下居民投放的垃圾系统要实时判断袋子是可回收、厨余、有害还是其他判断错就拿不到积分居民当场就会投诉。翻了一圈现成方案发现卡住进度的往往不是模型结构而是训练用的垃圾分类数据集及代码怎么组织、怎么预处理、怎么调参。数据没理清后面再新的模型也白搭。这篇笔记写给三类人正在做课程设计或毕业课题的学生想快速搭出图像分类Demo的初学者以及准备把垃圾分类模型推到边缘设备上的工程师。我会按自己实际跑通的路径来讲从前置的数据集选型、目录规范到预处理、训练代码、超参数再到部署前必须处理的坑。新手照着能一步步走通熟手可以直接跳到参数和边界部分对照。看完你不需要再去翻零散的博客片段照着做一次跑通是大概率事件。2. 选数据集和定目录结构这一步决定后面 80% 的翻车率初次接触垃圾分类数据集及代码的人最容易犯的错是上来就下载一个几十GB的全量压缩包结果类别名是英文缩写、图片分辨率五花八门、坏图一大堆光清洗就耗掉一周。我一般会先逼自己回答一个问题任务到底是几分类是只识别单品种比如塑料瓶、易拉罐、纸板、玻璃瓶还是做回收四分类把可回收、有害、厨余、其他分开。这两个方向的选数据策略完全不同。2.1 公开数据集怎么选三类来源与适用边界常见做法是先用小型公开数据集跑通流程再针对真实场景补数据。目前能拿到的垃圾分类相关图像数据大体可以归成三类。数据来源适用场景需要注意的点某高校实验室整理的图像集课程设计、算法对比图片干净、类别明确但总量有限某平台开源的环保主题图像集学术实验、类别覆盖测试标签体系不一致需要自己映射自己拍摄加定向补充边缘部署、真实场景标注工作量大必须人工复核第一类适合做可行性验证。我在模拟项目X里用它跑基线几千张图十几个常见垃圾类别一张入门显卡训练一轮只要十几分钟非常适合把整体代码链路调通。第二类覆盖面更广但很多图片是从视频帧里抽出来的模糊图和重复图比例偏高必须做清洗和去重否则模型会把视频帧里的运动模糊当成真实特征学进去。第三类看着最贴近部署场景实际坑最多。不同时段的光照、居民投放时的随意姿态、摄像头本身的畸变都会让模型中看不中用。我见过不少团队在实验室里跑出 95% 的准确率一上现场摄像头就掉到 60% 多根因就是把宝全押在公开数据上没有为部署场景补拍。选型的核心逻辑是先用小数据集跑通代码再针对场景做增量一上来追求数据量大清洗工作会淹没主要精力。2.2 目录结构约定ImageFolder 格式与标签文件无论最终用哪类数据我都建议把所有图片整理成 ImageFolder 标准结构。这个格式最大的好处是 PyTorch 的 torchvision.datasets.ImageFolder 可以直接读取连解析器都不用自己写。data/garbage/ ├── train/ │ ├── recyclable/ │ │ ├── plastic_bottle_001.jpg │ │ ├── plastic_bottle_002.jpg │ │ └── ... │ ├── kitchen/ │ ├── harmful/ │ └── other/ └── val/ ├── recyclable/ ├── kitchen/ ├── harmful/ └── other/目录名就是类别名类别顺序由数据集在读取时按字母序自动生成。这里有个容易踩的坑训练集和验证集的目录名必须完全一致。如果一边叫 recyclable另一边叫 recycleImageFolder 会生成两套不同的标签映射训练代码照样跑但验证结果全部失效因为模型输出的是训练集的类别索引而验证集用另一套索引去算准确率错位得一塌糊涂。还需要单独维护一份类别映射文件推理时用。常见做法是训练完把 class_to_idx 存成 JSON这份文件是模型和业务含义之间的翻译官。import json # 训练完拿到 train_dataset 后把类别映射存下来 class_to_idx train_dataset.class_to_idx with open(garbage_class_to_idx.json, w, encodingutf-8) as f: json.dump(class_to_idx, f, ensure_asciiFalse, indent2)推理端从 JSON 反向构建 idx_to_class 字典再配合 softmax 输出取最大概率的位置才能把模型输出的数字翻译回中文类别名。这份映射文件要跟着模型一起发布千万别只留在训练环境里否则部署端加载模型后完全不知道每个输出位代表什么。到这里数据侧准备才算结束。3. 预处理与数据划分让垃圾分类数据集真正喂得进模型数据整理成 ImageFolder 之后还不能直接丢给 DataLoader。我见过的翻车案例里相当一部分不是模型选错而是预处理时把标签和图像弄拧了或者没有做类别均衡处理导致模型把所有样本都预测成多数类。这一章先讲最关键的类别不平衡再讲具体变换参数怎么设。3.1 类别不平衡的三种处理手段垃圾分类数据集的天然问题就是类别不平衡。可回收类里的塑料瓶、纸板样本量大有害类里的电池、灯管样本量小直接硬训模型大概率学会偷懒无脑输出多数类。先跑一个统计脚本确认分布再决定要不要处理。import os from collections import Counter data_root data/garbage/train counts Counter() for cls_name in os.listdir(data_root): cls_path os.path.join(data_root, cls_name) if os.path.isdir(cls_path): counts[cls_name] len(os.listdir(cls_path)) for cls_name, num in counts.most_common(): print(f{cls_name}: {num})对照输出结果如果最大类样本数是最小类的 3 倍以上就有必要处理。我一般按优先级做三件事一是对少数类做离线复制和轻微扰动比如随机旋转、水平翻转各生成一份副本二是用 WeightedRandomSampler 在每轮迭代里让少数类有更高采样概率三是如果业务方允许把细粒度类别合并成回收四分类从根源上弱化不平衡。第一件事最直接但要注意别把同一张图的扰动版本同时分进训练集和验证集否则验证指标会虚高线上表现对不上。第二件事实现简单权重按样本数倒数计算但配合数据增强时可能让同一张图在连续几个迭代里反复出现反而容易过拟合少数类。第三件事改动最大需要业务方确认但对落地部署往往是最稳的选择模型要学的类别边界变少又贴合居民实际投放习惯。3.2 数据增强与归一化参数怎么设预处理逻辑直接写在 Dataset 的 transform 里。我的习惯是训练和验证各用一套验证集绝不做随机增强因为评估要的是确定性和可比性。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomResizedCrop(224, scale(0.8, 1.0)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.2, contrast0.2, 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((256, 256)), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]) ])Resize 先统一到 256再用 RandomResizedCrop 随机裁剪到 224等价于给模型提供不同缩放和裁剪位置的样本对拍摄距离不固定的场景明显更友好。ColorJitter 的亮度、对比度、饱和度扰动用来模拟垃圾房里不同时段的光照差异数值不要开太大0.2 左右足够开过头反而会把垃圾本身的颜色特征破坏比如把绿色玻璃瓶和绿色易拉罐搅在一起。Normalize 里的 mean 和 std 是 ImageNet 预训练权重的标准值。只要你是用 PyTorch 官方预训练模型这两组值不要自己重新统计否则微调时收敛速度会明显拖慢。验证集的 CenterCrop 保证每次评估都吃同一块区域指标才横纵可比。还有一个容易被忽略的细节如果原始图片长宽比很悬殊直接 Resize 会拉伸变形。我会先做一次短边等比缩放再把边缘补成正方形最后进流程。否则模型会把形状畸变当成类别特征学进去部署时一遇到正常比例图片反而认不出来。4. 用代码把模型跑起来迁移学习是当前最可靠的起点数据准备妥当之后就进入标题里「代码」最核心的部分。我不建议从头训练一个卷积网络垃圾分类这种任务数据量再扩充也就几千到几万张从头训练既慢又容易过拟合。常见做法是加载在 ImageNet 上预训练过的分类模型换掉最后一层全连接再决定冻结还是微调。4.1 用迁移学习搭一个基线模型下面这段代码是模拟项目X里验证过的写法用 ResNet18 做骨架输出类别数从训练数据集自动读入。import torch.nn as nn from torchvision import models def build_model(num_classes): model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) in_features model.fc.in_features model.fc nn.Sequential( nn.Dropout(p0.3), nn.Linear(in_features, num_classes) ) return modelweights 参数必须显式指定预训练权重。如果随手写成 weightsNone代码不会报错但随机初始化的特征提取层在小数据量上训练效果会退化到接近从零开始之前省下的时间全白费。Dropout 放在全连接之前用来抑制微调阶段的过拟合这是我在迁移学习里最常加的一层对轻量分类任务帮助明显。换掉 model.fc 之后整个网络依然端到端可训练。想进一步缩短时间可以先把除 fc 外的参数 requires_grad 置为 False只训分类头十几个 epoch再解冻全部层改用小学习率微调。两种做法在垃圾分类这种中等规模数据上都可行我倾向后者分类头先收敛再让骨干网络慢慢适配垃圾图像的纹理和色彩整体训练曲线更平缓。4.2 训练循环与关键超参数训练循环本身不复杂真正影响最终精度的是几个超参数的选择。下面这套代码是我调过之后能直接复用的版本参数按经验值给。import torch from torch.utils.data import DataLoader from torchvision import datasets train_dataset datasets.ImageFolder(data/garbage/train, transformtrain_transform) val_dataset datasets.ImageFolder(data/garbage/val, transformval_transform) train_loader DataLoader(train_dataset, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) val_loader DataLoader(val_dataset, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue) model build_model(num_classeslen(train_dataset.classes)) optimizer torch.optim.AdamW(model.parameters(), lr3e-4, weight_decay1e-4) scheduler torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max20) criterion nn.CrossEntropyLoss() for epoch in range(20): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(cuda), labels.to(cuda) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() * images.size(0) scheduler.step() val_acc evaluate(model, val_loader) print(fepoch {epoch1}: loss {running_loss/len(train_loader.dataset):.4f}, val_acc {val_acc:.4f})batch_size 32 是显存和收敛速度的折中。如果你的显卡只有 6GB 显存可以降到 16同时把学习率等比下调到 1.5e-4否则直接套用 3e-4 会震荡。AdamW 配合 weight_decay 1e-4 比纯 Adam 更适合迁移学习正则能压住最后一层过拟合。CosineAnnealingLR 的 T_max 设为总 epoch 数让学习率从 3e-4 平滑降到接近 0比固定学习率在中后段更容易挤出最后两三个点的精度。evaluate 函数在每个 epoch 结束后执行监控 val_acc 是判断收敛的核心依据。如果训练 loss 还在降val_acc 却连续五个 epoch 不涨甚至下跌这就是过拟合信号优先把 Dropout 概率调到 0.5再考虑减少训练轮数。我见过不少初学者疯狂加 epoch指望精度自己涨上来结果只是把验证集上的噪声曲线记得更牢。5. 训练翻车排查垃圾分类数据集及代码最容易踩的五个坑这一章写训练阶段的血泪经验。每一条都是我在不同项目里真实遇到过、并且能稳定复现的问题按「现象 → 原因 → 解决」的方式列出排查时可以逐个对照。5.1 验证集 loss 正常下降推理时却全部输出同一个类别现象训练曲线一切正常loss 在降验证准确率也有 70% 以上。但把训练好的模型拿出来对单张图片推理不管输入什么输出都是同一个类比如永远是可回收。原因验证集和训练集分布太接近或者验证集里多数类占比过高模型其实只学会了输出多数类验证指标被假性准确率撑住。另一个常见原因是类别映射错位推理时 idx_to_class 构建反了数字对不上中文名。解决先打印模型的原始 logits看最大值是不是集中在某一个类别上再交叉比对导出的 JSON 映射文件最后专门拿一批少数类图片做推理测试。如果少数类全错回到类别不平衡的处理手段重新训练。5.2 手机照片带 EXIF 旋转信息模型把横图竖图搞混现象手机拍的图片在本地预览方向正常模型推理时把竖着拍的瓶子识别成横着的整体准确率明显低于预期。原因部分手机在保存照片时会写入 EXIF 方向字段而不是直接旋转像素。OpenCV 的 imread 会忽略这个字段PIL 默认也忽略导致同一张图在不同读取方式下方向不一致训练和推理看到的不是同一种东西。解决在 Dataset 读取图片时统一用 ImageOps.exif_transpose 处理让所有图片按 EXIF 信息旋转后再进入预处理流程。这个操作必须在 Resize 之前做否则旋转的是已经缩放过的小图边缘信息会丢失。加了这一行之后我在模拟项目X里的验证准确率直接涨了 4 个百分点。5.3 显存不足但数据量和图片尺寸看起来都不大现象batch_size 设成 32图片是 224 分辨率6GB 显存的卡直接把 CUDA out of memory 抛到脸上。原因很多人把 train_transform 里的 Resize 设成 512 甚至 640实际送入网络的图比预期大好几倍显存翻倍消耗。也有可能是 num_workers 拉满后 DataLoader 预加载占满内存间接拖慢 GPU 释放。解决把 Resize 降回 256让随机裁剪回到 224显存实在不够就换 ResNet18 或 MobileNetV3不要硬顶 ResNet50。梯度累积是最后的备用方案它能把单步 batch 拆小但不会减少总计算量只适合大 batch 需求下补救。5.4 验证集指标虚高真实场景打骨折现象模型在验证集上跑出 95%上线到真实摄像头前只有 60%落差大到没法收场。原因划分数据集时用了纯随机切分同一个瓶子的多张连拍图、不同扰动版本同时落进训练集和验证集模型相当于提前见过答案。垃圾分类数据里同一个物体经常有多张几乎一样的图随机切分几乎必然泄漏。解决按图片文件名前缀或拍摄批次做分组切分保证同一物体的所有图只出现在一个集合里。我在模拟项目X里用拍摄批次编号作为分组键验证准确率虽然掉了 7 个百分点但真实场景表现反而提升了一大截。指标虚高骗自己不如让验证集苛刻一点。5.5 中文类别名在跨平台时引发标签错乱现象数据集目录名包含中文比如「可回收」训练代码在 Windows 上跑通过换到 Linux 服务器上报错或者验证时拿到的标签完全不匹配。原因不同系统默认编码对中文目录名处理不一致ImageFolder 生成的类别顺序也会受文件系统排序影响两次训练可能得到两套不同的标签映射。代码能跑是运气跑不报错但结果错是常态。解决目录名统一用英文小写加下划线recyclable、kitchen_waste 这种风格中文名单独放在 JSON 映射文件里展示层再翻译。这个习惯能省掉大量跨机器排查时间也是我接手任何图像项目的第一道规范化动作。6. 把模型导出成 ONNX 并接上真实摄像头让垃圾分类数据集及代码真正落地训练到这里只是完成了前半程。如果只停留在跑完一个 Demo垃圾分类数据集及代码的投入还没换回实际价值。我建议的收尾是导出 ONNX用 CPU 推理脚本验证一遍再接摄像头做实时预测。import torch model build_model(num_classes4) model.load_state_dict(torch.load(garbage_model.pth, map_locationcpu)) model.eval() dummy_input torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, garbage_model.onnx, input_names[input], output_names[logits], dynamic_axes{input: {0: batch}, logits: {0: batch}}, opset_version13 )opset_version 不用追新13 已经覆盖常见部署框架兼容性反而更好。动态轴加上之后推理端 batch 可以从 1 变成任意值无需重新导出。ONNX 模型体积和 .pth 差不多但部署端不再需要安装 PyTorchCPU 推理速度比 eager 模式快不少对社区回收站这种低成本边缘设备很适合。一个非常实用的技巧是给模型加置信度阈值softmax 输出的最大概率低于 0.5 时直接返回「无法识别」而不是硬塞一个类别。真实场景里居民手里可能拿着数据集中不存在的杂物强行分类比承认不认识更难处理。我习惯先在现场拍几十张照统计正常样本的置信度分布用中位数减一个余量作为阈值这样摄像头前出现未知物品时系统会提示重新投放而不是把纸盒认成塑料瓶。落地这一圈走完后我最大的习惯改变是拿到任何垃圾分类项目先检查数据来源和目录规范再谈模型结构。数据没理清再新的模型也救不回来阈值没调好再准的模型上线也会被用户投诉。希望这篇笔记里踩过的坑能帮你把项目一次做对。希望帮到你。本文还有配套的精品资源点击获取