垃圾分类图像分类实战：3类瓶子数据集从训练到调优

简介这份深度学习数据集面向图像分类初学者与算法实践者聚焦垃圾分类场景中的瓶子识别任务可用于训练卷积神经网络完成塑料瓶、玻璃瓶、金属瓶等类别的自动判别适合课程设计、模型练手与分类算法对比实验。资源包共约2000个文件以1998张jpg图像为主体另附1个py脚本与1个json类别文件压缩包大小约39.65MB图像按类别目录存放同类样本集中于同一文件夹便于直接接入ImageFolder等标准加载方式。数据已预先划分为训练集约2200张、验证集约200张、测试集约100张划分比例合理可支撑模型训练、调参与最终评估的完整流程。目前已有269人学习下载读者可借助该数据集快速搭建分类基线、验证数据增强与迁移学习效果并通过json类别文件核对标签映射减少自行标注与整理目录的时间成本。1. 垃圾分类图像分类数据集3 类瓶子的 2500 张图能跑出什么结果手上这份垃圾分类图像分类数据集第一次打开目录的时候我愣了一下——文件名全是can_1722347393584_jpg.rf.190f23ba1fed08a779304c3fe011df96.jpg这种带哈希后缀的长串一眼看过去完全分不清哪张是塑料瓶、哪张是玻璃瓶。但翻了两层目录就明白了数据按类别分文件夹存放同一类的图全在一个目录下这是图像分类任务最省心的组织方式ImageFolder或者flow_from_directory直接就能读不用自己写解析脚本。这份资源解决的是「想练图像分类但找不到干净小数据集」的问题。2500 张左右的图3 个类别塑料瓶、玻璃瓶、金属瓶训练/验证/测试按 2200/200/100 切好拿来跑通一个 CNN 或者微调一个轻量模型刚刚好。适合刚入门深度学习、想完整走一遍「读数据→训模型→看混淆矩阵」流程的人也适合需要快速验证某个 backbone 在细粒度瓶类分类上表现的熟手。它不追求 SOTA追求的是让你半小时内看到 loss 往下掉。2. 目录结构与类别映射先搞清楚 json 类别文件和三个 split 怎么对2.1 目录长什么样json 类别文件怎么用数据集按目录保存相同数据放在同一目录下这是摘要里明确写的。实际拿到手大概率是这种结构dataset/ ├── train/ │ ├── plastic/ │ ├── glass/ │ └── metal/ ├── val/ │ ├── plastic/ │ ├── glass/ │ └── metal/ └── test/ ├── plastic/ ├── glass/ └── metal/类别个数是 3具体类别参考 json 类别文件。这里有个血泪经验很多人拿到数据集直接os.listdir去数类别结果顺序是随机的训练时 label 0 可能是塑料瓶推理时 label 0 变成了玻璃瓶混淆矩阵整个错位。正确做法是先读 json把{0: plastic, 1: glass, 2: metal}这种映射固定下来训练和推理共用同一份。import json import os # 读类别映射训练和推理必须共用这一份 with open(dataset/class_names.json, r, encodingutf-8) as f: class_map json.load(f) # 按 key 排序保证 label 顺序稳定 class_names [class_map[k] for k in sorted(class_map.keys(), keyint)] print(class_names) # [plastic, glass, metal] num_classes len(class_names)这段代码的关键在sorted(..., keyint)。json 的 key 是字符串直接sorted会得到0,1,2没问题但如果类别超过 10 个10会排在2前面所以强制转 int 排序。参数上class_names的顺序就是后续模型输出的 index 顺序写进配置文件别在代码里硬编码。2.2 三个 split 的数量核对与路径校验摘要给的数字是训练集 2200 张左右、验证集 200 张左右、测试集 100 张左右。拿到手第一件事不是急着训是数一遍确认没有空目录、没有损坏文件。from pathlib import Path from PIL import Image root Path(dataset) for split in [train, val, test]: total 0 for cls in class_names: cls_dir root / split / cls imgs list(cls_dir.glob(*.jpg)) # 顺手校验能不能打开损坏图直接揪出来 bad [p.name for p in imgs if not _is_valid(p)] total len(imgs) print(f{split}/{cls}: {len(imgs)} 张, 损坏 {len(bad)}) print(f{split} 合计: {total}) def _is_valid(p): try: Image.open(p).verify() return True except Exception: return False逻辑说明glob(*.jpg)只抓 jpg如果你的数据里有 png 或 jpeg 后缀这里会漏数改成glob(*)再过滤后缀更稳。Image.open(p).verify()是轻量校验只读文件头不解码全图速度快但它会把文件指针弄乱校验完如果要继续用这张图得重新 open。参数上没太多可调的重点是养成「先数再训」的习惯我见过太多人训到一半报Truncated File才发现有张图是半截下载的。数量对不上怎么办如果训练集只有 2000 出头先别慌摘要写的是「左右」。但如果某个类别只有几十张那就是类别不平衡后面训练要加WeightedRandomSampler或者 class weight这个在第 4 章会展开。3. 从零跑通训练ImageFolder 轻量 CNN 的最小可复现流程3.1 用 ImageFolder 和 DataLoader 把数据喂进去图像分类任务里torchvision.datasets.ImageFolder是最省事的读法它自动按子目录名映射 label目录名排序后的 index 就是 label。前提是你的目录结构就是 2.1 那种split/类别/图片。import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader # 训练集做增强验证测试只做 resize 归一化 train_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.RandomHorizontalFlip(p0.5), transforms.RandomRotation(15), transforms.ColorJitter(brightness0.2, contrast0.2), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) eval_tf transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean[0.485, 0.456, 0.406], std[0.229, 0.224, 0.225]), ]) train_ds datasets.ImageFolder(dataset/train, transformtrain_tf) val_ds datasets.ImageFolder(dataset/val, transformeval_tf) test_ds datasets.ImageFolder(dataset/test, transformeval_tf) # 关键确认 ImageFolder 的类别顺序和 json 一致 assert train_ds.classes class_names, f类别顺序不一致: {train_ds.classes} train_loader DataLoader(train_ds, batch_size32, shuffleTrue, num_workers4) val_loader DataLoader(val_ds, batch_size32, shuffleFalse, num_workers4) test_loader DataLoader(test_ds, batch_size32, shuffleFalse, num_workers4)逻辑说明ImageFolder的classes属性是它自己按目录名排序得到的必须和 2.1 读 json 得到的class_names对齐否则 label 全错。assert这行别省我吃过亏——有一次 json 里写的是plastic/glass/metal目录名却是bottle_plastic/bottle_glass/bottle_metal排序后顺序变了训出来的模型把塑料瓶全判成玻璃瓶。参数说明Resize((224,224))是给 ImageNet 预训练模型用的标准输入如果你自己搭小 CNN可以降到 128 省显存。Normalize的 mean/std 是 ImageNet 统计值用预训练权重时必须一致从零训可以换成自己数据集的统计值但差别不大。batch_size32在 8G 显存上跑 224 分辨率够用爆显存就降到 16。num_workers4在 Windows 上如果报错就设 0。3.2 搭一个能跑通的轻量 CNN 并训练不追求精度的话一个四层卷积的小网络就够把这三个类分开。瓶子类之间的差异主要在材质反光和形状浅层特征就能抓到。import torch.nn as nn class BottleCNN(nn.Module): def __init__(self, num_classes3): super().__init__() self.features nn.Sequential( nn.Conv2d(3, 32, 3, padding1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), # 224 - 112 nn.Conv2d(32, 64, 3, padding1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2), # 112 - 56 nn.Conv2d(64, 128, 3, padding1), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(2), # 56 - 28 nn.Conv2d(128, 256, 3, padding1), nn.BatchNorm2d(256), nn.ReLU(), nn.AdaptiveAvgPool2d(1), # 任意尺寸 - 1x1 ) self.classifier nn.Linear(256, num_classes) def forward(self, x): x self.features(x) x x.flatten(1) return self.classifier(x) device torch.device(cuda if torch.cuda.is_available() else cpu) model BottleCNN(num_classes3).to(device) criterion nn.CrossEntropyLoss() optimizer torch.optim.AdamW(model.parameters(), lr1e-3, weight_decay1e-4)逻辑说明AdaptiveAvgPool2d(1)把任意空间尺寸压成 1x1这样输入分辨率变了也不用改全连接层。BatchNorm2d放在卷积和 ReLU 之间是标准顺序能加速收敛。AdamW比Adam多了正确的 weight decay 解耦小数据集上更稳。训练循环def train_one_epoch(model, loader, criterion, optimizer, device): model.train() total_loss, correct, total 0, 0, 0 for imgs, labels in loader: imgs, labels imgs.to(device), labels.to(device) optimizer.zero_grad() out model(imgs) loss criterion(out, labels) loss.backward() optimizer.step() total_loss loss.item() * imgs.size(0) correct (out.argmax(1) labels).sum().item() total imgs.size(0) return total_loss / total, correct / total for epoch in range(20): tr_loss, tr_acc train_one_epoch(model, train_loader, criterion, optimizer, device) print(fepoch {epoch}: loss{tr_loss:.4f} acc{tr_acc:.4f})参数说明lr1e-3是 AdamW 在小 CNN 上的常用起点loss 不降就降到 1e-4。weight_decay1e-4防过拟合2500 张图不算多正则要跟上。epoch 数看验证集一般 15 到 30 之间收敛。loss.item() * imgs.size(0)是按样本数加权平均比直接平均每个 batch 的 loss 更准因为最后一个 batch 可能不满。3.3 验证集评估与混淆矩阵训练完不看混淆矩阵等于白训。三个类别的分类任务最容易出的问题是某两类互相混。from sklearn.metrics import confusion_matrix, classification_report import numpy as np model.eval() all_preds, all_labels [], [] with torch.no_grad(): for imgs, labels in val_loader: imgs imgs.to(device) preds model(imgs).argmax(1).cpu().numpy() all_preds.extend(preds) all_labels.extend(labels.numpy()) print(classification_report(all_labels, all_preds, target_namesclass_names)) print(confusion_matrix(all_labels, all_preds))逻辑说明model.eval()关掉 dropout 和 batchnorm 的训练模式torch.no_grad()省显存。classification_report给出每个类的 precision/recall/f1比只看总 acc 有用得多。混淆矩阵对角线是判对的非对角线看哪两类在混。如果玻璃瓶和塑料瓶混得多说明模型没学到材质差异要么加数据增强里的ColorJitter强度要么换预训练 backbone。4. 避坑与排查这份数据集上最容易翻车的五个点4.1 现象训练 acc 冲到 99%验证 acc 只有 60%原因2500 张图里训练集 2200 张验证集才 200 张如果训练集和验证集的图来自同一批拍摄场景模型可能记住了背景而不是瓶子本身。更常见的是数据增强只加在训练集验证集分布和训练集差太多。解决先检查训练集和验证集是不是同一来源。如果是把验证集比例提到 20%或者做交叉验证。增强别加太猛RandomRotation(15)已经够加到 45 度瓶子都躺平了反而学不到。另外确认验证集的transform里没有RandomHorizontalFlip这类随机操作。4.2 现象某个类别 recall 特别低比如金属瓶总是被判成玻璃瓶原因金属瓶和玻璃瓶在反光上确实像如果数据集里金属瓶的样本数明显少于另外两类模型会偏向多数类。解决先数每个类的数量。如果金属瓶只有 500 张而塑料瓶有 900 张用WeightedRandomSampler给少数类加权from torch.utils.data import WeightedRandomSampler targets [s[1] for s in train_ds.samples] class_count np.bincount(targets) class_weight 1.0 / class_count sample_weight [class_weight[t] for t in targets] sampler WeightedRandomSampler(sample_weight, num_sampleslen(sample_weight), replacementTrue) train_loader DataLoader(train_ds, batch_size32, samplersampler, num_workers4)参数说明replacementTrue表示有放回采样少数类会被重复抽到。num_samples设成训练集大小保证每个 epoch 看到的样本数和原来一致。注意用了sampler就不能再设shuffleTrue两者冲突。4.3 现象DataLoader 报RuntimeError: DataLoader worker (pid xxx) is killed by signal原因num_workers开太大共享内存不够或者 Windows 上多进程 spawn 的问题。解决Linux 上把num_workers降到 2 或 0同时设torch.multiprocessing.set_sharing_strategy(file_system)。Windows 上直接把num_workers0虽然慢但稳。另外确认if __name__ __main__:包住了训练入口Windows 多进程必须这样。4.4 现象图片文件名带.rf.哈希后缀某些库读不出来原因这份数据集的文件名是can_1722347393584_jpg.rf.190f23ba1fed08a779304c3fe011df96.jpg这种格式中间有多个点。大部分库按最后一个点取后缀没问题但有些老版本的工具会截断错误。解决用Path(p).suffix取后缀它只认最后一个点。如果某个库报错批量重命名成短名字import uuid from pathlib import Path for p in Path(dataset/train/plastic).glob(*.jpg): new_name p.parent / f{uuid.uuid4().hex[:8]}.jpg p.rename(new_name)注意重命名后要重新生成 json 映射或者确认 label 不受影响因为 label 来自目录名不是文件名所以安全。4.5 现象测试集只有 100 张评估结果波动大这次 85% 下次 78%原因100 张测试集每张图权重 1%错 5 张就是 5 个百分点。小测试集上的单次评估不可靠。解决测试集评估跑三次每次用不同的随机种子做数据增强比如不同的 crop取平均。或者把验证集和测试集合并成 300 张做最终评估训练时从训练集里再切一部分做验证。别在小测试集上反复调参调着调着就过拟合测试集了。5. 进阶技巧用预训练 backbone 把验证集 acc 从 80% 推到 95%小 CNN 从零训在这份 2500 张的数据集上验证集 acc 大概能到 80% 上下。想再往上走最划算的一步是换预训练 backbone。不用自己搭torchvision.models里现成的就行。import torchvision.models as models import torch.nn as nn # 用 resnet18 预训练权重只换最后的分类头 model models.resnet18(weightsmodels.ResNet18_Weights.IMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, 3) model model.to(device) # 分层学习率backbone 小分类头大 optimizer torch.optim.AdamW([ {params: model.fc.parameters(), lr: 1e-3}, {params: [p for n, p in model.named_parameters() if not n.startswith(fc)], lr: 1e-4}, ], weight_decay1e-4)逻辑说明ResNet18_Weights.IMAGENET1K_V1是 ImageNet 预训练权重特征提取层已经学会了边缘、纹理、材质这些通用特征瓶子分类正好用得上。model.fc换成 3 类输出。分层学习率是关键——backbone 用 1e-4 微调别把预训练学到的特征冲掉分类头是随机初始化的用 1e-3 快速学。训练策略上前 5 个 epoch 可以冻结 backbone 只训分类头等分类头稳定了再解冻一起微调# 阶段一冻结 backbone for name, param in model.named_parameters(): if not name.startswith(fc): param.requires_grad False # 训 5 个 epoch 后解冻 for name, param in model.named_parameters(): param.requires_grad True参数说明冻结阶段lr1e-3只更新 fc5 个 epoch 足够。解冻后整体lr降到 1e-4再训 10 到 15 个 epoch。这份数据集上ResNet18 微调后验证集 acc 通常能到 93% 到 96%测试集因为只有 100 张波动大但混淆矩阵里金属瓶和玻璃瓶的混淆会明显减少。还有一个容易忽略的点输入分辨率。ResNet18 预训练是 224但瓶子图像如果原图很大缩到 224 会丢细节。可以试试 288 或 320显存够的话。改分辨率后Normalize不用变AdaptiveAvgPool会自动适配。最后说个验证方法训完别只看 acc把测试集里判错的图导出来看一眼。我一般会写个小脚本把pred ! label的图复制到一个errors/目录按「真实类_预测类」分子目录。翻一遍错图经常能发现是标注问题或者图本身模糊到人都分不清。从那以后我每次训完分类模型都强制走一遍错图检查比盯着 acc 数字有用得多。希望帮到你。本文还有配套的精品资源点击获取