真实废弃物图像分类实战：PyTorch数据加载、ResNet训练与避坑指南

简介面向图像分类初学者与智能垃圾分拣项目开发者这是一份真实废弃物图像分类数据集覆盖纸板、食品有机物、玻璃、金属、杂项垃圾、纸张、塑料、纺织品与植被9个类别约4800张已标注图片。数据已做预处理并划分训练集与测试集可直接作为分类网络输入配套show脚本可快速可视化样本json文件记录类别映射便于训练与验证。压缩包共2000个文件以1998个jpg图片为主另含1个Python脚本和1个json配置整体155.99MB适合离线下载使用。目前已有65人学习下载。对于垃圾分类识别、轻量化模型调优或毕业设计实验这份数据能省去自行采集与清洗的时间结合作者博客中的图像分类、分割网络改进专栏可以较快跑通完整流程并对比不同网络效果。1. 真实废弃物图像分类数据集4800张实拍图先别急着训网络做图像分类踩过最深的坑是拿 ImageNet 的套路直接套在生活里随手拍的垃圾照片上结果训练集准确率 92%验证集只有 61%。这套真实废弃物图像分类数据集一共约 4800 张实拍图分纸板、食品有机物、玻璃、金属、杂项垃圾、纸张、塑料、纺织品、植被 9 类作者按类别整理好训练集和测试集JSON 写明了类别映射还带可视化脚本数据经过预处理可以直接当分类网络输入。它适合两类人拿来做毕业设计或课程设计需要真实场景的分类数据已经在做检测任务想先用分类模型做候选框预筛。别嫌量小真实场景的垃圾照片背景杂、光照乱、类间相似度高4800 张配合迁移学习足够把 ResNet 级别的模型训起来也为后面接 YOLOv8 这类检测器打底。2. 目录与JSON先摸清9类样本分布再写第一行加载代码真实废弃物分类数据集的目录组织方式很直观train 和 test 两个大目录各自下面按类别建子文件夹。图像直接以类名加编号命名比如 Vegetation_110.jpg 就是植被类的第 110 张样本。JSON 文件负责把类别名与数字标签对应起来训练加载时不需要自己去数文件夹顺序避免不同机器上 os.listdir 排序不一致导致的标签错乱。这个设计很小但很实用能省掉不少标签对齐的麻烦。2.1 用JSON做类别映射加载的第一步是建立稳定的标签索引拿到数据集首先要做的事不是立刻训练而是把 JSON 读出来看一眼类别索引是否唯一、顺序是否固定。常见做法是写一个小脚本打印全部类别和样本数确认每个类在训练集和测试集里都有覆盖避免某个类别只有训练样本没有测试样本那会让最后的评估数字变得不可信。import json from pathlib import Path data_root Path(./waste_dataset) meta json.loads((data_root / label_map.json).read_text()) print(类别数量:, len(meta)) for name, idx in meta.items(): train_cnt len(list((data_root / train / name).glob(*.jpg))) test_cnt len(list((data_root / test / name).glob(*.jpg))) print(idx, name, train:, train_cnt, test:, test_cnt)逻辑说明meta.items()拿到类别名与数字索引的键值对glob(*.jpg)统计每个类别文件夹下的图片数量注意如果你的数据里有 PNG 格式需要把后缀也加进 glob 匹配里。参数上值得注意两点一是类别名最好保持和文件夹同名后面 Dataset 加载时才不会找不到路径二是索引从 0 开始还是从 1 开始要以 JSON 为准我见过不少数据集把背景类编为 0这里如果植被是第 8 号索引CrossEntropy 默认就用 0 到 8 的连续整数不需要手动改。跑完这段脚本你会对类别分布有一个直观认识。如果发现某个类的 test 数量只有个位数后续评估时这个类的准确率会非常不稳定训练时就要考虑要不要从训练集手动切一部分过去。这一步花五分钟能避免后面训练完才发现评估样本不够的尴尬。2.2 show脚本可视化巡检标注的准确性比数量更值得花时间作者在资源里放了一个 show 脚本作用是把同一目录下的图片拼成网格一并输出方便快速浏览。拿到数据后的第一件事我建议直接跑它先人工看一遍再训练。运行方式一般是python show.py --data_dir ./waste_dataset/train --class_name 玻璃 --grid 4x4逻辑说明--data_dir指向训练集目录--class_name选择你要巡检的类别--grid指定每张拼接图的格子数4x4 就是每次看 16 张图。巡检时重点看三样东西一是图片里到底是不是对应的废弃物二是图片是否存在拉伸或变形三是同类别内部的拍摄角度差异大不大。真实垃圾照片里最容易出现的问题是「标签标得没毛病但图片主体被手或阴影挡住」。这一类样本不进训练集反而是好事进了会让模型学到不该学的特征。如果发现某个类别里有大量遮挡样本建议先从训练集里剔除等模型第一版训完看结果再决定要不要加回来。2.3 用Dataset类包一层路径、标签、预处理一次解决光有 JSON 和目录还不够训练时需要一个统一的 Dataset 类。这里给一个 PyTorch 实现我一般会把它放在 data.py 里。import json from pathlib import Path from PIL import Image, ImageOps import torch from torch.utils.data import Dataset class WasteDataset(Dataset): def __init__(self, data_root, splittrain, transformNone): self.root Path(data_root) / split meta json.loads((Path(data_root) / label_map.json).read_text()) self.classes list(meta.keys()) # 类别名列表 self.transform transform self.samples [] # (路径, 标签) 列表 for cls_name, label in meta.items(): for img_path in (self.root / cls_name).glob(*.jpg): self.samples.append((str(img_path), label)) def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label self.samples[idx] img Image.open(path).convert(RGB) img ImageOps.exif_transpose(img) # 修正手机照片的EXIF旋转 if self.transform: img self.transform(img) return img, torch.tensor(label, dtypetorch.long)逻辑说明__init__阶段就把所有样本路径和标签读进内存4800 条字符串的开销很小训练时不会频繁访问文件系统也不依赖文件夹的实际枚举顺序。__getitem__里对 PIL Image 做了convert(RGB)防止灰度图或带透明通道的 PNG 混进来导致通道数不一致。ImageOps.exif_transpose是个容易被忽略但很关键的细节它把 EXIF 里的旋转信息直接应用到像素上不影响后续增强的一致性。返回值里 label 包成torch.long因为 CrossEntropyLoss 要求标签是整型张量。transform 参数留给训练脚本注入训练和测试用不同的预处理。3. 分类训练队列预处理、增强与ResNet基线数据加载写好后下一步是确定训练用的预处理流水线。这批数据作者已经做过一轮清洗图像尺寸比较规整但直接拿原始尺寸喂网络仍然不现实。我一般会把图像缩放到 256x256再随机裁剪到 224x224这是 ResNet、EfficientNet、ViT 这类分类网络最通用的输入尺寸。归一化直接用 ImageNet 的均值与标准差即 mean[0.485, 0.456, 0.406]std[0.229, 0.224, 0.225]。原因很简单加载预训练权重时模型期望的输入分布和 ImageNet 统计接近用同一套归一化可以让初始特征直接生效。3.1 训练与测试的两套预处理增强要克制别把垃圾图变成抽象画训练集的预处理我一般用这组操作RandomResizedCrop(224, scale(0.6, 1.0)) 随机截取主体并缩放到 224RandomHorizontalFlip() 随机水平翻转ColorJitter(brightness0.2, contrast0.2, saturation0.2) 做颜色扰动。测试集只用 Resize(256) 加 CenterCrop(224)任何随机操作都不加保证评估结果可复现。from torchvision import transforms train_transform transforms.Compose([ transforms.Resize((256, 256)), transforms.RandomResizedCrop(224, scale(0.6, 1.0)), transforms.RandomHorizontalFlip(), 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]), ]) test_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, 256))直接强制缩放简单粗暴如果希望保持宽高比可以换成Resize(256)加CenterCrop(224)的组合。训练集里先 Resize 再 RandomResizedCrop 相当于在 256 分辨率下随机取一块区域再放大增加尺度的多样性。参数上有几个可以讨论的点scale 下界为什么取 0.6 而不是更激进的 0.3因为垃圾图像的主体占比通常不小裁剪比例太低会让模型看到一堆背景反而学不到物体特征ColorJitter 的幅度控制在 0.2 以内防止把玻璃和塑料的色泽扰动得过度接近两者本来就容易混淆。一个常见的误用是把训练增强直接套在验证集上这会让你看到的验证准确率上下乱跳后面的避坑章节会专门展开。3.2 基线训练脚本ResNet18加ImageNet预训练权重约30分钟出结果下面是完整的训练脚本按单张消费级显卡 8GB 显存来写。batch_size 取 32输入 224x224 的 RGB 图ResNet18 在 8GB 显存下毫无压力更大的 batch 收益不大反而可能让 AdamW 的稳定性变差。import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision.models import resnet18 from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR from tqdm import tqdm from data import WasteDataset, train_transform, test_transform def main(): train_set WasteDataset(./waste_dataset, train, train_transform) test_set WasteDataset(./waste_dataset, test, test_transform) train_loader DataLoader(train_set, batch_size32, shuffleTrue, num_workers4, pin_memoryTrue) test_loader DataLoader(test_set, batch_size32, shuffleFalse, num_workers4, pin_memoryTrue) model resnet18(weightsIMAGENET1K_V1) model.fc nn.Linear(model.fc.in_features, 9) model model.cuda() criterion nn.CrossEntropyLoss() optimizer AdamW(model.parameters(), lr3e-4, weight_decay1e-4) scheduler CosineAnnealingLR(optimizer, T_max30, eta_min1e-6) best_acc 0.0 for epoch in range(30): model.train() total_loss, correct, total 0.0, 0, 0 loop tqdm(train_loader, descfEpoch {epoch1}/30) for imgs, labels in loop: imgs, labels imgs.cuda(), labels.cuda() optimizer.zero_grad() out model(imgs) loss criterion(out, labels) loss.backward() optimizer.step() total_loss loss.item() * len(imgs) pred out.argmax(dim1) correct (pred labels).sum().item() total len(imgs) loop.set_postfix(lossf{loss.item():.4f}) scheduler.step() model.eval() val_correct, val_total 0, 0 with torch.no_grad(): for imgs, labels in test_loader: imgs, labels imgs.cuda(), labels.cuda() out model(imgs) val_correct (out.argmax(1) labels).sum().item() val_total len(imgs) val_acc val_correct / val_total print(fEpoch {epoch1}: train_acc{correct/total:.4f} val_acc{val_acc:.4f}) if val_acc best_acc: best_acc val_acc torch.save(model.state_dict(), best_waste_model.pt) print(saved best model.) if __name__ __main__: main()逻辑说明最后一个全连接层用 9 替换分类数必须和 JSON 里对齐训练前最好打印model.fc确认in_features是 512。优化器选 AdamW 而非 SGD权重衰减设 1e-4对垃圾图像这类样本量小的任务过拟合压力主要来自数据多样性不足1e-4 的正则在 30 个 epoch 下基本不会拖慢收敛。CosineAnnealingLR 把学习率从 3e-4 平滑衰减到 1e-6配合 AdamW 在训练后半段做精细调整。每个 epoch 结束用测试集做一次验证按验证集准确率存档最优权重这个习惯能避免训练到最后过拟合时把最优权重覆盖掉。硬件正常的情况下30 个 epoch 在 ResNet18 上大约需要 30 分钟上下。3.3 用混淆矩阵和单类准确率找真问题只看一个总体准确率很误导。9 个类别如果纸板和纸张本来就分不清总体准确率再高也不代表模型可用。我在训练完第一版后一定会做两件事打印每个类别的 precision、recall保存一张混淆矩阵图。import numpy as np import matplotlib.pyplot as plt from sklearn.metrics import confusion_matrix, classification_report import torch def evaluate_model(model, loader, classes): model.eval() y_true, y_pred [], [] with torch.no_grad(): for imgs, labels in loader: imgs imgs.cuda() out model(imgs) y_true labels.cpu().tolist() y_pred out.argmax(1).cpu().tolist() print(classification_report(y_true, y_pred, target_namesclasses, digits3)) cm confusion_matrix(y_true, y_pred) return cm cm evaluate_model(model, test_loader, test_set.classes) plt.imshow(cm, cmapBlues) plt.colorbar() plt.tight_layout() plt.savefig(confusion_matrix.png, dpi150)逻辑说明classification_report会输出每个类别的 precision、recall、F1重点看那些偏低的具体数字混淆矩阵保存成图片后直接用看图工具打开对角线越亮说明分类越干净非对角线的高亮块就是下一步要处理的硬骨头。这一步判断出的「哪两类最常被搞混」比任何玄学调参都重要后面避坑章节的第一条就是从这个矩阵里发现的。4. 避坑指南真实废弃物分类的5条踩坑记录这套数据集用的是真实照片意味着它不像 CIFAR 那样「干净」。下面每条都是我在类似数据集上实际翻过车的经验按现象、原因、解决三步整理希望能帮你少走弯路。4.1 玻璃和塑料互相认错标签边界本身就是模糊的现象训练完第一版损失降得挺快但混淆矩阵里玻璃和塑料之间的错误特别集中单类召回率一个 71%、一个 66%肉眼能看到不少玻璃瓶被预测成塑料。原因这两种废弃物的物理特性本来就接近透明或半透明的材质在复杂背景里容易反光、透出底色加上拍摄角度不同同一个塑料瓶在不同照片里可能呈现「玻璃质感」。更麻烦的是像饮料瓶这类物体标签是塑料但瓶身透明度和高光表现和玻璃瓶几乎一样。解决我没有去改网络结构而是做了三件事。第一调整类别粒度如果业务允许把玻璃瓶和塑料瓶合并看成一个「透明容器」类再做二级判断第二训练时加大 ColorJitter 的饱和度扰动让模型不能依赖颜色这一单一维度同时增加 RandomErasing 随机遮挡逼模型去学轮廓而不是高光第三针对混淆严重的样本单独抽出来人工复核了一遍标签发现大概 3% 的图片标签本身就标得不准修正后再训玻璃的 recall 提升了 5 个点。4.2 模型记住的是背景不是垃圾本身现象训练集准确率蹭蹭涨验证集死活上不去典型的过拟合特征。用 Grad-CAM 把注意力可视化出来发现模型高亮区域全在图里的草地和水泥地上垃圾主体反而被忽略。原因真实垃圾照片的拍摄背景高度重复。植被类大概率在户外绿地拍摄纸张类大概率在桌面或地板上拍如果训练集里植被类 95% 的背景是草地模型根本不需要看叶子和树枝看绿色背景就能猜对标签换到测试集里自然露馅。解决这个坑靠增强方案就能缓解把 RandomResizedCrop 的 scale 下界从 0.6 改到 0.4强制模型看更多局部细节再加一步 RandomErasing随机抹掉图像中央或边缘的一块区域从数据层面切断「看到草地就答植被」的捷径。更彻底的做法是训练过程中定期把 Grad-CAM 输出存下来人工确认注意力是否回到物体上这一步我用得很勤因为背景干扰不可能完全靠参数消除。4.3 类别不均衡导致少数类直接躺平现象训练过程中发现金属类和杂项垃圾类的 loss 一直比其他类高训练结束后这两个类的 precision 只有 50% 上下。原因9 个类别数量并不均匀植被和塑料的样本明显多金属、纺织品偏少模型在数据多的类别上收敛得很好少数类学不到足够判别特征。解决先看类别统计确定哪些类少于 300 张然后训练时给 CrossEntropyLoss 传一个类别权重权重大小按样本总数的倒数归一化来算如果权重加成后少数类还是弱就把这些类的训练图做更多增强旋转、缩放、颜色抖动都加大或者把损失函数换成 Focal Loss把难样本的梯度权重抬起来。我在这个数据集上把少数类的权重调到 2.0 左右金属类 recall 从 58% 提到 74%代价是植被类降了约 1 个点总体准确率基本持平但每个类的表现均匀了不少。4.4 验证集评分反复横跳评估模式与随机种子现象同一份最优权重两次跑测试集准确率分别是 86.4% 和 84.1%一度怀疑程序有 bug后来发现是评估流程的问题。原因三个容易踩的点凑在一起。训练完权重后忘记调回model.eval()BN 层还在用 batch 内统计量导致两次评估数值不同DataLoader 的 num_workers 在 Windows 上没设固定种子图像读取顺序不同GPU 计算的精度波动被放大另外随机增强混进了测试集虽然测试集 transform 只做 Resize 和 CenterCrop但如果图省事把训练 transform 直接用在验证阶段结果一定反复横跳。解决固定全局种子torch.manual_seed(0)num_workers 在 Linux 上设为 4在 Windows 上先设为 0 跑通再逐步加大评估前一定加model.eval()并用torch.no_grad()包住测试集 transform 单独定义绝不能复用训练增强。从那以后我每次评估都强制走这三步数值稳定在 ±0.2% 以内。4.5 EXIF旋转和坏图预处理阶段最容易被忽略的元数据坑现象训练到第 10 个 epoch 突然报错RuntimeError: CUDA error: device-side assert triggered一看是标签越界。原因罪魁祸首不是标签而是数据集里有几张手机拍的图带有 EXIF Orientation 信息PIL 默认不处理旋转导致图片方向不对个别图片被系统缩略图替换读出来是损坏的 RGB 文件最终在 DataLoader 里产生异常。更隐蔽的是这类坏图不一定每次都会报错如果它恰好被 DataLoader 丢弃你看到的症状就是训练集样本数莫名少几十张。解决在 Dataset 的__getitem__里加ImageOps.exif_transpose(img)处理 EXIF 旋转同时用 try/except 捕获PIL.UnidentifiedImageError遇到坏图直接跳过并打印路径。这是我翻过不小的车才加上的防御性代码强烈建议每个 Dataset 都内置。5. 进阶迁移学习、TTA与半监督式增量迭代到这里你已经有了一个能跑通且知道坑在哪的基线。最后一节聊三个能实际提升成绩的习惯性做法都不需要改网络结构。5.1 迁移学习是默认选项4800 张图从零训练比较大的网络很容易过拟合直接加载 ImageNet 预训练权重是成本最低的提点方式。torchvision 里resnet18(weightsIMAGENET1K_V1)即可。训练时不要对所有层使用同一个学习率我一般把 backbone 的学习率设为分类头的十分之一用参数分组实现。原因很直接分类头是随机初始化的需要较快更新backbone 的预训练特征已经足够好用偏小的学习率做微调即可能明显减少早期几个 epoch 的 loss 震荡。作者还维护着一套图像分类与分割网络改进的系列博客想追更精深的改进方向可以顺路去翻但基线先按这章来别一上来就上高级技巧。5.2 测试时增强不用改模型白捡一到两个点TTA 是性价比很高的技巧用多组预处理对同一张图做预测最后投票决定类别。常见组合是原图、水平翻转图、中心裁剪放大图。实现不复杂把每个结果 argmax 后做一个多数投票。在垃圾图像这种背景复杂、方向不固定的场景里TTA 通常能带来 0.5 到 1.5 个点的提升并且降低误分类的方差代价只是推理时间翻两三倍离线评估时很划算。5.3 半监督式增量把「模型错得自信」的样本回收再利用训练完第一版后跑一遍测试集把预测概率最高但预测错的样本找出来。这些是错得最自信的图通常是标注错误或类别相近的边界样本。把这些图单独移到一个 review 文件夹人工看一遍标签错的直接修正类别真的模糊就归并到大类修正后的样本补充回训练集再训一轮。这个习惯我每次训练分类模型都会做在 4800 张的中小数据集上第二轮训练基本能提升 2 到 3 个百分点比换更大的模型划算得多。我现在每次拿到新的分类数据集流程已经固定成先跑 show 脚本看标注质量再确认 JSON 映射训练第一轮看混淆矩阵最后做一次增量复核不急着换模型结构。这套顺序帮我避开了太多「模型没毛病、数据有毛病」的返工希望帮到你。本文还有配套的精品资源点击获取