简介:这份深度学习图像分类数据集面向从事计算机视觉入门与垃圾分类识别实践的开发者、学生及算法爱好者,围绕塑料瓶、玻璃瓶、金属瓶等可回收物类别构建,可直接用于训练与评估卷积神经网络分类模型。资源包共约2000个文件,以1998张jpg图像为主体,另附1个py脚本与1个json类别文件,压缩包大小约39.65MB,数据按类别目录存放,相同类别图像归入同一文件夹,便于直接读取与标签映射。数据已预先划分为训练集约2200张、验证集约200张、测试集约100张,开箱即可投入模型训练与调参。目前已有268人学习下载,适合希望快速搭建垃圾分类识别基线、验证网络结构或完成课程设计的读者参考使用。
1. 垃圾分类图像数据集:从「能跑通」到「能落地」差在哪
垃圾分类图像分类这件事,很多人第一次做都会低估它。你拿一个公开数据集,用 ResNet 或 EfficientNet 跑一遍,验证集准确率轻松上 95%,感觉任务已经解决了。但真把模型放到实际场景里——比如一个社区智能回收箱的摄像头前——准确率可能直接掉到 60% 以下。原因不复杂:公开数据集里的图像往往是摆拍、背景干净、光照均匀,而真实场景里瓶子被捏扁、标签被撕掉、油污遮挡、逆光拍摄,分布差异巨大。
这个方向的核心技术链路其实很清晰:收集或下载垃圾分类图像数据集,做清洗和标注,选一个图像分类模型(CNN 为主,也可以上最新的图像分类模型如 ConvNeXt、EfficientNetV2),训练、评估、导出部署。但每一步都有大量细节决定最终能不能用。适合谁看?如果你正在做深度学习实战项目案例、需要交课程设计、或者想把这个方向做成一个真正能演示的系统,这篇笔记会帮你把从数据集到模型的完整路径走一遍,重点讲那些跑通之后才会暴露的坑。
2. 垃圾分类数据集怎么选、怎么造、怎么洗
2.1 公开数据集盘点与选型判断
垃圾分类方向最常被引用的公开数据集是 TrashNet,大约 2500 张图像,分 6 类:玻璃、纸、纸板、塑料、金属、其他垃圾。它的优点是干净、标注质量高、适合入门;缺点是量太小、场景单一,直接拿来训练很容易过拟合,验证集分数好看但泛化差。
另一个常见来源是 Kaggle 上的垃圾分类竞赛数据集,通常包含训练集和测试集,类别划分更细,有的按可回收/不可回收二分,有的按具体材质分十几类。这类数据集图像数量从几千到上万不等,质量参差不齐,需要自己做清洗。
如果你要做更贴近实际的系统,我一般会建议:公开数据集打底,再自己补拍一批真实场景图像。补拍时注意几个维度——不同光照(顺光、逆光、室内荧光灯)、不同角度(俯拍、侧拍)、不同状态(完整、捏扁、带标签、去标签)、不同背景(纯色桌面、地面、垃圾桶内壁)。每个类别至少补 200 到 300 张,才能让模型学到真正的判别特征而不是背景捷径。
注意:不要直接用 ImageNet 预训练模型在 TrashNet 上微调后就宣称「垃圾分类准确率 98%」。那个数字只在 TrashNet 的测试集上成立,换一批真实照片就会翻车。
2.2 数据清洗与标注规范
拿到数据集后第一步不是训练,是清洗。常见问题包括:重复图像、标签错误、极端模糊、类别不均衡。重复图像会导致训练集和验证集泄漏,验证分数虚高。可以用感知哈希(pHash)快速去重:
import imagehash from PIL import Image import os def deduplicate(image_dir, hash_size=8, threshold=5): """基于感知哈希去除重复或高度相似的图像""" hashes = {} duplicates = [] for fname in os.listdir(image_dir): if not fname.lower().endswith(('.jpg', '.jpeg', '.png')): continue path = os.path.join(image_dir, fname) try: img = Image.open(path) h = imagehash.phash(img, hash_size=hash_size) except Exception: continue # 与已有哈希比较汉明距离 is_dup = False for existing_h, existing_f in hashes.items(): if h - existing_h <= threshold: duplicates.append((path, existing_f)) is_dup = True break if not is_dup: hashes[h] = path return duplicates # 用法:dups = deduplicate('./data/train') # 返回的 duplicates 列表里每项是 (重复图路径, 原始图路径)这段代码的逻辑是:对每张图计算 pHash,得到一个 64 位指纹,然后比较汉明距离。距离小于等于阈值就认为是重复或高度相似。hash_size控制指纹精度,8 对应 64 位,一般够用;threshold设 5 比较保守,能抓住明显重复但不误杀相似不同图。跑完后人工抽查一遍 duplicates 列表,确认没有误判再删除。
标注方面,垃圾分类图像分类任务通常不需要边界框,只需要类别标签。但如果你后续想升级成检测(比如同时定位和分类),那就需要 YOLO 格式的标注。常见做法是先用分类标签训练一个基线,再决定要不要上检测。标注工具用 LabelImg 或 CVAT 都行,关键是制定清晰的标注规范:瓶子带盖算塑料还是其他?纸杯有塑料涂层算纸还是其他?这些边界情况必须在标注前统一,否则模型学到的决策边界会混乱。
2.3 数据增强策略与划分比例
垃圾分类图像分类的数据增强不能随便套 ImageNet 那一套。水平翻转通常安全,但垂直翻转对垃圾图像不一定合理(瓶子倒过来还是瓶子,但某些类别可能不自然)。颜色抖动要谨慎,因为颜色是垃圾分类的重要线索——你把塑料瓶的蓝色调成绿色,模型可能就懵了。比较稳妥的增强组合是:随机裁剪、小角度旋转(±15 度)、轻微亮度对比度调整、随机擦除。
数据划分比例建议训练:验证:测试 = 7:1.5:1.5,并且要按「拍摄批次」划分而不是随机划分。什么意思?如果你同一天同一场景拍了 100 张瓶子,随机划分会让训练集和验证集里都有这个场景的图,验证分数会偏高。正确做法是把不同天、不同场景的图分到不同集合,这样验证分数才能反映真实泛化能力。
import splitfolders # 按 7:1.5:1.5 划分,seed 固定保证可复现 splitfolders.ratio( './data/raw', output='./data/split', seed=42, ratio=(0.7, 0.15, 0.15), group_prefix=None # 如果有分组前缀可以在这里指定 )splitfolders这个库的好处是会自动创建 train/val/test 三个目录,每个类别一个子文件夹,直接兼容 PyTorch 的ImageFolder。seed固定后每次划分结果一致,方便复现。如果你有分组信息(比如每张图属于哪个拍摄批次),可以用group_prefix参数确保同一组的图不会被分到不同集合。
3. 用 CNN 和最新图像分类模型跑通基线
3.1 从零训练还是微调:判断依据
垃圾分类图像分类,我一般不会从零训练。原因很简单:公开数据集量级通常几千到几万,从零训练一个 CNN 很容易过拟合,而且收敛慢。微调预训练模型是更务实的选择。PyTorch 的 torchvision 里自带 ResNet、EfficientNet、ConvNeXt 等模型,加载 ImageNet 预训练权重后替换最后一层全连接,就能直接用于垃圾分类。
选哪个骨干?如果追求精度且算力够,ConvNeXt-Tiny 或 EfficientNetV2-S 是当前比较新的选择,在图像分类任务上表现稳定。如果算力有限或者要部署到边缘设备,MobileNetV3 或 EfficientNet-B0 更合适。ResNet-50 是经典基线,适合快速验证想法。
import torch import torch.nn as nn from torchvision import models def build_model(num_classes, backbone='resnet50', pretrained=True): """构建垃圾分类分类模型,替换最后一层""" if backbone == 'resnet50': model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT if pretrained else None) in_features = model.fc.in_features model.fc = nn.Sequential( nn.Dropout(0.3), nn.Linear(in_features, num_classes) ) elif backbone == 'efficientnet_v2_s': model = models.efficientnet_v2_s(weights=models.EfficientNet_V2_S_Weights.DEFAULT if pretrained else None) in_features = model.classifier[1].in_features model.classifier = nn.Sequential( nn.Dropout(0.3), nn.Linear(in_features, num_classes) ) elif backbone == 'convnext_tiny': model = models.convnext_tiny(weights=models.ConvNeXt_Tiny_Weights.DEFAULT if pretrained else None) in_features = model.classifier[2].in_features model.classifier[2] = nn.Linear(in_features, num_classes) else: raise ValueError(f'Unsupported backbone: {backbone}') return model # 用法:model = build_model(num_classes=6, backbone='convnext_tiny')这段代码的关键点在于:不同骨干网络的分类头结构不一样,ResNet 是model.fc,EfficientNetV2 是model.classifier[1],ConvNeXt 是model.classifier[2]。替换时只改最后一层,前面的特征提取层保留预训练权重。加 Dropout 是为了在小数据集上抑制过拟合,0.3 是一个比较安全的起点,如果过拟合严重可以加到 0.5。
3.2 训练循环与关键超参数
训练垃圾分类模型时,有几个超参数比别的更值得调。学习率:微调时用 1e-4 到 1e-3 之间,太大容易破坏预训练权重,太小收敛慢。批量大小:根据显存来,16 或 32 都行,但要注意批量大小和学习率通常要联动调整。优化器:AdamW 比 SGD 更省心,权重衰减设 1e-4 到 1e-2 之间。学习率调度:CosineAnnealingLR 或 ReduceLROnPlateau 都可以,前者更平滑。
import torch.optim as optim from torch.optim.lr_scheduler import CosineAnnealingLR def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss, correct, total = 0.0, 0, 0 for images, labels in loader: images, labels = images.to(device), labels.to(device) optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() total_loss += loss.item() * images.size(0) _, preds = outputs.max(1) correct += (preds == labels).sum().item() total += labels.size(0) return total_loss / total, correct / total # 初始化 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = build_model(num_classes=6, backbone='convnext_tiny').to(device) criterion = nn.CrossEntropyLoss(label_smoothing=0.1) optimizer = optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4) scheduler = CosineAnnealingLR(optimizer, T_max=30)label_smoothing=0.1是一个容易被忽略但很有效的技巧,它把硬标签变成软标签,能缓解过拟合和标签噪声问题。垃圾分类数据里经常有标注模糊的样本,标签平滑能让模型不那么自信,泛化更好。CosineAnnealingLR的T_max设成总 epoch 数,学习率会按余弦曲线从初始值降到接近零,训练后期更稳定。
3.3 评估指标:准确率之外还要看什么
垃圾分类图像分类不能只看总体准确率。如果「其他垃圾」类别占了 40%,模型全预测成「其他垃圾」也有 40% 准确率,但这显然没用。必须看每个类别的精确率、召回率和 F1。特别是召回率——对于可回收物,漏检比误检代价更高,因为漏检意味着可回收物被扔进其他垃圾,失去了回收价值。
from sklearn.metrics import classification_report, confusion_matrix import numpy as np def evaluate(model, loader, device, class_names): model.eval() all_preds, all_labels = [], [] with torch.no_grad(): for images, labels in loader: images = images.to(device) outputs = model(images) _, preds = outputs.max(1) all_preds.extend(preds.cpu().numpy()) all_labels.extend(labels.numpy()) print(classification_report(all_labels, all_preds, target_names=class_names, digits=4)) cm = confusion_matrix(all_labels, all_preds) return np.array(all_preds), np.array(all_labels), cmclassification_report会输出每个类别的 precision、recall、f1-score 和支持样本数。重点看哪些类别的 recall 明显偏低,然后去混淆矩阵里找它被误分成了什么。比如塑料瓶经常被误分成玻璃瓶,那就要检查两类图像在视觉上是不是真的难以区分,或者训练样本里是不是有标注错误。
4. 垃圾分类模型训练避坑与排查
4.1 验证集准确率虚高但实际翻车
现象:验证集准确率 97%,但拿手机拍几张真实垃圾照片测试,错得离谱。原因:训练集和验证集来自同一批图像,分布太接近,模型学到了背景捷径而不是垃圾本身的特征。比如所有塑料瓶都在白色桌面上拍,模型可能学会了「白色背景=塑料」。解决:按拍摄批次划分数据集,确保验证集和训练集来自不同场景;同时补拍真实场景图像,用手机在垃圾桶旁边拍几百张加入训练。
4.2 某个类别召回率极低
现象:玻璃类别的召回率只有 0.5,其他类别都在 0.9 以上。原因:玻璃和塑料在图像上确实相似,尤其是透明瓶子;或者玻璃类别的训练样本太少,模型没学够。解决:先检查混淆矩阵,看玻璃被误分成了什么。如果是塑料,可以增加玻璃类别的样本量,或者在数据增强时对玻璃类别做更强的颜色和光照扰动。另一个思路是引入注意力机制,让模型关注瓶身纹理而不是整体轮廓。
4.3 训练损失震荡不收敛
现象:训练 loss 上下跳动,验证 loss 不下降。原因:学习率太大,或者批量大小太小导致梯度噪声大。解决:把学习率降到 1e-5 试试,或者换用 AdamW 并加 warmup。如果批量大小只有 8 或 16,可以尝试梯度累积,模拟更大的批量。
# 梯度累积示例:每 4 个 batch 更新一次参数 accumulation_steps = 4 optimizer.zero_grad() for i, (images, labels) in enumerate(loader): outputs = model(images.to(device)) loss = criterion(outputs, labels.to(device)) / accumulation_steps loss.backward() if (i + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()4.4 模型导出后推理结果不一致
现象:PyTorch 里测试正常,导出 ONNX 或 TorchScript 后推理结果变了。原因:预处理不一致。训练时用了归一化(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),导出后推理时忘了做同样的归一化,或者输入尺寸不对。解决:把预处理逻辑和模型绑定在一起导出,或者写一个明确的预处理函数,确保训练和推理用同一套参数。
4.5 数据不平衡导致模型偏向多数类
现象:多数类召回率很高,少数类几乎全错。原因:类别样本数差距大,损失函数被多数类主导。解决:用 WeightedRandomSampler 做重采样,或者在 CrossEntropyLoss 里传 class_weights。class_weights 可以按类别频率的倒数来设。
from torch.utils.data import WeightedRandomSampler import numpy as np # 假设 labels 是训练集所有标签的列表 class_counts = np.bincount(labels) class_weights = 1.0 / class_counts sample_weights = class_weights[labels] sampler = WeightedRandomSampler(sample_weights, num_samples=len(sample_weights), replacement=True) # 然后 DataLoader 里传 sampler=sampler,不要再传 shuffle=True5. 把模型推到能演示的程度:导出、量化与一个实用技巧
训练完一个垃圾分类图像分类模型只是起点。如果你要把它做成一个能演示的系统,下一步是导出和优化。PyTorch 模型导出 ONNX 是最通用的做法,导出后可以用 ONNX Runtime 推理,速度比原生 PyTorch 快不少,而且方便集成到其他语言的服务里。
import torch.onnx model.eval() dummy_input = torch.randn(1, 3, 224, 224).to(device) torch.onnx.export( model, dummy_input, 'trash_classifier.onnx', input_names=['input'], output_names=['output'], dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}, opset_version=13 )导出时dynamic_axes让批量维度可变,这样同一个模型可以处理单张图也可以处理一批图。opset_version=13兼容性比较好,大部分推理引擎都支持。导出后一定要用 ONNX Runtime 跑一遍验证,对比 PyTorch 和 ONNX 的输出差异,如果差异大于 1e-3 就要检查导出过程。
如果目标设备是边缘端,还可以做动态量化,把模型大小压缩到原来的四分之一左右,推理速度也能提升。量化后的模型精度通常会掉 1 到 2 个百分点,对垃圾分类这种任务一般可以接受。
import torch.quantization model_cpu = model.to('cpu') model_cpu.eval() quantized_model = torch.quantization.quantize_dynamic( model_cpu, {torch.nn.Linear}, dtype=torch.qint8 ) torch.save(quantized_model.state_dict(), 'trash_classifier_quantized.pth')最后分享一个我在实际项目里养成的习惯:每次训练完,不要只看验证集指标,一定要拿一批「从来没见过的真实照片」做一次盲测。我一般会准备 50 张手机随手拍的垃圾照片,不参与任何训练和调参,只在最后评估时用一次。如果这 50 张的准确率能到 80% 以上,这个模型才算真正能演示。如果只有 50% 以下,说明数据集和真实场景的差距还很大,需要回去补数据而不是调模型。这个习惯帮我省了很多次「演示时翻车」的尴尬。希望帮到你。
本文还有配套的精品资源,点击获取